| from typing import List, Optional |
|
|
| import torch |
| import torch.nn as nn |
| import torch.distributed as dist |
|
|
| from transformers.modeling_utils import PreTrainedModel |
| from transformers.utils import logging |
| try: |
| from .configuration_mimo_v2_pro import MiMoV2ProConfig |
| except ImportError: |
| from configuration_mimo_v2_pro import MiMoV2ProConfig |
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| class FakeRMSNorm(nn.Module): |
| def __init__(self, hidden_size, eps=1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.empty(hidden_size)) |
|
|
| def forward(self, x): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class FakeMLP(nn.Module): |
| def __init__(self, config: MiMoV2ProConfig, is_expert=False): |
| super().__init__() |
| if is_expert: |
| self.intermediate_size = config.moe_intermediate_size |
| else: |
| self.intermediate_size = config.intermediate_size |
| self.hidden_size = config.hidden_size |
|
|
| self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) |
| self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) |
|
|
| def forward(self, x): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class FakeMoEGate(nn.Module): |
| def __init__(self, config: MiMoV2ProConfig): |
| super().__init__() |
| self.weight = nn.Parameter( |
| torch.empty((config.n_routed_experts, config.hidden_size), dtype=torch.float32) |
| ) |
| if config.topk_method == "noaux_tc": |
| self.e_score_correction_bias = nn.Parameter( |
| torch.empty((config.n_routed_experts), dtype=torch.float32) |
| ) |
|
|
| def forward(self, hidden_states): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class FakeMoE(nn.Module): |
| def __init__(self, config: MiMoV2ProConfig): |
| super().__init__() |
| self.experts = nn.ModuleList( |
| [FakeMLP(config, is_expert=True) for _ in range(config.n_routed_experts)] |
| ) |
| self.gate = FakeMoEGate(config) |
|
|
| def forward(self, hidden_states): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class FakeFullAttention(nn.Module): |
| """Fake Attention for full attention layers (hybrid_layer_pattern == 0).""" |
|
|
| def __init__(self, config: MiMoV2ProConfig): |
| super().__init__() |
| q_size = config.num_attention_heads * config.head_dim |
| k_size = config.num_key_value_heads * config.head_dim |
| v_size = config.num_key_value_heads * config.v_head_dim |
| o_size = config.num_attention_heads * config.v_head_dim |
|
|
| self.qkv_proj = nn.Linear( |
| config.hidden_size, q_size + k_size + v_size, bias=config.attention_bias |
| ) |
| self.o_proj = nn.Linear(o_size, config.hidden_size, bias=False) |
|
|
| if getattr(config, "add_full_attention_sink_bias", False): |
| self.attention_sink_bias = nn.Parameter( |
| torch.empty(config.num_attention_heads) |
| ) |
|
|
| def forward(self, *args, **kwargs): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class FakeSWAAttention(nn.Module): |
| """Fake Attention for SWA layers (hybrid_layer_pattern == 1).""" |
|
|
| def __init__(self, config: MiMoV2ProConfig): |
| super().__init__() |
| q_size = config.swa_num_attention_heads * config.swa_head_dim |
| k_size = config.swa_num_key_value_heads * config.swa_head_dim |
| v_size = config.swa_num_key_value_heads * config.swa_v_head_dim |
| o_size = config.swa_num_attention_heads * config.swa_v_head_dim |
|
|
| self.qkv_proj = nn.Linear( |
| config.hidden_size, q_size + k_size + v_size, bias=config.attention_bias |
| ) |
| self.o_proj = nn.Linear(o_size, config.hidden_size, bias=False) |
|
|
| if getattr(config, "add_swa_attention_sink_bias", False): |
| self.attention_sink_bias = nn.Parameter( |
| torch.empty(config.swa_num_attention_heads) |
| ) |
|
|
| def forward(self, *args, **kwargs): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class FakeDecoderLayer(nn.Module): |
| def __init__(self, config: MiMoV2ProConfig, layer_idx: int): |
| super().__init__() |
| self.input_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| self.post_attention_layernorm = FakeRMSNorm( |
| config.hidden_size, eps=config.layernorm_epsilon |
| ) |
|
|
| is_moe = ( |
| getattr(config, "n_routed_experts", None) is not None |
| and config.moe_layer_freq[layer_idx] |
| ) |
| self.mlp = FakeMoE(config) if is_moe else FakeMLP(config) |
|
|
| is_swa = config.hybrid_layer_pattern[layer_idx] == 1 |
| self.self_attn = FakeSWAAttention(config) if is_swa else FakeFullAttention(config) |
|
|
| def forward(self, *args, **kwargs): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class FakeMTPLayer(nn.Module): |
| def __init__(self, config: MiMoV2ProConfig): |
| super().__init__() |
| self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) |
| self.enorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| self.hnorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| self.final_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| self.input_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| self.pre_mlp_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| self.self_attn = FakeSWAAttention(config) |
| self.mlp = FakeMLP(config) |
|
|
| def forward(self, *args, **kwargs): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class MiMoV2ProModel(PreTrainedModel): |
| config_class = MiMoV2ProConfig |
|
|
| def __init__(self, config: MiMoV2ProConfig): |
| super().__init__(config) |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList( |
| [FakeDecoderLayer(config, i) for i in range(config.num_hidden_layers)] |
| ) |
| self.norm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| self.mtp = nn.ModuleDict( |
| {"layers": nn.ModuleList([FakeMTPLayer(config)])} |
| ) |
|
|
| def forward(self, *args, **kwargs): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
|
|
|
|
| class MiMoV2ProForCausalLM(PreTrainedModel): |
| config_class = MiMoV2ProConfig |
| _keys_to_ignore_on_load_unexpected = [ |
| r"model.layers\.\d+\.self_attn\.rotary_emb\.inv_freq" |
| ] |
|
|
| def __init__(self, config: MiMoV2ProConfig): |
| super().__init__(config) |
| self.model = MiMoV2ProModel(config) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
|
|
| def forward(self, *args, **kwargs): |
| raise NotImplementedError("This is a fake model and does not support forward pass.") |
