Update modeling_quiet.py

modeling_quiet.py

@@ -18,9 +18,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """ PyTorch Quiet model."""
-import inspect
 import math
-import pdb
 import warnings
 from collections import defaultdict
 from typing import List, Optional, Tuple, Union
@@ -31,8 +29,7 @@ import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.generation.utils import GenerationMixin
-from transformers.generation.stopping_criteria import StoppingCriteriaList, validate_stopping_criteria
-from transformers import TextStreamer, AutoTokenizer
+from transformers import AutoTokenizer
 import transformers
 
 from transformers.activations import ACT2FN
@@ -43,8 +40,6 @@ from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
-    is_flash_attn_2_available,
-    is_flash_attn_greater_or_equal_2_10,
     logging,
     replace_return_docstrings,
 )
@@ -240,7 +235,6 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
     """
 
-    # pdb.set_trace()
     batch, num_key_value_heads, slen, head_dim = hidden_states.shape
     if n_rep == 1:
         return hidden_states
@@ -332,7 +326,7 @@ class QuietAttention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
         if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            cache_kwargs = {"sin": sin, "cos": cos} 
             key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         # repeat k/v heads if n_kv_heads < n_heads
@@ -377,8 +371,7 @@ class QuietAttention(nn.Module):
                 )
 
             attn_weights = attn_weights + attention_mask
-
-        # upcast attention to fp32
+            
         attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
         attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
         attn_output = torch.matmul(attn_weights, value_states)
@@ -851,16 +844,12 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         self.model = QuietModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        # self.router_aux_loss_coef = config.router_aux_loss_coef
-        # self.num_experts = config.num_experts
-        # self.num_experts_per_tok = config.num_experts_per_tok
         self.max_thoughts = config.max_thoughts
         self.merged_lm_and_talk_heads = config.merged_lm_and_talk_heads
         self.use_concat_talk_head = config.use_concat_talk_head
         self.use_shallow_talk = config.use_shallow_talk
         self.use_complex_talk_head = config.use_complex_talk_head
         self.use_weighted_talk_head = config.use_weighted_talk_head
-        # the weighted head will output a single value, so it can't be passed to the lm head
         assert not (self.use_weighted_talk_head and self.use_shallow_talk)
 
         self.n_ahead = 1
@@ -931,7 +920,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         self.thinking_threshold = 0.5
         self.thinking_usefulness_loss_weight = 1e-2
 
-        # Not used in the paper:
         self.use_thought_prefix = False
         self.use_reparam_for_thought_embeddings = False
         self.use_upper_triangular = False
@@ -939,7 +927,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         self.comparison_mode = False
         self.gumbel_detach = False
     
-        # For visualization
         self.eval_mode = False
 
         num_talk = 1
@@ -968,7 +955,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         # Add dropout regularization
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
 
-        # Initialize weights and apply final processing
         self.post_init()
 
     def get_input_embeddings(self):
@@ -1219,20 +1205,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             n_passes_to_restore = self.n_passes
             self.n_ahead_talk = 1
             self.n_passes = 1
-        
-        # aux_loss = None
-        # output_router_logits = output_router_logits if output_router_logits is not None else self.config.output_router_logits
-        # if output_router_logits:
-        #     router_logits = outputs.router_logits if return_dict else outputs[-1]
-        #     if router_logits is not None:
-        #         aux_loss = load_balancing_loss_func(
-        #             router_logits,
-        #             self.num_experts,
-        #             self.num_experts_per_tok,
-        #             attention_mask,
-        #         )
-        #         if labels is not None:
-        #             loss += self.router_aux_loss_coef * aux_loss.to(loss.device)
         if input_ids.dim() == 1:
             input_ids = input_ids.unsqueeze(0)
             attention_mask = attention_mask.unsqueeze(0) if attention_mask is not None else None
@@ -1300,7 +1272,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                 self.start_token_id = self.tokenizer.bos_token_id
                 self.tokenizer_has_start_thought_token = False
             elif self.use_start_thought_token:
-                # base_start_id = self.tokenizer.convert_tokens_to_ids(self.initial_start_token)
                 base_start_id = self.tokenizer.encode(self.initial_start_token, add_special_tokens=False)[0]
                 if self.initialize_thought_embedding_to_normal:
                     self.start_embedding.data = torch.zeros_like(self.start_embedding.data)
@@ -1313,7 +1284,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                 self.end_token_id = self.tokenizer.eos_token_id
                 self.tokenizer_has_end_thought_token = False
             elif self.use_end_thought_token:
-                # base_end_id = self.tokenizer.convert_tokens_to_ids(self.initial_end_token)
                 base_end_id = self.tokenizer.encode(self.initial_end_token, add_special_tokens=False)[0]
                 if self.initialize_thought_embedding_to_normal:
                     self.end_embedding.data = torch.zeros_like(self.end_embedding.data)
@@ -1332,7 +1302,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             else:
                 # convert to identity transform
                 def lambda_transform(cur_head):
-                    # pdb.set_trace()
                     if cur_head.weight.data.shape[0] != cur_head.weight.data.shape[1]:
                         return torch.cat([
                         torch.eye(
@@ -1360,28 +1329,23 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                     self.talk_head[-1].weight.data = lambda_transform(self.talk_head[0])
 
         loss = None
-        prev_rm_tokens = None
         cur_rm_tokens = None
-        prev_rm_logits = None
         prev_sample_probs = None
         did_skip_sampling = None
         skip_sampling = None
         sample_probs = None
         hidden_states = None
         logits = None
-        talk_kl_penalty = None
         rm_logits = None
         residual_logits = None
         probabilities_2d = None
         prev_probabilities_2d = None
         policy_reward = None
-        logits_to_output = None
         batch_size, seq_len = input_ids.shape
         base_input_ids = input_ids.clone()
         loss_list = []
         dqn_loss_list = []
         sampled_token_history = []
-        sample_probs_history = []
         action_loglikelihoods_list = []
 
         temperature = self.temperature
@@ -1397,7 +1361,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             if self.train_only_thinking_embedding:
                 base_embeddings = base_embeddings.detach()
                 
-        # # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        # 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
@@ -1442,15 +1406,12 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                     base_attention_mask = base_attention_mask.view(1, 1, seq_len, seq_len)
                     base_attention_mask = base_attention_mask.repeat(input_ids.shape[0], 1, 1, 1)
                     attention_mask = base_attention_mask
-                    # breakpoint()
                 elif attention_mask.dim() == 2:
                     if seq_len + past_key_values_length != attention_mask.shape[-1]:
-                        # breakpoint()
                         attention_mask = torch.cat(
                             [torch.ones((attention_mask.shape[0], past_key_values_length), dtype=attention_mask.dtype, device=attention_mask.device), attention_mask],
                             dim=-1
                         )
-                    # # if the attention mask 
                     attention_mask = _prepare_4d_causal_attention_mask(
                         attention_mask,
                         (batch_size, seq_len),
@@ -1460,7 +1421,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                     )
 
             outputs = self.model(
-                # input_ids=input_ids,
                 attention_mask=attention_mask,
                 position_ids=position_ids,
                 past_key_values=past_key_values,
@@ -1468,14 +1428,13 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                 use_cache=use_cache,
                 output_attentions=output_attentions,
                 output_hidden_states=output_hidden_states,
-                # output_router_logits=output_router_logits,
                 return_dict=return_dict,
             )
 
             prev_hidden_states = hidden_states
             hidden_states = outputs[0]
-            prev_rm_logits = rm_logits  # for policy gradient
-            prev_rm_tokens = cur_rm_tokens  # for policy gradient
+            prev_rm_logits = rm_logits 
+            prev_rm_tokens = cur_rm_tokens
 
             if ahead_idx == 0:
                 hidden_states_lm = hidden_states
@@ -1521,7 +1480,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                 assert sum([self.cumulative_residual, self.clever_residual, self.skip_residual, self.no_residual]) == 1
                 if self.clever_residual:
                     if ahead_idx >= self.n_ahead - 1:
-                        # get the logits shifted according to the current talk ahead
                         cur_base_logits = torch.cat([
                             base_logits[..., ahead_idx - self.n_ahead + 1:, :],
                             base_logits[..., :ahead_idx - self.n_ahead + 1, :]
@@ -1566,7 +1524,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 
             attempted = False
             talk_loss_list = []
-            if self.original_mode or (self.n_ahead == 1) or (self.comparison_mode and ahead_idx == 0):# or (self.optimize_lm_head_only_at_start and ahead_idx == 0):
+            if self.original_mode or (self.n_ahead == 1) or (self.comparison_mode and ahead_idx == 0):
                 loss = None
                 attempted = True
 
@@ -1597,7 +1555,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             
             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())
                 rm_logits = apply_head(self.lm_head, rm_hidden_states, detach=self.optimize_lm_head_only_at_start)
                     
                 # don't allow it to predict the thinking token
@@ -1626,9 +1583,8 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                     probabilities_2d[:, override_token] = 1.0
                     skip_sampling = True
                 elif ahead_idx >= self.n_ahead - 1:
-                    if labels is not None:  # we're in the talk phase
+                    if labels is not None:
                         cur_talk_n = ahead_idx - (self.n_ahead - 1) + 1
-                        # print("Setting rm to labels", cur_talk_n, "during", ahead_idx)
                         shift_labels = labels[..., cur_talk_n:].contiguous().to(probabilities_2d.device)
                         padding = torch.full_like(
                             labels[..., :cur_talk_n],
@@ -1640,11 +1596,9 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                             [shift_labels, padding],
                             dim=-1
                         )
-                        
-                        # print((new_rm_tokens > self.vocab_size - 1).any().item())
+
                         new_rm_tokens = torch.clamp(new_rm_tokens, 0, self.vocab_size - 1)
 
-                        # Now safely convert rm tokens to one-hot
                         probabilities_2d = F.one_hot(new_rm_tokens, num_classes=self.vocab_size).reshape(-1, self.vocab_size).to(probabilities_2d.dtype)
                     else:
                         continue
@@ -1704,7 +1658,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                         new_attention = original_attention
                     else:
                         original_attention = original_attention == attention_mask.max()
-                        # because eye isn't implemented for BF16, we need to handle the case
                         if not attention_mask.dtype == torch.bfloat16:
                             new_attention = torch.eye(
                                 seq_len, dtype=attention_mask.dtype, device=attention_mask.device
@@ -1742,9 +1695,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                     # if shift_labels.min() == self.tokenizer.pad_token_id:
                     shift_labels = torch.where(shift_labels == self.tokenizer.pad_token_id, -100, shift_labels)
                     unreduced_loss = loss_fct(shift_logits, shift_labels)
-                    # print("Loss:", unreduced_loss.item())  # Print the loss before checking for NaN values
                     if torch.any(unreduced_loss != unreduced_loss):
-                        # pdb.set_trace()
                         raise ValueError("NaN loss")
                     unreduced_loss = unreduced_loss.reshape(logits.shape[0], -1)
                     loss_list.append(unreduced_loss)