| |
| |
| |
| |
|
|
| """PyTorch implementation of the MiniMax M2 architecture for Hugging Face Transformers.""" |
|
|
| from __future__ import annotations |
|
|
| import copy |
| import time |
| from typing import Optional, Tuple, Union |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from transformers.activations import ACT2FN |
| from transformers.cache_utils import Cache, DynamicCache |
| from transformers.generation import GenerationMixin |
| from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask |
| from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast |
| from transformers.modeling_utils import PreTrainedModel |
| from transformers.utils import logging |
|
|
| from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, repeat_kv, rotate_half |
|
|
| from .configuration_minimax_m2 import MiniMaxM2Config |
|
|
| logger = logging.get_logger(__name__) |
|
|
| _CONFIG_FOR_DOC = "MiniMaxM2Config" |
| _CHECKPOINT_FOR_DOC = "MiniMaxAI/MiniMax-M2" |
|
|
|
|
| def load_balancing_loss_func( |
| gate_logits: Union[torch.Tensor, Tuple[torch.Tensor, ...]], |
| num_experts: int, |
| top_k: int, |
| attention_mask: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| if gate_logits is None: |
| return torch.tensor(0.0) |
| if isinstance(gate_logits, torch.Tensor): |
| logits = gate_logits |
| else: |
| logits = torch.cat([layer_gate.to(gate_logits[0].device) for layer_gate in gate_logits], dim=0) |
|
|
| routing_weights = torch.softmax(logits, dim=-1, dtype=torch.float32) |
| _, selected = torch.topk(routing_weights, top_k, dim=-1) |
| expert_mask = torch.nn.functional.one_hot(selected, num_experts) |
|
|
| if attention_mask is None: |
| tokens_per_expert = torch.mean(expert_mask.float(), dim=0) |
| router_prob_per_expert = torch.mean(routing_weights, dim=0) |
| else: |
| batch_size, seq_len = attention_mask.shape |
| num_layers = logits.shape[0] // (batch_size * seq_len) |
|
|
| expanded_mask = ( |
| attention_mask[None, :, :, None, None] |
| .expand(num_layers, batch_size, seq_len, top_k, num_experts) |
| .reshape(-1, top_k, num_experts) |
| .to(logits.device) |
| ) |
| tokens_per_expert = torch.sum(expert_mask.float() * expanded_mask, dim=0) / torch.sum(expanded_mask, dim=0) |
|
|
| router_mask = ( |
| attention_mask[None, :, :, None] |
| .expand(num_layers, batch_size, seq_len, num_experts) |
| .reshape(-1, num_experts) |
| .to(logits.device) |
| ) |
| router_prob_per_expert = torch.sum(routing_weights * router_mask, dim=0) / torch.sum(router_mask, dim=0) |
|
|
| loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) |
| return loss * num_experts |
|
|
|
|
| def apply_rotary_pos_emb_partial( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| cos: torch.Tensor, |
| sin: torch.Tensor, |
| rotary_dim: int, |
| unsqueeze_dim: int = 2, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| cos = cos.unsqueeze(unsqueeze_dim)[..., :rotary_dim] |
| sin = sin.unsqueeze(unsqueeze_dim)[..., :rotary_dim] |
| q_rot = q[..., :rotary_dim] |
| k_rot = k[..., :rotary_dim] |
|
|
| q_rot = (q_rot * cos) + (rotate_half(q_rot) * sin) |
| k_rot = (k_rot * cos) + (rotate_half(k_rot) * sin) |
|
|
| q = torch.cat((q_rot, q[..., rotary_dim:]), dim=-1) |
| k = torch.cat((k_rot, k[..., rotary_dim:]), dim=-1) |
| return q, k |
|
|
|
|
| class MiniMaxM2RMSNorm(nn.Module): |
| def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.variance_epsilon = eps |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| input_dtype = hidden_states.dtype |
| hidden_states = hidden_states.to(torch.float32) |
| variance = hidden_states.pow(2).mean(-1, keepdim=True) |
| hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
| return (self.weight * hidden_states).to(input_dtype) |
|
|
|
|
| class MiniMaxM2MLP(nn.Module): |
| def __init__(self, config: MiniMaxM2Config) -> None: |
| super().__init__() |
| self.hidden_size = config.hidden_size |
| self.intermediate_size = config.intermediate_size |
|
|
| self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) |
| self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) |
| self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) |
| self.act_fn = ACT2FN[config.hidden_act] |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| gate = self.act_fn(self.w1(hidden_states)) |
| up = self.w3(hidden_states) |
| hidden_states = gate * up |
| hidden_states = self.w2(hidden_states) |
| return hidden_states |
|
|
|
|
| class MiniMaxM2SparseMoeBlock(nn.Module): |
| def __init__(self, config: MiniMaxM2Config) -> None: |
| super().__init__() |
| self.hidden_dim = config.hidden_size |
| self.experts = nn.ModuleList([MiniMaxM2MLP(config) for _ in range(config.num_local_experts)]) |
| self.num_experts = config.num_local_experts |
| self.top_k = config.num_experts_per_tok |
| self.jitter_noise = config.router_jitter_noise |
| self.use_routing_bias = config.use_routing_bias |
| self.scoring_func = getattr(config, "scoring_func", "softmax") |
| self.use_grouped_topk = getattr(config, "use_grouped_topk", False) |
| self.num_expert_group = getattr(config, "num_expert_group", None) |
| self.topk_group = getattr(config, "topk_group", None) |
| self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) |
|
|
| if self.use_grouped_topk: |
| if self.num_expert_group is None or self.num_expert_group <= 0: |
| self.num_expert_group = 1 |
| if self.topk_group is None or self.topk_group <= 0: |
| self.topk_group = min(self.num_expert_group, self.top_k) |
| else: |
| self.num_expert_group = 1 |
| self.topk_group = 1 |
|
|
| self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False) |
| if self.use_routing_bias: |
| self.e_score_correction_bias = nn.Parameter(torch.zeros(self.num_experts, dtype=torch.float32)) |
| else: |
| self.register_parameter("e_score_correction_bias", None) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| batch_size, seq_len, hidden_dim = hidden_states.shape |
| if self.training and self.jitter_noise > 0: |
| noise = torch.empty_like(hidden_states).uniform_( |
| 1.0 - self.jitter_noise, |
| 1.0 + self.jitter_noise, |
| ) |
| hidden_states = hidden_states * noise |
|
|
| hidden_states = hidden_states.view(-1, hidden_dim) |
| gate_dtype = self.gate.weight.dtype |
| router_logits = self.gate(hidden_states.to(gate_dtype)).to(torch.float32) |
| if self.e_score_correction_bias is not None: |
| |
| correction_bias = self.e_score_correction_bias.to(router_logits.device, router_logits.dtype) |
| else: |
| correction_bias = None |
|
|
| if self.scoring_func == "sigmoid": |
| scores = torch.sigmoid(router_logits) |
| elif self.scoring_func == "softmax": |
| scores = torch.softmax(router_logits, dim=-1) |
| else: |
| raise ValueError(f"Unsupported scoring function: {self.scoring_func}") |
|
|
| if correction_bias is not None: |
| original_scores = scores |
| scores = scores + correction_bias |
| else: |
| original_scores = scores |
| topk_scores: torch.Tensor |
| if self.use_grouped_topk and self.num_expert_group > 1: |
| experts_per_group = scores.size(-1) // self.num_expert_group |
| scores_grouped = scores.view(scores.size(0), self.num_expert_group, experts_per_group) |
| if correction_bias is not None: |
| topk_in_group = min(2, experts_per_group) |
| if topk_in_group > 0: |
| group_scores = scores_grouped.topk(topk_in_group, dim=-1)[0].sum(dim=-1) |
| else: |
| group_scores = torch.zeros_like(scores_grouped[..., 0]) |
| else: |
| group_scores = scores_grouped.max(dim=-1).values |
| group_mask = torch.zeros_like(group_scores) |
| selected_groups = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=True).indices |
| group_mask.scatter_(1, selected_groups, 1.0) |
| mask = group_mask.unsqueeze(-1).expand(-1, -1, experts_per_group).reshape(scores.size()) |
| masked_scores = scores.masked_fill(mask == 0, float("-inf")) |
| topk_scores, selected_experts = torch.topk(masked_scores, self.top_k, dim=-1, sorted=True) |
| else: |
| topk_scores, selected_experts = torch.topk(scores, self.top_k, dim=-1, sorted=True) |
|
|
| if correction_bias is not None: |
| routing_weights = original_scores.gather(1, selected_experts) |
| else: |
| routing_weights = topk_scores |
|
|
| routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True).clamp(min=1e-12) |
| if self.routed_scaling_factor != 1.0: |
| routing_weights = routing_weights * self.routed_scaling_factor |
| routing_weights = routing_weights.to(hidden_states.dtype) |
| selected_experts = selected_experts.to(torch.long) |
|
|
| final_hidden_states = torch.zeros_like(hidden_states) |
| expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) |
| expert_hit = torch.nonzero(expert_mask.sum(dim=(-1, -2)) > 0, as_tuple=False).flatten() |
|
|
| for expert_idx in expert_hit.tolist(): |
| expert_layer = self.experts[expert_idx] |
| idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) |
| token_states = hidden_states.index_select(0, top_x) |
| expert_output = expert_layer(token_states) * routing_weights[top_x, idx].unsqueeze(-1) |
| final_hidden_states.index_add_(0, top_x, expert_output.to(final_hidden_states.dtype)) |
|
|
| final_hidden_states = final_hidden_states.view(batch_size, seq_len, hidden_dim) |
| return final_hidden_states, router_logits |
|
|
|
|
| class MiniMaxM2Attention(nn.Module): |
| def __init__(self, config: MiniMaxM2Config, layer_idx: int) -> None: |
| super().__init__() |
| self.config = config |
| self.layer_idx = layer_idx |
|
|
| self.head_dim = config.head_dim |
| self.num_heads = config.num_attention_heads |
| self.num_key_value_heads = config.num_key_value_heads |
| self.num_key_value_groups = self.num_heads // max(1, self.num_key_value_heads) |
| self.rotary_dim = config.rotary_dim |
| self.scaling = self.head_dim**-0.5 |
| self.attention_dropout = config.attention_dropout |
| self.is_causal = True |
|
|
| max_position_embeddings = getattr(config, "max_position_embeddings", 8192) |
| max_model_len = getattr(config, "max_model_len", None) |
| if max_model_len is not None: |
| max_position_embeddings = max(max_position_embeddings, max_model_len) |
|
|
| attn_window_size = getattr(config, "attn_window_size", None) |
| if isinstance(attn_window_size, list): |
| sliding_window = attn_window_size[layer_idx] |
| else: |
| sliding_window = attn_window_size |
| if sliding_window is not None and sliding_window <= 0: |
| sliding_window = None |
| self.sliding_window = sliding_window |
|
|
| swa_rope_theta = getattr(config, "swa_rope_theta", -1.0) |
| rope_theta = config.rope_theta |
| if self.sliding_window is not None and swa_rope_theta > 0: |
| rope_theta = swa_rope_theta |
|
|
| self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) |
| self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) |
| self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) |
| self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) |
|
|
| self.use_qk_norm = config.use_qk_norm |
| if self.use_qk_norm: |
| self.q_norm = MiniMaxM2RMSNorm(self.num_heads * self.head_dim, eps=config.rms_norm_eps) |
| self.k_norm = MiniMaxM2RMSNorm(self.num_key_value_heads * self.head_dim, eps=config.rms_norm_eps) |
|
|
| rope_config = copy.deepcopy(config) |
| rope_config.hidden_size = config.hidden_size |
| rope_config.num_attention_heads = config.num_attention_heads |
| rope_config.partial_rotary_factor = float(config.rotary_dim) / float(self.head_dim) |
| rope_config.rope_theta = rope_theta |
| rope_config.max_position_embeddings = max_position_embeddings |
| self.rotary_emb = LlamaRotaryEmbedding(rope_config) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[Cache] = None, |
| use_cache: Optional[bool] = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| output_attentions: bool = False, |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
| bsz, q_len, _ = hidden_states.size() |
|
|
| query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
|
|
| if self.use_qk_norm: |
| q_flat = query_states.transpose(1, 2).reshape(bsz * q_len, -1) |
| k_flat = key_states.transpose(1, 2).reshape(bsz * q_len, -1) |
| q_flat = self.q_norm(q_flat) |
| k_flat = self.k_norm(k_flat) |
| query_states = q_flat.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| key_states = k_flat.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
|
|
| if position_embeddings is None: |
| cos, sin = self.rotary_emb(value_states, position_ids) |
| else: |
| cos, sin = position_embeddings |
|
|
| query_states, key_states = apply_rotary_pos_emb_partial( |
| query_states.transpose(1, 2), key_states.transpose(1, 2), cos, sin, self.rotary_dim |
| ) |
| query_states = query_states.transpose(1, 2) |
| key_states = key_states.transpose(1, 2) |
|
|
| if past_key_values is not None: |
| cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} |
| key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) |
|
|
| key_states = repeat_kv(key_states, self.num_key_value_groups) |
| value_states = repeat_kv(value_states, self.num_key_value_groups) |
|
|
| attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * self.scaling |
| if attention_mask is not None: |
| attn_weights = attn_weights + attention_mask |
|
|
| if self.sliding_window is not None and past_key_values is None: |
| query_positions = torch.arange(q_len, device=hidden_states.device).view(1, 1, q_len, 1) |
| key_positions = torch.arange(key_states.shape[-2], device=hidden_states.device).view(1, 1, 1, -1) |
| window_mask = key_positions < (query_positions - self.sliding_window) |
| if window_mask.any(): |
| attn_weights = attn_weights.masked_fill(window_mask, float("-inf")) |
|
|
| attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) |
| if self.training and self.attention_dropout > 0: |
| attn_weights = F.dropout(attn_weights, p=self.attention_dropout) |
|
|
| attn_output = torch.matmul(attn_weights, value_states) |
| attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1) |
| attn_output = self.o_proj(attn_output) |
|
|
| if not output_attentions: |
| attn_weights = None |
| return attn_output, attn_weights |
|
|
|
|
| class MiniMaxM2LogitsProcessor(nn.Module): |
| def __init__(self, config: MiniMaxM2Config) -> None: |
| super().__init__() |
| self.scale = getattr(config, "logits_scale", 1.0) |
|
|
| def forward(self, lm_head: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: |
| logits = lm_head(hidden_states) |
| if self.scale != 1.0: |
| logits = logits * self.scale |
| return logits |
|
|
|
|
| class MiniMaxM2DecoderLayer(nn.Module): |
| def __init__(self, config: MiniMaxM2Config, layer_idx: int) -> None: |
| super().__init__() |
| self.hidden_size = config.hidden_size |
| self.self_attn = MiniMaxM2Attention(config, layer_idx) |
| self.block_sparse_moe = MiniMaxM2SparseMoeBlock(config) |
| self.input_layernorm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| self.post_attention_layernorm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_values: Optional[Cache] = None, |
| use_cache: Optional[bool] = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| output_attentions: bool = False, |
| residual: Optional[torch.Tensor] = None, |
| ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]: |
| residual_input = hidden_states if residual is None else residual |
| hidden_states = self.input_layernorm(hidden_states) |
|
|
| attn_output, attn_weights = self.self_attn( |
| hidden_states=hidden_states, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| use_cache=use_cache, |
| cache_position=cache_position, |
| position_embeddings=position_embeddings, |
| output_attentions=output_attentions, |
| ) |
| hidden_states = residual_input + attn_output |
|
|
| residual_post_attn = hidden_states |
| hidden_states = self.post_attention_layernorm(hidden_states) |
| moe_output, router_logits = self.block_sparse_moe(hidden_states) |
| hidden_states = residual_post_attn + moe_output |
|
|
| return hidden_states, hidden_states, router_logits, attn_weights |
|
|
|
|
| class MiniMaxM2PreTrainedModel(PreTrainedModel): |
| config_class = MiniMaxM2Config |
| base_model_prefix = "model" |
| supports_gradient_checkpointing = True |
| _no_split_modules = ["MiniMaxM2DecoderLayer"] |
| _supports_flash_attn = False |
| _supports_sdpa = False |
| _supports_attention_backend = False |
|
|
| def _init_weights(self, module: nn.Module) -> None: |
| if isinstance(module, nn.Linear): |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) |
| if module.bias is not None: |
| module.bias.data.zero_() |
| elif isinstance(module, nn.Embedding): |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) |
| if module.padding_idx is not None: |
| module.weight.data[module.padding_idx].zero_() |
|
|
| def _remap_qkv_weights(self, state_dict): |
| num_q = self.config.num_attention_heads * self.config.head_dim |
| num_kv = self.config.num_key_value_heads * self.config.head_dim |
|
|
| for layer_idx in range(self.config.num_hidden_layers): |
| prefix = f"model.layers.{layer_idx}.self_attn" |
| weight_key = f"{prefix}.qkv_proj.weight" |
| if weight_key in state_dict: |
| qkv_weight = state_dict.pop(weight_key) |
| q_weight, k_weight, v_weight = qkv_weight.split([num_q, num_kv, num_kv], dim=0) |
| state_dict.setdefault(f"{prefix}.q_proj.weight", q_weight) |
| state_dict.setdefault(f"{prefix}.k_proj.weight", k_weight) |
| state_dict.setdefault(f"{prefix}.v_proj.weight", v_weight) |
|
|
| def load_state_dict(self, state_dict, strict: bool = True): |
| if not isinstance(state_dict, dict): |
| raise TypeError(f"Expected state_dict to be dict, got {type(state_dict)}") |
|
|
| filtered_state_dict = {} |
| drop_suffixes = ("weight_scale_inv", "weight_scale", "input_scale", "scales", "amax") |
| for key, value in state_dict.items(): |
| if key.endswith(drop_suffixes) or "fp8" in key: |
| continue |
| filtered_state_dict[key] = value |
|
|
| self._remap_qkv_weights(filtered_state_dict) |
|
|
| if logger.isEnabledFor(logging.INFO): |
| logger.info( |
| "MiniMaxM2: loading %d tensors (filtered from %d original).", |
| len(filtered_state_dict), |
| len(state_dict), |
| ) |
|
|
| load_start = time.perf_counter() |
| result = super().load_state_dict(filtered_state_dict, strict=strict) |
| load_elapsed = time.perf_counter() - load_start |
| if logger.isEnabledFor(logging.INFO): |
| logger.info("MiniMaxM2: state_dict load finished in %.2f seconds.", load_elapsed) |
|
|
| return result |
|
|
|
|
| class MiniMaxM2Model(MiniMaxM2PreTrainedModel): |
| def __init__(self, config: MiniMaxM2Config) -> None: |
| super().__init__(config) |
| self.padding_idx = config.pad_token_id |
| self.vocab_size = config.vocab_size |
|
|
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) |
| self.layers = nn.ModuleList( |
| [MiniMaxM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] |
| ) |
| self.norm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| self.gradient_checkpointing = False |
|
|
| self.post_init() |
|
|
| def get_input_embeddings(self) -> nn.Module: |
| return self.embed_tokens |
|
|
| def set_input_embeddings(self, value: nn.Module) -> None: |
| self.embed_tokens = value |
|
|
| 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[Cache] = None, |
| inputs_embeds: Optional[torch.Tensor] = None, |
| cache_position: Optional[torch.LongTensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: bool = False, |
| output_hidden_states: bool = False, |
| output_router_logits: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| ) -> Union[MoeModelOutputWithPast, Tuple]: |
| if (input_ids is None) == (inputs_embeds is None): |
| raise ValueError("You must specify exactly one of input_ids or inputs_embeds.") |
|
|
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| use_cache = use_cache if use_cache is not None else self.config.use_cache |
| output_router_logits = ( |
| output_router_logits if output_router_logits is not None else self.config.output_router_logits |
| ) |
|
|
| if inputs_embeds is None: |
| inputs_embeds = self.embed_tokens(input_ids) |
|
|
| if use_cache and past_key_values is None: |
| past_key_values = DynamicCache(config=self.config) |
|
|
| if cache_position is None: |
| past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 |
| cache_position = torch.arange( |
| past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device |
| ) |
|
|
| if position_ids is None: |
| position_ids = cache_position.unsqueeze(0) |
|
|
| if self.config.sliding_window is not None: |
| causal_mask = create_sliding_window_causal_mask( |
| config=self.config, |
| input_embeds=inputs_embeds, |
| attention_mask=attention_mask, |
| cache_position=cache_position, |
| past_key_values=past_key_values, |
| position_ids=position_ids, |
| ) |
| else: |
| causal_mask = create_causal_mask( |
| config=self.config, |
| input_embeds=inputs_embeds, |
| attention_mask=attention_mask, |
| cache_position=cache_position, |
| past_key_values=past_key_values, |
| position_ids=position_ids, |
| ) |
|
|
| hidden_states = inputs_embeds |
|
|
| all_hidden_states = () if output_hidden_states else None |
| all_attentions = () if output_attentions else None |
| all_router_logits = () if output_router_logits else None |
|
|
| residual = None |
| for decoder_layer in self.layers: |
| if output_hidden_states: |
| all_hidden_states = all_hidden_states + (hidden_states,) |
|
|
| layer_outputs = decoder_layer( |
| hidden_states, |
| attention_mask=causal_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| use_cache=use_cache, |
| cache_position=cache_position, |
| position_embeddings=None, |
| output_attentions=output_attentions, |
| residual=residual, |
| ) |
|
|
| hidden_states, residual, router_logits, attn_weights = layer_outputs |
|
|
| if output_router_logits: |
| all_router_logits = all_router_logits + (router_logits,) |
| if output_attentions: |
| all_attentions = all_attentions + (attn_weights,) |
|
|
| hidden_states = self.norm(hidden_states) |
|
|
| if output_hidden_states: |
| all_hidden_states = all_hidden_states + (hidden_states,) |
|
|
| if not return_dict: |
| outputs = (hidden_states, past_key_values) |
| if output_hidden_states: |
| outputs += (all_hidden_states,) |
| if output_attentions: |
| outputs += (all_attentions,) |
| if output_router_logits: |
| outputs += (all_router_logits,) |
| return outputs |
|
|
| return MoeModelOutputWithPast( |
| last_hidden_state=hidden_states, |
| past_key_values=past_key_values, |
| hidden_states=all_hidden_states, |
| attentions=all_attentions, |
| router_logits=all_router_logits, |
| ) |
|
|
|
|
| class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin): |
| def __init__(self, config: MiniMaxM2Config) -> None: |
| super().__init__(config) |
| self.model = MiniMaxM2Model(config) |
| self.vocab_size = config.vocab_size |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.router_aux_loss_coef = config.router_aux_loss_coef |
| self.num_experts = config.num_local_experts |
| self.num_experts_per_tok = config.num_experts_per_tok |
| self.logits_processor = MiniMaxM2LogitsProcessor(config) |
|
|
| self.post_init() |
|
|
| def get_input_embeddings(self) -> nn.Module: |
| return self.model.embed_tokens |
|
|
| def set_input_embeddings(self, value: nn.Module) -> None: |
| self.model.embed_tokens = value |
|
|
| def get_output_embeddings(self) -> nn.Module: |
| return self.lm_head |
|
|
| def set_output_embeddings(self, new_embeddings: nn.Module) -> None: |
| self.lm_head = new_embeddings |
|
|
| def prepare_inputs_for_generation( |
| self, |
| input_ids: torch.LongTensor, |
| past_key_values: Optional[Cache] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| inputs_embeds: Optional[torch.Tensor] = None, |
| **kwargs, |
| ): |
| if past_key_values is not None: |
| input_ids = input_ids[:, -1:] |
| if attention_mask is not None: |
| attention_mask = attention_mask[:, -past_key_values.get_seq_length() - 1 :] |
|
|
| return { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| "past_key_values": past_key_values, |
| "inputs_embeds": inputs_embeds, |
| } |
|
|
| 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[Cache] = None, |
| inputs_embeds: Optional[torch.Tensor] = None, |
| labels: Optional[torch.LongTensor] = None, |
| cache_position: Optional[torch.LongTensor] = None, |
| use_cache: Optional[bool] = None, |
| output_attentions: bool = False, |
| output_hidden_states: bool = False, |
| output_router_logits: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| logits_to_keep: Union[int, torch.Tensor] = 0, |
| ) -> Union[MoeCausalLMOutputWithPast, Tuple]: |
| 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 |
| ) |
|
|
| model_outputs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| inputs_embeds=inputs_embeds, |
| cache_position=cache_position, |
| use_cache=use_cache, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| output_router_logits=output_router_logits, |
| return_dict=True, |
| ) |
|
|
| hidden_states = model_outputs.last_hidden_state |
| slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) and logits_to_keep > 0 else slice(None) |
| logits = self.logits_processor(self.lm_head, hidden_states[:, slice_indices, :]) |
|
|
| loss = None |
| if labels is not None: |
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| loss_fct = nn.CrossEntropyLoss() |
| loss = loss_fct(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1)) |
|
|
| aux_loss = None |
| if output_router_logits and model_outputs.router_logits is not None: |
| aux_loss = load_balancing_loss_func( |
| model_outputs.router_logits, |
| num_experts=self.num_experts, |
| top_k=self.num_experts_per_tok, |
| attention_mask=attention_mask, |
| ) |
| if loss is not None: |
| loss = loss + self.router_aux_loss_coef * aux_loss.to(loss.device) |
|
|
| if not return_dict: |
| output = (logits,) + (model_outputs.past_key_values,) |
| if output_hidden_states: |
| output += (model_outputs.hidden_states,) |
| if output_attentions: |
| output += (model_outputs.attentions,) |
| if output_router_logits: |
| output += (model_outputs.router_logits,) |
| return ((loss,) + output) if loss is not None else output |
|
|
| return MoeCausalLMOutputWithPast( |
| loss=loss, |
| aux_loss=aux_loss, |
| logits=logits, |
| past_key_values=model_outputs.past_key_values, |
| hidden_states=model_outputs.hidden_states, |
| attentions=model_outputs.attentions, |
| router_logits=model_outputs.router_logits, |
| ) |
|
|
| |
| |
| |
|
|
| MiniMaxRMSNorm = MiniMaxM2RMSNorm |
| MiniMaxSparseMoeBlock = MiniMaxM2SparseMoeBlock |
| MiniMaxAttention = MiniMaxM2Attention |
| MiniMaxDecoderLayer = MiniMaxM2DecoderLayer |
| MiniMaxMLP = MiniMaxM2MLP |
| MiniMaxPreTrainedModel = MiniMaxM2PreTrainedModel |
| MiniMaxModel = MiniMaxM2Model |
|
|
|
|
| class MiniMaxForCausalLM(MiniMaxM2ForCausalLM): |
| """Alias for compatibility with checkpoints exporting MiniMaxForCausalLM.""" |
|
|
|
|
| __all__ = [ |
| "MiniMaxM2RMSNorm", |
| "MiniMaxM2SparseMoeBlock", |
| "MiniMaxM2Attention", |
| "MiniMaxM2DecoderLayer", |
| "MiniMaxM2Model", |
| "MiniMaxM2ForCausalLM", |
| "MiniMaxM2PreTrainedModel", |
| "MiniMaxRMSNorm", |
| "MiniMaxSparseMoeBlock", |
| "MiniMaxAttention", |
| "MiniMaxDecoderLayer", |
| "MiniMaxPreTrainedModel", |
| "MiniMaxModel", |
| "MiniMaxMLP", |
| "MiniMaxForCausalLM", |
| ] |
|
|
