fix four files

config.json

@@ -6,7 +6,8 @@
     "auto_map": {
         "AutoConfig": "configuration_bailing_moe.BailingMoeConfig",
         "AutoModel": "modeling_bailing_moe.BailingMoeModel",
-        "AutoModelForCausalLM": "modeling_bailing_moe.BailingMoeForCausalLM"
+        "AutoModelForCausalLM": "modeling_bailing_moe.BailingMoeForCausalLM",
+        "AutoModelForTokenClassification": "modeling_bailing_moe.BailingMoeForTokenClassification"
     },
     "eos_token_id": 126081,
     "pad_token_id": 126081,
@@ -40,6 +41,5 @@
     "embedding_dropout": 0.1,
     "norm_head": true,
     "norm_softmax": false,
-    "output_dropout": 0.1,
-    "head_dim": 0
+    "output_dropout": 0.1
 }

modeling_bailing_moe.py

@@ -24,7 +24,6 @@ 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
 
@@ -72,6 +71,81 @@ logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "BailingMoeConfig"
 
+# Copied from transformers.models.mixtral.modeling_mixtral.load_balancing_loss_func
+def load_balancing_loss_func(
+    gate_logits_and_topk: Union[torch.Tensor, Tuple[torch.Tensor], None],
+    num_experts: Optional[int] = None,
+    top_k=2,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, int]:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits_and_topk is None or not isinstance(gate_logits_and_topk, tuple):
+        return 0
+
+    if isinstance(gate_logits_and_topk, tuple):
+        # concatenated_gate_logits.shape = [batch_size * num_layers * seq_len, num_experts]
+        concatenated_gate_logits = torch.cat([layer_gate[0] for layer_gate in gate_logits_and_topk], dim=0)
+        # selected_experts.shape = [batch_size * num_layers * seq_len, top_k_experts]
+        selected_experts = torch.cat([layer_gate[1] for layer_gate in gate_logits_and_topk], dim=0)
+        selected_experts.to(concatenated_gate_logits.device)
+        
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+    
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+    if attention_mask is None:
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+        )
+
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+            .reshape(-1, num_experts)
+        )
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(router_per_expert_attention_mask, dim=0)
+        
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss
 
 def _get_unpad_data(attention_mask):
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
@@ -421,7 +495,7 @@ class BailingMoeSparseMoeBlock(nn.Module):
             y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(bsz, seq_len, h)
         if self.config.num_shared_experts is not None:
             y = y + self.shared_experts(identity)
-        return y, (router_logits.view(bsz, seq_len, -1), topk_idx.view(bsz, seq_len, -1))
+        return y, (router_logits, topk_idx)
 
     @torch.no_grad()
     def moe_infer(self, x, topk_ids, topk_weight):
@@ -1363,21 +1437,12 @@ class BailingMoeForCausalLM(BailingMoePreTrainedModel):
 
     def compute_logit(self, hidden_states):
         if self.norm_head:
-            if self.training:
-                norm_weight = (
-                    self.lm_head.weight / (torch.norm(self.lm_head.weight, p=2, dim=0, keepdim=True) + 1e-7).detach()
-                )
-                logits = F.linear(hidden_states, norm_weight, None)
-            else:
-                self.lm_head.weight.data = (
-                    self.lm_head.weight.data.float()
-                    / (torch.norm(self.lm_head.weight.data.float(), p=2, dim=0, keepdim=True) + 1e-7)
-                ).to(hidden_states.dtype)
-                logits = F.linear(hidden_states, self.lm_head.weight.data, None)
-                self.norm_head = False
+            weight_float = self.lm_head.weight.float()
+            norm = torch.norm(weight_float, p=2, dim=0, keepdim=True).clamp(min=1e-7)
+            norm_weight = (weight_float / norm).to(hidden_states.dtype)
+            logits = F.linear(hidden_states, norm_weight, None)
         else:
             logits = self.lm_head(hidden_states)
-        return logits
 
     @add_start_docstrings_to_model_forward(BAILINGMOE_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
@@ -1452,6 +1517,14 @@ class BailingMoeForCausalLM(BailingMoePreTrainedModel):
         loss = None
         aux_loss = None
 
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits if return_dict else outputs[-1],
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+
         if labels is not None:
             # Shift so that tokens < n predict n
             shift_logits = logits[..., :-1, :].contiguous()
@@ -1547,3 +1620,107 @@ class BailingMoeForCausalLM(BailingMoePreTrainedModel):
                 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
             )
         return reordered_past
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaForTokenClassification with Llama->Qwen2Moe, LLAMA->QWEN2MOE
+class BailingMoeForTokenClassification(BailingMoePreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        
+        self.model = BailingMoeModel(config)
+        if getattr(config, "classifier_dropout", None) is not None:
+            classifier_dropout = config.classifier_dropout
+        elif getattr(config, "hidden_dropout", None) is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.score = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(BAILINGMOE_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+
+        outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            return_dict=return_dict,
+        )
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        logits = self.score(sequence_output)
+
+        loss = None
+        aux_loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config)
+        
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits if return_dict else outputs[-1],
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            if output_router_logits:
+                output = (aux_loss,) + output
+            return (loss,) + output if loss is not None else output
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )

special_tokens_map.json

@@ -1,10 +1,15 @@
 {
+    "additional_special_tokens": [
+        "<|arithmetic_start|>",
+        "<|arithmetic_end|>",
+        "<role>",
+        "</role>",
+        "<|number_end|>",
+        "<|number_start|>"
+    ],
     "bos_token": "<|startoftext|>",
     "cls_token": "[CLS]",
     "eos_token": "<|endoftext|>",
     "gmask_token": "[gMASK]",
-    "additional_special_tokens": [
-        "<role>",
-        "</role>"
-    ]
+    "pad_token": "<|endoftext|>"
 }

tokenizer_config.json

@@ -1,15 +1,25 @@
 {
     "add_bos_token": false,
     "add_eos_token": false,
+    "additional_special_tokens": [
+        "<role>",
+        "</role>",
+        "<|arithmetic_start|>",
+        "<|arithmetic_end|>",
+        "<|number_start|>",
+        "<|number_end|>"
+    ],
     "bos_token": "<|startoftext|>",
+    "chat_template": "{% for message in messages %}{% set role = message['role'] | lower %}{% if role == 'user' %}{% set role = 'HUMAN' %}{% endif %}{% set role = role | upper %}{{ '<role>' + role + '</role>' + message['content'] }}{% if role == 'ASSISTANT' %}{{ eos_token }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<role>ASSISTANT</role>' }}{% endif %}",
     "clean_up_tokenization_spaces": false,
     "cls_token": "[CLS]",
     "eos_token": "<|endoftext|>",
     "gmask_token": "[gMASK]",
     "merges_file": null,
     "model_max_length": 1000000000000000019884624838656,
+    "pad_token": "<|endoftext|>",
     "tokenizer_class": "PreTrainedTokenizerFast",
+    "trust_remote_code": true,
     "vocab_file": null,
-    "pad_token": "<|endoftext|>",
     "fast_tokenizer": true
 }