| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | from transformers import PreTrainedConfig |
| 6 | |
| 7 | |
| 8 | class MiniMaxM3TextConfig(PreTrainedConfig): |
| 9 | """Config for the MiniMax M3 text backbone (MiniMaxM3SparseForCausalLM). |
| 10 | |
| 11 | Defaults mirror the ``text_config`` of the MiniMax-M3-preview checkpoint. |
| 12 | """ |
| 13 | |
| 14 | model_type = "minimax_m3_text" |
| 15 | architectures = ["MiniMaxM3SparseForCausalLM"] |
| 16 | |
| 17 | def __init__( |
| 18 | self, |
| 19 | vocab_size: int = 200064, |
| 20 | hidden_size: int = 6144, |
| 21 | intermediate_size: int = 3072, |
| 22 | dense_intermediate_size: int = 12288, |
| 23 | shared_intermediate_size: int = 3072, |
| 24 | num_hidden_layers: int = 60, |
| 25 | num_attention_heads: int = 64, |
| 26 | num_key_value_heads: int = 4, |
| 27 | head_dim: int = 128, |
| 28 | max_position_embeddings: int = 524288, |
| 29 | rms_norm_eps: float = 1e-6, |
| 30 | use_gemma_norm: bool = True, |
| 31 | attention_output_gate: bool = False, |
| 32 | rope_theta: float = 5000000, |
| 33 | rotary_dim: int = 64, |
| 34 | partial_rotary_factor: float = 0.5, |
| 35 | hidden_act: str = "swigluoai", |
| 36 | swiglu_alpha: float = 1.702, |
| 37 | # SwiGLU-OAI uses the (up + 1) bias, i.e. beta=1.0 (matches the |
| 38 | # reference: gate * sigmoid(gate * alpha) * (up + 1)). The checkpoint |
| 39 | # config omits swiglu_beta, so this default must stay 1.0. |
| 40 | swiglu_beta: float = 1.0, |
| 41 | swiglu_limit: float = 7.0, |
| 42 | use_qk_norm: bool = True, |
| 43 | qk_norm_type: str = "per_head", |
| 44 | num_local_experts: int = 128, |
| 45 | num_experts_per_tok: int = 4, |
| 46 | n_shared_experts: int = 1, |
| 47 | scoring_func: str = "sigmoid", |
| 48 | use_routing_bias: bool = True, |
| 49 | routed_scaling_factor: float = 2.0, |
| 50 | num_mtp_modules: int = 1, |
| 51 | moe_layer_freq: list[int] | None = None, |
| 52 | sparse_attention_config: dict[str, Any] | None = None, |
| 53 | tie_word_embeddings: bool = False, |
| 54 | **kwargs, |
| 55 | ): |
| 56 | self.vocab_size = vocab_size |
| 57 | self.hidden_size = hidden_size |
| 58 | self.intermediate_size = intermediate_size |
| 59 | self.dense_intermediate_size = dense_intermediate_size |
| 60 | self.shared_intermediate_size = shared_intermediate_size |
| 61 | self.num_hidden_layers = num_hidden_layers |
| 62 | self.num_attention_heads = num_attention_heads |
| 63 | self.num_key_value_heads = num_key_value_heads |
| 64 | self.head_dim = head_dim |
| 65 | self.max_position_embeddings = max_position_embeddings |
| 66 | self.rms_norm_eps = rms_norm_eps |
| 67 | self.use_gemma_norm = use_gemma_norm |
| 68 | self.attention_output_gate = attention_output_gate |
| 69 | self.rope_theta = rope_theta |
| 70 | self.rotary_dim = rotary_dim |
| 71 | self.partial_rotary_factor = partial_rotary_factor |
| 72 | self.hidden_act = hidden_act |
| 73 | self.swiglu_alpha = swiglu_alpha |
| 74 | self.swiglu_beta = swiglu_beta |
| 75 | self.swiglu_limit = swiglu_limit |
| 76 | self.use_qk_norm = use_qk_norm |
| 77 | self.qk_norm_type = qk_norm_type |
| 78 | self.num_local_experts = num_local_experts |
| 79 | self.num_experts_per_tok = num_experts_per_tok |
| 80 | self.n_shared_experts = n_shared_experts |
| 81 | self.scoring_func = scoring_func |
| 82 | self.use_routing_bias = use_routing_bias |
| 83 | self.routed_scaling_factor = routed_scaling_factor |
| 84 | self.num_mtp_modules = num_mtp_modules |
| 85 | # First 3 layers are dense; the remaining 57 are sparse MoE. |
| 86 | self.moe_layer_freq = ( |
| 87 | moe_layer_freq if moe_layer_freq is not None else [0] * 3 + [1] * 57 |
| 88 | ) |
| 89 | self.sparse_attention_config = ( |
| 90 | sparse_attention_config |
| 91 | if sparse_attention_config is not None |
| 92 | else { |
| 93 | "use_sparse_attention": True, |
| 94 | "sparse_index_dim": 128, |
| 95 | "sparse_num_index_heads": 4, |
| 96 | "sparse_topk_blocks": 16, |
| 97 | "sparse_block_size": 128, |
| 98 | "sparse_disable_index_value": [0] * 3 + [1] * 57, |
| 99 | "sparse_score_type": "max", |
| 100 | "sparse_init_block": 0, |
| 101 | "sparse_local_block": 1, |
| 102 | "sparse_attention_freq": [0] * 3 + [1] * 57, |
| 103 | } |
| 104 | ) |
| 105 | super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) |
| 106 | |
| 107 | |
| 108 | class MiniMaxM3MTPConfig(MiniMaxM3TextConfig): |
| 109 | """Config for a standalone MiniMax M3 MTP (multi-token prediction) head. |
| 110 | |
| 111 | The MTP transformer layer is structurally a single MiniMax M3 decoder |
| 112 | layer, so this reuses the text backbone schema. Standalone MTP checkpoints |
| 113 | use ``model_type='minimax_m3_mtp'`` and a single hidden layer. |
| 114 | """ |
| 115 | |
| 116 | model_type = "minimax_m3_mtp" |
| 117 | architectures = ["MiniMaxM3MTP"] |
| 118 | |
| 119 | def __init__(self, num_hidden_layers: int = 1, **kwargs): |
| 120 | super().__init__(num_hidden_layers=num_hidden_layers, **kwargs) |
| 121 | |
| 122 | |
| 123 | class MiniMaxM3Config(PreTrainedConfig): |
| 124 | """Top-level MiniMax M3 (VL) config. |
| 125 | |
| 126 | Holds the text backbone as ``text_config`` so that |
| 127 | ``config.get_text_config()`` extracts the MiniMaxM3SparseForCausalLM |
| 128 | backbone. Vision components are kept as a raw dict passthrough and are |
| 129 | not modeled here. |
| 130 | """ |
| 131 | |
| 132 | model_type = "minimax_m3_vl" |
| 133 | |
| 134 | def __init__( |
| 135 | self, |
| 136 | text_config: dict | MiniMaxM3TextConfig | None = None, |
| 137 | vision_config: dict | None = None, |
| 138 | **kwargs, |
| 139 | ): |
| 140 | if text_config is None: |
| 141 | text_config = MiniMaxM3TextConfig() |
| 142 | elif isinstance(text_config, dict): |
| 143 | text_config = MiniMaxM3TextConfig(**text_config) |
| 144 | self.text_config = text_config |
| 145 | self.vision_config = vision_config |
| 146 | |
| 147 | self.hidden_size = text_config.hidden_size |
| 148 | |
| 149 | super().__init__(**kwargs) |