Upload modeling_bailing_moe_v2.py with huggingface_hub

modeling_bailing_moe_v2.py

@@ -25,9 +25,7 @@ from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
-import torch.utils.checkpoint
 from torch import nn
-from torch.nn import CrossEntropyLoss
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
@@ -1157,11 +1155,11 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
-        self.num_mtp_layers = config.num_mtp_layers
+        self.num_nextn_predict_layers = config.num_nextn_predict_layers
 
         self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
         self.layers = []
-        for layer_idx in range(config.num_hidden_layers + config.num_mtp_layers):
+        for layer_idx in range(config.num_hidden_layers + config.num_nextn_predict_layers):
             layer_cls = BailingMoeV2DecoderLayer if layer_idx < config.num_hidden_layers else BailingMoeV2MTPLayer
             self.layers.append(layer_cls(config, layer_idx))
 
@@ -1267,8 +1265,8 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
         all_self_attns = () if output_attentions else None
         all_router_logits = () if output_router_logits else None
         next_decoder_cache = None
-        layers = self.layers[: -self.num_mtp_layers] if self.num_mtp_layers > 0 else self.layers
-        mtp_layers = self.layers[-self.num_mtp_layers :] if self.num_mtp_layers > 0 else None
+        layers = self.layers[: -self.num_nextn_predict_layers] if self.num_nextn_predict_layers > 0 else self.layers
+        mtp_layers = self.layers[-self.num_nextn_predict_layers :] if self.num_nextn_predict_layers > 0 else None
 
         for decoder_layer in layers:
             if output_hidden_states:
@@ -1391,7 +1389,7 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
         self.model = BailingMoeV2Model(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.num_mtp_layers = config.num_mtp_layers
+        self.num_nextn_predict_layers = config.num_nextn_predict_layers
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
 
         # Initialize weights and apply final processing
@@ -1491,18 +1489,21 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
             loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
 
         all_mtp_logits = None
-        if self.num_mtp_layers > 0:
+        if self.num_nextn_predict_layers > 0:
             mtp_hidden_states = outputs.mtp_hidden_states
-            shift_labels_mtp = labels.clone()
-            for i in range(self.num_mtp_layers):
+            shift_labels_mtp = None
+            for i in range(self.num_nextn_predict_layers):
                 mtp_hidden_states = mtp_hidden_states[i]
                 mtp_logits = self.lm_head(mtp_hidden_states).float()
                 if all_mtp_logits is None:
                     all_mtp_logits = []
                 all_mtp_logits.append(mtp_logits)
                 if labels is not None:
+                    if shift_labels_mtp is None:
+                        shift_labels_mtp = labels.clone()
                     shift_labels_mtp, _ = roll_tensor(shift_labels_mtp, shifts=-1, dims=-1, fill_value=-100)
-                    mtp_loss = self.loss_function(mtp_logits, shift_labels_mtp, self.config.vocab_size, **kwargs)
+                    mtp_logits_ = mtp_logits.view(-1, self.config.vocab_size)
+                    mtp_loss = self.loss_function(mtp_logits_, shift_labels_mtp.to(mtp_logits_.device).view(-1), self.config.vocab_size, **kwargs)
                     if loss is not None:
                         loss += self.mtp_loss_scaling_factor * mtp_loss
                     else:
@@ -1529,3 +1530,4 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
             attentions=outputs.attentions,
             router_logits=outputs.router_logits,
         )
+