fix model file and script

chat_template.jinja

@@ -0,0 +1 @@
+{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|role_start|>system<|role_end|>你是megrez-3x7b-a3b-instruct<|turn_end|>' }}{% endif %}{{ '<|role_start|>' + message['role'] + '<|role_end|>' + message['content'] + '<|turn_end|>' }}{% endfor %}{% if add_generation_prompt %}{{ '<|role_start|>assistant<|role_end|>' }}{% endif %}

config.json

@@ -29,7 +29,7 @@
   "norm_topk_prob": true,
   "num_attention_heads": 16,
   "num_experts_per_tok": 6,
-  "num_hidden_layers": 11,
+  "num_hidden_layers": 31,
   "num_key_value_heads": 4,
   "pad_token_id": 120002,
   "pre_gate": true,

model-00001-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a7d2425fd5f1de10e9dfef4107cd5a23e077adb5dbf72e3650b1cbfcba991d0
+size 4994539224

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:510db202ee9e0172fdd3ab71fd7ed3ade9b78f9dae68f576a8712394e1eedf39
+size 4995331096

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2bca00397338a1e77781693e9afa93625845c316f0930cb6499cfb4be60719c
+size 4958908880

modeling_megrez_moe.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved.
+# Copyright 2025 Infini-AI and The HuggingFace Inc. team. All rights reserved.
 #
 # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
 # and OPT implementations in this library. It has been modified from its
@@ -17,53 +17,30 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-""" PyTorch DeepSeek model."""
+"""PyTorch Megrez model."""
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
-import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
-from transformers.modeling_attn_mask_utils import (
-    AttentionMaskConverter,
-    _prepare_4d_attention_mask,
-    _prepare_4d_causal_attention_mask,
-)
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-    SequenceClassifierOutputWithPast,
-)
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+from transformers.modeling_outputs import (BaseModelOutputWithPast, CausalLMOutputWithPast,
+                                           SequenceClassifierOutputWithPast)
 from transformers.modeling_utils import PreTrainedModel
-from transformers.pytorch_utils import (
-    ALL_LAYERNORM_LAYERS,
-    is_torch_greater_or_equal_than_1_13,
-)
-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,
-)
-from transformers.utils.import_utils import is_torch_fx_available
-from .configuration_megrez_moe import MegrezMoeConfig
-import torch.distributed as dist
-import numpy as np
-
 from transformers.models.llama.modeling_llama import LlamaAttention, LlamaRotaryEmbedding
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
+from transformers.utils import (add_start_docstrings, add_start_docstrings_to_model_forward, logging,
+                                replace_return_docstrings)
+from transformers.utils.import_utils import is_torch_fx_available
 
-if is_flash_attn_2_available():
-    from flash_attn import flash_attn_func, flash_attn_varlen_func
-    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
-
+from .configuration_megrez_moe import MegrezMoeConfig
 
 # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
 # It means that the function will not be traced through and simply appear as a node in the graph.
@@ -79,18 +56,6 @@ logger = logging.get_logger(__name__)
 _CONFIG_FOR_DOC = "MegrezMoeConfig"
 
 
-def _get_unpad_data(attention_mask):
-    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    max_seqlen_in_batch = seqlens_in_batch.max().item()
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
-    return (
-        indices,
-        cu_seqlens,
-        max_seqlen_in_batch,
-    )
-
-
 class MegrezMoeRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
         """
@@ -214,8 +179,9 @@ class MegrezMoeMoE(nn.Module):
     A mixed expert module containing shared experts.
     """
 
-    def __init__(self, config, init_experts: bool = True):
+    def __init__(self, config, layer_number, init_experts: bool = True):
         super().__init__()
+        self.layer_number = layer_number
         self.config = config
         self.num_experts_per_tok = config.num_experts_per_tok
 
@@ -267,7 +233,6 @@ class MegrezMoeMoE(nn.Module):
             topk_idx, topk_weight = self.gate(pre_gate_hidden_states)
         else:
             topk_idx, topk_weight = self.gate(hidden_states)
-
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
         flat_topk_idx = topk_idx.view(-1)
         if self.training:
@@ -280,7 +245,9 @@ class MegrezMoeMoE(nn.Module):
         else:
             y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
-            y = y + self.shared_experts(identity)
+            shared_out = self.shared_experts(identity)
+            y = y + shared_out
+            # y = y + self.shared_experts(identity)
         return y
 
     @torch.no_grad()
@@ -369,7 +336,9 @@ class MegrezMoeDecoderLayer(nn.Module):
         self.config = config
         self.layer_number = layer_idx
 
-        self.experts_shared = config.experts_shared_frequency is not None and layer_idx >= self.config.first_k_dense_replace
+        self.experts_shared = (
+            config.experts_shared_frequency is not None and layer_idx >= self.config.first_k_dense_replace
+        )
 
         self.pre_gate = config.pre_gate
 
@@ -381,43 +350,11 @@ class MegrezMoeDecoderLayer(nn.Module):
             and layer_idx % config.moe_layer_freq == 0
         )
 
-        if self.experts_shared:
-            assert config.moe_layer_freq == 1
-            self.self_attn = torch.nn.ModuleList(
-                [LlamaAttention(config=config, layer_idx=layer_idx) for _ in range(config.experts_shared_frequency)]
-            )
-            for idx in range(config.experts_shared_frequency):
-                attn_layer_idx = (
-                    layer_idx - self.config.first_k_dense_replace
-                ) * self.config.experts_shared_frequency + self.config.first_k_dense_replace + idx
-                self.self_attn[idx].layer_idx = attn_layer_idx
-            self.input_layernorm = torch.nn.ModuleList(
-                [
-                    MegrezMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-                    for _ in range(config.experts_shared_frequency)
-                ]
-            )
-            self.post_attention_layernorm = torch.nn.ModuleList(
-                [
-                    MegrezMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-                    for _ in range(config.experts_shared_frequency)
-                ]
-            )
-
-            mlp = [MegrezMoeMoE(config, init_experts=True)]
-            for _ in range(1, config.experts_shared_frequency):
-                layer = MegrezMoeMoE(config, init_experts=False)
-                # layer.set_experts(mlp[0].experts)
-                mlp.append(layer)
-            self.mlp = torch.nn.ModuleList(mlp)
-
-        else:
-
-            self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
-
-            self.mlp = MegrezMoeMoE(config) if is_moe else MegrezMoeMLP(config)
-            self.input_layernorm = MegrezMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-            self.post_attention_layernorm = MegrezMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        init_experts = (layer_idx - config.first_k_dense_replace) % config.experts_shared_frequency == 0
+        self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
+        self.mlp = MegrezMoeMoE(config, layer_idx, init_experts) if is_moe else MegrezMoeMLP(config)
+        self.input_layernorm = MegrezMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = MegrezMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
         self,
@@ -457,63 +394,32 @@ class MegrezMoeDecoderLayer(nn.Module):
                 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
             )
 
-        if self.experts_shared:
-
-            for idx in range(self.config.experts_shared_frequency):
-                residual = hidden_states
-
-                hidden_states = self.input_layernorm[idx](hidden_states)
-
-                # Self Attention
-                hidden_states, self_attn_weights = self.self_attn[idx](
-                    hidden_states=hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_value,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                    position_embeddings=position_embeddings,
-                    **kwargs,
-                )
-                hidden_states = residual + hidden_states
-
-                # Fully Connected
-                residual = hidden_states
-                hidden_states = self.post_attention_layernorm[idx](hidden_states)
-                post_attention_layernorm_hidden_states = hidden_states
-                if idx > 0:
-                    self.mlp[idx].set_experts(self.mlp[0].experts)
-                hidden_states = self.mlp[idx](hidden_states, pre_gate_hidden_states=pre_gate_hidden_states)
-                if idx > 0:
-                    self.mlp[idx].set_experts(None)
-                hidden_states = residual + hidden_states
-                pre_gate_hidden_states = post_attention_layernorm_hidden_states
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
 
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        post_attention_layernorm_hidden_states = hidden_states
+        if isinstance(self.mlp, MegrezMoeMoE):
+            hidden_states = self.mlp(hidden_states, pre_gate_hidden_states=pre_gate_hidden_states)
         else:
-            residual = hidden_states
-
-            hidden_states = self.input_layernorm(hidden_states)
-
-            # Self Attention
-            hidden_states, self_attn_weights = self.self_attn(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                position_embeddings=position_embeddings,
-                **kwargs,
-            )
-            hidden_states = residual + hidden_states
-
-            # Fully Connected
-            residual = hidden_states
-            hidden_states = self.post_attention_layernorm(hidden_states)
-            post_attention_layernorm_hidden_states = hidden_states
             hidden_states = self.mlp(hidden_states)
-            hidden_states = residual + hidden_states
-            pre_gate_hidden_states = post_attention_layernorm_hidden_states
+        hidden_states = residual + hidden_states
+        pre_gate_hidden_states = post_attention_layernorm_hidden_states
 
         if self.pre_gate and self.layer_number < self.config.num_hidden_layers - 1:
             hidden_states = torch.cat([pre_gate_hidden_states, hidden_states], dim=0)
@@ -728,7 +634,6 @@ class MegrezMoeModel(MegrezMoePreTrainedModel):
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
-
         if self._use_flash_attention_2:
             # 2d mask is passed through the layers
             attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
@@ -750,11 +655,19 @@ class MegrezMoeModel(MegrezMoePreTrainedModel):
         next_decoder_cache = None
 
         position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
-
-        for decoder_layer in self.layers:
+        for layer_idx, decoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
+            shared_layer_idx = (
+                (layer_idx - self.config.first_k_dense_replace)
+                // self.config.experts_shared_frequency
+                * self.config.experts_shared_frequency
+                + self.config.first_k_dense_replace
+            )
+            if layer_idx >= self.config.first_k_dense_replace and shared_layer_idx != layer_idx:
+                decoder_layer.mlp.set_experts(self.layers[shared_layer_idx].mlp.experts)
+
             if self.gradient_checkpointing and self.training:
                 layer_outputs = self._gradient_checkpointing_func(
                     decoder_layer.__call__,
@@ -778,14 +691,14 @@ class MegrezMoeModel(MegrezMoePreTrainedModel):
                     position_embeddings=position_embeddings,
                     **flash_attn_kwargs,
                 )
-
+            if layer_idx >= self.config.first_k_dense_replace and shared_layer_idx != layer_idx:
+                decoder_layer.mlp.set_experts(None)
             hidden_states = layer_outputs[0]
 
             if output_attentions:
                 all_self_attns += (layer_outputs[1],)
 
         hidden_states = self.norm(hidden_states)
-
         # add hidden states from the last decoder layer
         if output_hidden_states:
             all_hidden_states += (hidden_states,)