Update modeling_quiet.py

modeling_quiet.py

@@ -1351,6 +1351,8 @@ class QuietForCausalLM(QuietPreTrainedModel):
                 return_dict=return_dict,
             )
             new_key_values = outputs.past_key_values
+            print(f"Hidden states contains NaN: {torch.isnan(hidden_states).any().item()}")
+
             hidden_states = outputs[0]
             logits = self.lm_head(hidden_states)
             logits = logits[:, -1, :]  # Only consider the last token
@@ -1691,6 +1693,8 @@ 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
@@ -1814,7 +1818,12 @@ 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:                    
@@ -1876,7 +1885,11 @@ 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