Update modeling_quiet.py

modeling_quiet.py

@@ -1110,6 +1110,7 @@ class QuietModel(QuietPreTrainedModel):
         next_decoder_cache = None
 
         for decoder_layer in self.layers:
+            print(f"Hidden states contains NaN before layer {i}:", torch.isnan(hidden_states).any().item())
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
@@ -1167,6 +1168,7 @@ def nonzero_mean(x, axis=None):
 
 def loss_mean(x):
     return x.sum() / (x != 0).sum()
+print(f"Hidden states contains NaN after layer {i}:", torch.isnan(hidden_states).any().item())
 
 class QuietForCausalLM(QuietPreTrainedModel):
     _tied_weights_keys = ["lm_head.weight"]
@@ -1915,8 +1917,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
                     else:
                         loss_logits = logits
                     shift_idx = 1 + max(0, ahead_idx - (self.n_ahead - 1))
-                    import pdb; pdb.set_trace()
-
                     # print("initial_loss_logits contains NaN:", torch.isnan(initial_loss_logits).any().item())
                     # print("logits contains NaN:", torch.isnan(logits).any().item())
                     # print("loss_logits contains NaN:", torch.isnan(loss_logits).any().item())