Update modeling_quiet.py

modeling_quiet.py

@@ -1616,27 +1616,24 @@ class QuietForCausalLM(QuietPreTrainedModel):
             base_embeddings = self.model.embed_tokens.weight
             if self.train_only_thinking_embedding:
                 base_embeddings = base_embeddings.detach()
-            # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-            fwd_iters = 1 if self.original_mode else self.n_ahead + self.n_ahead_talk - 1
-            print("Input IDs contains NaN:", torch.isnan(input_ids).any().item())
-            print("Attention mask contains NaN:", torch.isnan(attention_mask).any().item())
-            print("Labels contains NaN:", torch.isnan(labels).any().item() if labels is not None else False)
-            for ahead_idx in range(fwd_iters):
-                past_key_values_length = 0
-                if past_key_values is not None:
-                    use_legacy_cache = not isinstance(past_key_values, Cache)
-                    if use_legacy_cache:
-                        past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-                    past_key_values_length = past_key_values.get_usable_length(seq_len)
-
-                if position_ids is None:
-                    device = input_ids.device if input_ids is not None else inputs_embeds.device
-                    position_ids = torch.arange(
-                        past_key_values_length, seq_len + past_key_values_length, dtype=torch.long, device=device
-                    )
-                    position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
-                else:
-                    position_ids = position_ids.view(-1, seq_len).long()
+        # # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        fwd_iters = 1 if self.original_mode else self.n_ahead + self.n_ahead_talk - 1
+        for ahead_idx in range(fwd_iters):
+            past_key_values_length = 0
+            if past_key_values is not None:
+                use_legacy_cache = not isinstance(past_key_values, Cache)
+                if use_legacy_cache:
+                    past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+                past_key_values_length = past_key_values.get_usable_length(seq_len)
+
+            if position_ids is None:
+                device = input_ids.device if input_ids is not None else inputs_embeds.device
+                position_ids = torch.arange(
+                    past_key_values_length, seq_len + past_key_values_length, dtype=torch.long, device=device
+                )
+                position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
+            else:
+                position_ids = position_ids.view(-1, seq_len).long()
 
             if inputs_embeds is None:
                 contains_start = self.use_start_thought_token and (input_ids == self.start_token_id).any()
@@ -1697,8 +1694,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
             hidden_states = outputs[0]
             prev_rm_logits = rm_logits  # for policy gradient
             prev_rm_tokens = cur_rm_tokens  # for policy gradient
-            print("Hidden states contains NaN after loop:", torch.isnan(hidden_states).any().item())
-            print("Logits contains NaN after loop:", torch.isnan(logits).any().item())
+
             if ahead_idx == 0:
                 hidden_states_lm = hidden_states
                 logits = self.lm_head(hidden_states_lm)
@@ -2088,7 +2084,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
                                 # This will only happen when we force the next token to be the end of thought token
                                 break
                             dqn_loss_list.append(actor_loss.mean())
-        print("Loss list contains NaN before loss computation:", any(torch.isnan(loss).any() for loss in loss_list))
+
         if loss_list:
             if self.first_and_last_mode:
                 loss = sum(
@@ -2116,20 +2112,18 @@ class QuietForCausalLM(QuietPreTrainedModel):
                 loss = loss / len(loss_list)
             
             loss = loss * self.base_loss_beta
-            print("Loss contains NaN after loss computation:", torch.isnan(loss).any().item())
-print("DQN loss list contains NaN before loss computation:", any(torch.isnan(loss).any() for loss in dqn_loss_list))
+
         if dqn_loss_list:
             dqn_loss = sum(dqn_loss_list) / len(dqn_loss_list)
-print("DQN loss contains NaN after loss computation:", torch.isnan(dqn_loss).any().item())
             if self.include_policy_loss:
                 if loss is not None:
                     loss += dqn_loss * self.policy_loss_beta
                 else:
                     loss = dqn_loss * self.policy_loss_beta
+
         if not return_dict:
             output = (logits,) + outputs[1:]
             return (loss,) + output if loss is not None else output
-        print("DQN loss contains NaN after loss computation:", torch.isnan(dqn_loss).any().item())
     
         base_log_dict = {
             f"loss_{i}": nonzero_mean(loss_list[i]) for i in range(len(loss_list))