Update modeling_quiet.py

modeling_quiet.py

@@ -1693,8 +1693,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
             )
 
             prev_hidden_states = hidden_states
-            print(f"1696 Hidden states contains NaN: {torch.isnan(hidden_states).any().item()}")
-
             hidden_states = outputs[0]
             prev_rm_logits = rm_logits  # for policy gradient
             prev_rm_tokens = cur_rm_tokens  # for policy gradient
@@ -1818,12 +1816,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
             if not attempted or self.comparison_mode:
                 rm_hidden_states = hidden_states
                 # print("Magnitude of RM hidden states before RM head", rm_hidden_states.norm())
-                print(f"RM hidden states contains NaN: {torch.isnan(rm_hidden_states).any().item()}")
-
                 rm_logits = apply_head(self.lm_head, rm_hidden_states, detach=self.optimize_lm_head_only_at_start)
-
-                print(f"RM logits contains NaN: {torch.isnan(rm_logits).any().item()}")
-
                     
                 # don't allow it to predict the thinking token
                 if self.tokenizer_has_start_thought_token:                    
@@ -1885,11 +1878,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
 
                 if not contains_thought:
                     with torch.set_grad_enabled(not self.train_only_thinking_embedding):
-                        print(f"Probabilities_2d contains NaN: {torch.isnan(probabilities_2d).any().item()}")
-
                         inputs_embeds = probabilities_2d @ (self.model.embed_tokens.weight.to(probabilities.device).to(probabilities.dtype))
-                        print(f"Inputs_embeds contains NaN: {torch.isnan(inputs_embeds).any().item()}")
-
                 else:
                     thought_id = self.start_token_id if contains_start else self.end_token_id
                     cur_thought_embedding = start_embedding if contains_start else end_embedding
@@ -1926,6 +1915,8 @@ 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())