Update modeling_motif.py

modeling_motif.py

@@ -1040,13 +1040,12 @@ class MotifModel(MotifPreTrainedModel):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
-        self.multi_token_heads = config.multi_token_heads
 
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
         # NOTE: For multi-token models, the last decoder layers (one for each token index)
         # are implemented as a part of `MotifModelForCausalLM` to enable a custom forward-backward procedure.
 
-        num_hidden_layers = config.num_hidden_layers if self.multi_token_heads is None else config.num_hidden_layers - 1
+        num_hidden_layers = config.num_hidden_layers
         self.layers = nn.ModuleList([MotifDecoderLayer(config = config, layer_idx=layer_idx) for layer_idx in range(num_hidden_layers)])
         self._attn_implementation = config._attn_implementation
         RMSNorm = MorehRMSNorm
@@ -1338,16 +1337,8 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
         super().__init__(config)
         self.model = MotifModel(config)
         self.vocab_size = config.vocab_size
-        self.multi_token_heads = config.multi_token_heads
 
-        if self.multi_token_heads is None:
-            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        else:
-            self.tokenwise_last_layers = nn.ModuleList(
-                [MotifDecoderLayer(config, config.num_hidden_layers - 1) for _ in range(self.multi_token_heads)])
-            self.tokenwise_lm_heads = nn.ModuleList(
-                [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(self.multi_token_heads)])
-        self.should_skip_separate_backward_pass = self.multi_token_heads is not None
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -1374,101 +1365,7 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
     def get_decoder(self):
         return self.model
 
-    def multi_token_forward_backward(self,
-                                     hidden_states: torch.FloatTensor,
-                                     outputs: MotifModelOutputWithPast,
-                                     labels: torch.LongTensor,
-                                     position_ids: Optional[torch.LongTensor],
-                                     output_attentions: Optional[bool],
-                                     use_cache: Optional[bool],
-                                     cache_position: Optional[torch.LongTensor],
-                                     return_dict: Optional[bool],
-                                     num_logits_to_keep: int = 0) -> CausalLMOutputWithPast:
-        """ 
-        This implements the main forward-backward procedure for multi-token model training proposed in 
-        the paper https://arxiv.org/abs/2404.19737. 
-        Essentially, 
-        - The multi-token model tries to predict n (instead of 1) tokens at a time.
-        - Applying this only during training and using first-token prediction during inference is still helpful.
-        - The change in architecture: when using n-token prediction, each token index (between 1 and n) has its own 
-            (1) last attention layer and (2) lm head.
-        - The change in loss: sum of cross-entropy losses corresponding to each token index. 
-        - Custom forward-backward procedure for memory efficiency: refer to the implementation of `multi_head_forward_backward`.
-        """
-        if not return_dict:
-            raise NotImplementedError("return_dict must be True for multi-token training")
-
-        past_key_values = outputs.past_key_values
-        causal_mask = outputs.causal_mask
-        position_embeddings = outputs.position_embeddings
-
-        if labels is not None:
-            labels = labels.to(hidden_states.device)
-
-        def _tokenwise_forward(hidden_states: torch.Tensor, token_idx):
-            ## Model forward
-            layer = self.tokenwise_last_layers[token_idx]
-            lm_head = self.tokenwise_lm_heads[token_idx]
-
-            layer_outputs = layer(
-                hidden_states,
-                attention_mask=causal_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,  # TODO: update past_key_values?
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-            )
-            last_hidden_states = layer_outputs[0]
-            if num_logits_to_keep > 0:
-                assert labels is None
-                last_hidden_states = last_hidden_states[:, -num_logits_to_keep:, :]
-            tokenwise_logits = lm_head(last_hidden_states)
-
-            if labels is None:
-                return {
-                    "loss": None,
-                    "logits": tokenwise_logits,
-                }
-
-            ## Compute loss
-            shift_n = token_idx + 1
-            shift_logits = tokenwise_logits[..., :-shift_n, :].contiguous()
-            shift_labels = labels[..., shift_n:].contiguous()
-
-            loss_fct = CrossEntropyLoss()
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
-
-            tokenwise_loss = loss_fct(shift_logits, shift_labels)
-
-            return {
-                "loss": tokenwise_loss,
-                "logits": tokenwise_logits,
-            }
-
-        head_fns = [
-            lambda hidden_states, token_idx=token_idx: _tokenwise_forward(hidden_states, token_idx)
-            for token_idx in range(self.multi_token_heads)
-        ]
-        loss, logits = multi_head_forward_backward(hidden_states,
-                                                   head_fns,
-                                                   return_keys=("loss", "logits"),
-                                                   return_only_first_head=True)
-
-        if not return_dict:
-            output = (logits, ) + outputs[1:]
-            return (loss, ) + output
-
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
+    
     @add_start_docstrings_to_model_forward(MOTIF_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(
@@ -1524,8 +1421,6 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs_include_causal_mask = self.multi_token_heads is not None
-        outputs_include_position_embeddings = self.multi_token_heads is not None
         outputs: MotifModelOutputWithPast = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -1537,23 +1432,10 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
             cache_position=cache_position,
-            outputs_include_causal_mask=outputs_include_causal_mask,
-            outputs_include_position_embeddings=outputs_include_position_embeddings,
         )
 
         hidden_states = outputs[0]
 
-        if self.multi_token_heads is not None:
-            return self.multi_token_forward_backward(hidden_states,
-                                                     outputs,
-                                                     labels,
-                                                     position_ids,
-                                                     output_attentions,
-                                                     use_cache,
-                                                     cache_position,
-                                                     return_dict,
-                                                     num_logits_to_keep=num_logits_to_keep)
-
         # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
         hidden_states = hidden_states
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])