diff --git a/docs/api/build_from_gguf.md b/docs/api/build_from_gguf.md index 5dae11c51..9b82d7c4b 100644 --- a/docs/api/build_from_gguf.md +++ b/docs/api/build_from_gguf.md @@ -15,7 +15,7 @@ from mobius import build_from_gguf | Census | Total | Closure | |---|---:|---| -| Architectures | 147 | graph verdicts: {'deferred': 88, 'rejected': 2, 'supported': 57}; importable: 55; quantized import: {'rejected': 11, 'supported': 136}; runtime: {'deferred': 144, 'rejected': 2, 'supported': 1} | +| Architectures | 147 | graph verdicts: {'deferred': 87, 'rejected': 2, 'supported': 58}; importable: 56; quantized import: {'rejected': 11, 'supported': 136}; runtime: {'deferred': 144, 'rejected': 2, 'supported': 1} | | Active stored qtypes | 25 | 24 have an import route; 1 are explicitly deferred with no route | | Serialized projector strings | 60 | {'graph-importable': 2, 'runtime-supported': 0} | | Tokenizer pre identifiers | 87 | 56 semantic groups; all default to deferred and become materializable only from a validated embedded `tokenizer.huggingface.json` or an exact pinned source in runtime evidence | @@ -428,7 +428,7 @@ before graph construction or durable output. | `mimo2` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiMo2 requires fused-QKV dense MTP blocks, attention sinks, interleaved sliding KV cache, and three chained heads selected by offsets. Mobius permits one head and cannot preserve that state or FP8 converter transform. | | `minicpm` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiniCPM requires architecture-specific embedding, residual, and logit scales, Q/K permutation, optional long/short RoPE tensors, and a conditional dense-or-MoE loader. The existing MiniCPM graph does not prove this complete GGUF contract. | | `minicpm3` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | The pinned MiniCPM3 graph uses MLA Q/KV LoRA projections, separate NoPE/RoPE query and key channels, and embedding, residual, and LM-head scales. The current Mobius MiniCPM graph does not represent that exact topology or its scales. | -| `minimax-01` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | The pinned loader schedule is not periodic and its Lightning Attention decay, scaling, residual multipliers, and recurrent rollback semantics are not represented by the current MiniMax graph. | +| `minimax-01` | — | model=`minimax`; tensor=`minimax` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Graph import is exact, but released ORT GenAI packaging cannot represent the heterogeneous KV/recurrent state slots or bounded rollback snapshots; runtime packaging remains tracked by #605. | | `minimax-m2` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiniMax-M2 uses full-vector Q/K norms, partial RoPE, and all-layer correction-biased routed experts under metadata-selected gating. Mobius has no exact graph or suffix-safe expert import for that topology. | | `minimax-m3` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiniMax-M3 adds F32 sparse-indexer tensors and a second index-key cache with position/cell maps, block masks, rollback, and reorder semantics alongside main K/V state. Mobius has no MSA cache task or sparse-index operators; dense fallback would change the model. | | `mistral3` | — | none (fails before config extraction) | exact-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | The pinned Mistral3 loader selects dense or routed-expert text blocks from metadata and applies architecture-specific output temperature scaling. A VLM package additionally requires the deferred Pixtral clip sidecar and exact patch/merge/token contract. The existing Hugging Face Mistral3 graph does not cover that conditional GGUF closure. | @@ -593,8 +593,11 @@ real-artifact full-logit and stateful-generation parity. - `static_cache=True` and non-hybrid task dispatch are rejected for these mixed state ABIs. -`minimax-01` remains deferred before config extraction because its pinned -Lightning schedule and decay/scaling semantics do not match the current graph. +`minimax-01` graph import supports its exact pinned Lightning schedule, +decay/scaling semantics, recurrent state, and mixed full-attention cache. +Runtime packaging remains deferred because the released schema cannot represent +that heterogeneous state ABI or bounded rollback snapshots; this is tracked by +[`onnxruntime/mobius#605`](https://github.com/onnxruntime/mobius/issues/605). PLaMo2 has a dedicated alternating Mamba1/attention graph and strict GGUF tensor closure. Its mixed per-layer recurrent/KV runtime package remains deferred to [`onnxruntime/mobius#605`](https://github.com/onnxruntime/mobius/issues/605). diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index 3708a7900..ca6ba3868 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -47,6 +47,7 @@ LongcatFlashConfig, Mamba2Config, MambaConfig, + MiniMaxConfig, MllamaConfig, MMSConfig, MoonshineConfig, @@ -127,6 +128,7 @@ "Mamba2Config", "MambaConfig", "MllamaConfig", + "MiniMaxConfig", "MMSConfig", "MoonshineConfig", "MuseGlimmerConfig", diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index e222457bc..24f24bdf7 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -486,6 +486,16 @@ class ArchitectureConfig(BaseModelConfig): topk_method: str = "greedy" first_k_dense_replace: int = 0 n_shared_experts: int | None = None + disable_qmoe: bool = False + + # MiniMax-01 hybrid attention and normalized-residual scaling. + lightning_norm_eps: float | None = None + full_attn_alpha_factor: float = 1.0 + full_attn_beta_factor: float = 1.0 + linear_attn_alpha_factor: float = 1.0 + linear_attn_beta_factor: float = 1.0 + mlp_alpha_factor: float = 1.0 + mlp_beta_factor: float = 1.0 # Multi-head Latent Attention (MLA) config — DeepSeek-V2/V3 q_lora_rank: int | None = None @@ -1435,6 +1445,64 @@ class CausalLMConfig(ArchitectureConfig): """ +@dataclasses.dataclass +class MiniMaxConfig(CausalLMConfig): + """Exact configuration for MiniMax-Text-01 and MiniMax-M1 backbones.""" + + @classmethod + def from_transformers(cls, config, parent_config=None) -> MiniMaxConfig: + base = ArchitectureConfig.from_transformers(config, parent_config) + raw_schedule = getattr(config, "attn_type_list", None) + if raw_schedule is None: + raise ValueError("MiniMax-01 config requires an explicit attn_type_list") + if len(raw_schedule) != base.num_hidden_layers: + raise ValueError( + "MiniMax-01 attn_type_list must contain exactly " + f"{base.num_hidden_layers} entries, got {len(raw_schedule)}" + ) + if any(value not in (0, 1, False, True) for value in raw_schedule): + raise ValueError("MiniMax-01 attn_type_list entries must be 0 or 1") + if not bool(getattr(config, "postnorm", True)): + raise ValueError("MiniMax-01 requires postnorm=true") + if int(getattr(config, "shared_intermediate_size", 0) or 0): + raise ValueError( + "MiniMax-01 shared experts are not supported by the pinned GGUF architecture" + ) + + beta_names = ( + "layernorm_full_attention_beta", + "layernorm_linear_attention_beta", + "layernorm_mlp_beta", + ) + betas = {name: float(getattr(config, name, 1.0)) for name in beta_names} + if any(not math.isclose(value, 1.0) for value in betas.values()): + raise ValueError(f"MiniMax-01 beta residual factors must all equal 1.0: {betas}") + + fields = _shallow_fields(base) + fields.update( + model_type="minimax", + layer_types=[ + "full_attention" if int(value) == 1 else "lightning_attention" + for value in raw_schedule + ], + hidden_act="silu", + norm_topk_prob=True, + disable_qmoe=True, + lightning_norm_eps=float(getattr(config, "lightning_norm_eps", 1e-6)), + full_attn_alpha_factor=float( + getattr(config, "layernorm_full_attention_alpha", 1.0) + ), + full_attn_beta_factor=betas["layernorm_full_attention_beta"], + linear_attn_alpha_factor=float( + getattr(config, "layernorm_linear_attention_alpha", 1.0) + ), + linear_attn_beta_factor=betas["layernorm_linear_attention_beta"], + mlp_alpha_factor=float(getattr(config, "layernorm_mlp_alpha", 1.0)), + mlp_beta_factor=betas["layernorm_mlp_beta"], + ) + return cls(**fields) + + @dataclasses.dataclass class EncoderConfig(ArchitectureConfig): """Configuration for encoder-only models (BERT, ViT, etc.).""" diff --git a/src/mobius/_configs/_base_test.py b/src/mobius/_configs/_base_test.py index 51f5ec1c0..a9db12108 100644 --- a/src/mobius/_configs/_base_test.py +++ b/src/mobius/_configs/_base_test.py @@ -7,7 +7,48 @@ import types -from mobius._configs import ArchitectureConfig, NemotronParseConfig +import pytest + +from mobius._configs import ArchitectureConfig, MiniMaxConfig, NemotronParseConfig + + +def test_minimax_config_extracts_exact_schedule_head_geometry_and_residuals(): + config = types.SimpleNamespace( + model_type="MiniMaxText01", + hidden_size=48, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + rotary_dim=8, + rope_theta=10_000_000.0, + vocab_size=64, + attn_type_list=[0, 1], + num_local_experts=2, + num_experts_per_tok=1, + rms_norm_eps=1e-5, + postnorm=True, + shared_intermediate_size=0, + layernorm_full_attention_alpha=3.5, + layernorm_full_attention_beta=1.0, + layernorm_linear_attention_alpha=3.5, + layernorm_linear_attention_beta=1.0, + layernorm_mlp_alpha=3.5, + layernorm_mlp_beta=1.0, + ) + + extracted = MiniMaxConfig.from_transformers(config) + + assert extracted.model_type == "minimax" + assert extracted.head_dim == 16 + assert extracted.partial_rotary_factor == pytest.approx(0.5) + assert extracted.layer_types == ["lightning_attention", "full_attention"] + assert extracted.lightning_norm_eps == pytest.approx(1e-6) + assert extracted.full_attn_alpha_factor == pytest.approx(3.5) + assert extracted.linear_attn_alpha_factor == pytest.approx(3.5) + assert extracted.mlp_alpha_factor == pytest.approx(3.5) + assert extracted.disable_qmoe def test_nemotron_parse_maps_raw_mbart_decoder_attention_heads(): diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 278150b13..f2b2af5c5 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -32,6 +32,7 @@ Lfm2Config, Lfm2MoeConfig, Lfm2VlConfig, + MiniMaxConfig, MMSConfig, MoonshineConfig, MuseGlimmerConfig, @@ -602,7 +603,9 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "granitemoeshared": ModelRegistration(GraniteMoECausalLMModel), "hunyuan_v1_moe": ModelRegistration(HunYuanMoEV1CausalLMModel), "jetmoe": ModelRegistration(JetMoeCausalLMModel), - "minimax": ModelRegistration(MiniMaxCausalLMModel), + "minimax": ModelRegistration(MiniMaxCausalLMModel, config_class=MiniMaxConfig), + "MiniMaxText01": ModelRegistration(MiniMaxCausalLMModel, config_class=MiniMaxConfig), + "minimax_text_01": ModelRegistration(MiniMaxCausalLMModel, config_class=MiniMaxConfig), "mixtral": ModelRegistration(MoECausalLMModel), "olmoe": ModelRegistration(MoECausalLMModel), "phimoe": ModelRegistration(Phi3MoECausalLMModel), diff --git a/src/mobius/components/_lightning_attention.py b/src/mobius/components/_lightning_attention.py index 3b75e59ae..d3f9149fd 100644 --- a/src/mobius/components/_lightning_attention.py +++ b/src/mobius/components/_lightning_attention.py @@ -59,31 +59,38 @@ class LightningAttention(nn.Module): layer_idx: Zero-based layer index, used to compute slope_rate. """ - def __init__(self, config: ArchitectureConfig, layer_idx: int): + def __init__( + self, + config: ArchitectureConfig, + layer_idx: int, + linear_class: type | None = None, + ): super().__init__() + linear_class = linear_class or Linear self.num_heads = config.num_attention_heads - self.head_dim = config.hidden_size // config.num_attention_heads + self.head_dim = config.head_dim self.hidden_size = config.hidden_size self._dtype = config.dtype # Fused QKV projection — SiLU applied to all 3*H*d_k before split - self.qkv_proj = Linear( + self.qkv_proj = linear_class( config.hidden_size, self.num_heads * self.head_dim * 3, bias=False, ) - self.out_proj = Linear( + self.o_proj = linear_class( self.num_heads * self.head_dim, config.hidden_size, bias=False, ) # output_gate: sigmoid gate applied to the normalized attention output - self.output_gate = Linear( + self.output_gate = linear_class( config.hidden_size, self.num_heads * self.head_dim, bias=False, ) - self.norm = RMSNorm(self.num_heads * self.head_dim, eps=config.rms_norm_eps) + norm_eps = config.lightning_norm_eps or config.rms_norm_eps + self.norm = RMSNorm(self.num_heads * self.head_dim, eps=norm_eps) # Per-head log-space decay values (negative, so exp < 1) # HF: slope_rate[h] = base^(h+1) * factor @@ -97,6 +104,7 @@ def forward( op: OpBuilder, hidden_states: ir.Value, recurrent_state: ir.Value, + attention_mask: ir.Value | None = None, ): """Lightning Attention forward. @@ -117,14 +125,28 @@ def forward( qkv = self.qkv_proj(op, hidden_states) qkv = op.Swish(qkv) - # Split into Q, K, V: each (B, T, num_heads * head_dim) - head_total = self.num_heads * self.head_dim + # MiniMax stores Q/K/V adjacent within each head: + # (B, T, H * 3D) -> (B, T, H, 3D) -> three (B, T, H, D) tensors. + qkv = op.Reshape(qkv, [0, 0, self.num_heads, 3 * self.head_dim]) query, key, value = op.Split( qkv, - op.Constant(value_ints=[head_total, head_total, head_total]), + op.Constant(value_ints=[self.head_dim, self.head_dim, self.head_dim]), axis=-1, _outputs=3, ) + query = op.Reshape(query, [0, 0, -1]) + key = op.Reshape(key, [0, 0, -1]) + value = op.Reshape(value, [0, 0, -1]) + if attention_mask is not None: + # Only newly processed tokens contribute to the recurrent state. + current_mask = op.Slice( + attention_mask, + op.Neg(seq_dim), + op.Constant(value_ints=[9223372036854775807]), + op.Constant(value_ints=[1]), + ) + current_mask = op.Unsqueeze(op.CastLike(current_mask, value), [2]) + value = op.Mul(value, current_mask) # Static decay tensor: (B, T, num_heads) with constant per-head values. # Each decay[h] = -slope_rate[h] in log-space → exp(decay[h]) < 1. @@ -144,8 +166,6 @@ def forward( decay = op.Expand(decay_1, expand_to) # (B, T, num_heads) # LinearAttention "gated": S_t = exp(g_t) * S_{t-1} + k_t ⊗ v_t - # scale = 1/sqrt(head_dim) for proper scaling - scale = 1.0 / math.sqrt(self.head_dim) attn_out, new_state = op.LinearAttention( query, key, @@ -155,7 +175,7 @@ def forward( update_rule="gated", q_num_heads=self.num_heads, kv_num_heads=self.num_heads, - scale=scale, + scale=1.0, _domain=DOMAIN, _outputs=2, ) @@ -166,7 +186,7 @@ def forward( gate = op.Sigmoid(self.output_gate(op, hidden_states)) attn_out = op.Mul(gate, attn_out) - output = self.out_proj(op, attn_out) + output = self.o_proj(op, attn_out) return output, new_state @@ -174,8 +194,14 @@ def _compute_decay_log(layer_idx: int, num_layers: int, num_heads: int) -> list[ """Compute per-head log-space decay values for Lightning Attention. Returns negative values so that exp(decay[h]) = exp(-slope_rate[h]) < 1. - Matches HF ``MiniMaxLightningAttention.get_slope_rate()``. + Matches the pinned llama.cpp MiniMax-01 slope calculation. """ - base = 1.0 / (2.0 ** (8.0 / num_heads)) - factor = 1.0 - layer_idx / (num_layers - 1.0 + 1e-5) + 1e-5 - return [-(base ** (h + 1)) * factor for h in range(num_heads)] + if num_layers <= 1: + raise ValueError("MiniMax Lightning Attention requires at least two layers") + + if num_heads <= 0: + raise ValueError("MiniMax Lightning Attention requires at least one head") + start = 2.0 ** (-(2.0 ** -(math.log2(num_heads) - 3.0))) + slopes = [start ** (head + 1) for head in range(num_heads)] + factor = 1.0 - layer_idx / (num_layers - 1.0) + 1e-5 + return [-slope * factor for slope in slopes] diff --git a/src/mobius/components/_lightning_attention_test.py b/src/mobius/components/_lightning_attention_test.py new file mode 100644 index 000000000..a47fd1060 --- /dev/null +++ b/src/mobius/components/_lightning_attention_test.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import math + +import pytest + +from mobius.components._lightning_attention import _compute_decay_log + + +def test_minimax_decay_matches_pinned_64_head_formula(): + actual = _compute_decay_log(layer_idx=7, num_layers=80, num_heads=64) + factor = 1.0 - 7.0 / 79.0 + 1e-5 + expected = [-(2.0 ** (-(head + 1) / 8.0)) * factor for head in range(64)] + + assert actual == pytest.approx(expected) + + +def test_minimax_decay_supports_non_power_of_two_head_counts(): + actual = _compute_decay_log(layer_idx=0, num_layers=2, num_heads=6) + start = 2.0 ** (-(2.0 ** -(math.log2(6) - 3.0))) + expected = [-(start ** (head + 1)) * 1.00001 for head in range(6)] + + assert actual == pytest.approx(expected) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index de3420205..9aa47a249 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -429,7 +429,11 @@ def __init__( assert config.num_experts_per_tok is not None self.num_experts = config.num_local_experts self.top_k = config.num_experts_per_tok - self._qmoe_quantization = _supported_qmoe_quantization(config.quantization) + self._qmoe_quantization = ( + None + if getattr(config, "disable_qmoe", False) + else _supported_qmoe_quantization(config.quantization) + ) # Clipped-SwiGLU attributes (QMoE's ``activation_alpha``/``activation_beta``/ # ``swiglu_limit``). Left ``None`` by default so existing callers get a # byte-identical QMoE call (the attributes are simply omitted, even though diff --git a/src/mobius/integrations/gguf/_arch_registry.py b/src/mobius/integrations/gguf/_arch_registry.py index ff2af418c..ca6318a36 100644 --- a/src/mobius/integrations/gguf/_arch_registry.py +++ b/src/mobius/integrations/gguf/_arch_registry.py @@ -1191,14 +1191,31 @@ ), GGUFArchitectureSpec( gguf_arch="minimax-01", - config=Support.DEFERRED, - tensor_map=Support.DEFERRED, - graph=Support.DEFERRED, + model_type="minimax", + config_key_map="minimax", + config_postprocessor="minimax", + tensor_map_recipe=("minimax",), + required_metadata=( + "context_length", + "embedding_length", + "block_count", + "feed_forward_length", + "attention.head_count", + "attention.head_count_kv", + "attention.key_length", + "attention.value_length", + "attention.layer_norm_rms_epsilon", + "rope.freq_base", + "rope.dimension_count", + "expert_count", + "expert_used_count", + "residual_scale", + ), runtime=Support.DEFERRED, reason=( - "The pinned loader schedule is not periodic and its Lightning Attention " - "decay, scaling, residual multipliers, and recurrent rollback semantics " - "are not represented by the current MiniMax graph." + "Graph import is exact, but released ORT GenAI packaging cannot represent " + "the heterogeneous KV/recurrent state slots or bounded rollback snapshots; " + "runtime packaging remains tracked by #605." ), ), GGUFArchitectureSpec( diff --git a/src/mobius/integrations/gguf/_arch_registry_test.py b/src/mobius/integrations/gguf/_arch_registry_test.py index eca74635b..67774a0a6 100644 --- a/src/mobius/integrations/gguf/_arch_registry_test.py +++ b/src/mobius/integrations/gguf/_arch_registry_test.py @@ -57,7 +57,7 @@ #: Number of importable architectures. Pinned so that adding support is a #: deliberate act that also updates the documented support matrix, and so that #: accidentally losing an architecture is a failure rather than a silence. -_EXPECTED_SUPPORTED_COUNT = 55 +_EXPECTED_SUPPORTED_COUNT = 56 _FINAL_CENSUS_CLOSURE = frozenset( { "afmoe", @@ -167,6 +167,7 @@ "llada", "llada-moe", "llama", + "minimax-01", "modern-bert", "muse-glimmer", "nemotron", diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index a1bd627b0..9fe57b49c 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -616,6 +616,7 @@ def _validate_gguf_model( _raise_for_unsupported_auxiliary_quantization(gguf_model) _raise_for_invalid_falcon_h1_tensor_contract(gguf_model) _raise_for_invalid_plamo2_tensor_contract(gguf_model) + _raise_for_invalid_minimax_tensor_contract(gguf_model) _raise_for_invalid_hybrid_tensor_contract(gguf_model) _raise_for_invalid_t5_tensor_contract(gguf_model) _raise_for_malformed_recurrent_tensors(gguf_model) @@ -634,6 +635,128 @@ def _validate_gguf_model( inspect_gguf_tokenizer(gguf_model.metadata, source=source) +def _raise_for_invalid_minimax_tensor_contract(gguf_model) -> None: + """Validate MiniMax-01 metadata, per-layer families, and exact tensor shapes.""" + if gguf_model.architecture != "minimax-01": + return + + from mobius.integrations.gguf._config_mapping import _derive_hybrid_layout + + metadata = gguf_model.metadata + layers, layer_types, mtp_count = _derive_hybrid_layout( + "minimax-01", metadata, gguf_model.tensor_names + ) + assert layer_types is not None + if mtp_count: + raise ValueError("MiniMax-01 GGUF does not support appended MTP blocks") + + hidden = int(metadata["minimax-01.embedding_length"]) + intermediate = int(metadata["minimax-01.feed_forward_length"]) + heads = int(metadata["minimax-01.attention.head_count"]) + kv_heads = int(metadata["minimax-01.attention.head_count_kv"]) + head_dim = int(metadata["minimax-01.attention.key_length"]) + value_dim = int(metadata["minimax-01.attention.value_length"]) + rope_dim = int(metadata["minimax-01.rope.dimension_count"]) + experts = int(metadata["minimax-01.expert_count"]) + top_k = int(metadata["minimax-01.expert_used_count"]) + residual_scale = float(metadata["minimax-01.residual_scale"]) + norm_eps = float(metadata["minimax-01.attention.layer_norm_rms_epsilon"]) + rope_freq_base = float(metadata["minimax-01.rope.freq_base"]) + vocab = int(metadata.get("minimax-01.vocab_size", 0)) + if not vocab: + vocab = len(metadata.get("tokenizer.ggml.tokens", ())) + if ( + min(hidden, intermediate, heads, kv_heads, head_dim, experts, top_k, vocab) <= 0 + or heads % kv_heads + or value_dim != head_dim + or rope_dim <= 0 + or rope_dim > head_dim + or rope_dim % 2 + or experts <= 1 + or top_k > experts + or not math.isfinite(residual_scale) + or residual_scale <= 0 + or not math.isfinite(norm_eps) + or norm_eps <= 0 + or not math.isfinite(rope_freq_base) + or rope_freq_base <= 0 + ): + raise ValueError("MiniMax-01 GGUF has inconsistent architecture metadata") + if any( + key in metadata + for key in ( + "minimax-01.expert_shared_count", + "minimax-01.expert_shared_feed_forward_length", + ) + ): + raise ValueError("MiniMax-01 pinned GGUF does not support shared experts") + + q_width = heads * head_dim + kv_width = kv_heads * head_dim + required: dict[str, tuple[int, ...]] = { + "token_embd.weight": (vocab, hidden), + "output_norm.weight": (hidden,), + } + optional: dict[str, tuple[int, ...]] = {"output.weight": (vocab, hidden)} + for layer, layer_type in enumerate(layer_types): + prefix = f"blk.{layer}." + required.update( + { + prefix + "attn_norm.weight": (hidden,), + prefix + "attn_output.weight": (hidden, q_width), + prefix + "ffn_norm.weight": (hidden,), + prefix + "ffn_gate_inp.weight": (experts, hidden), + prefix + "ffn_gate_exps.weight": (experts, intermediate, hidden), + prefix + "ffn_up_exps.weight": (experts, intermediate, hidden), + prefix + "ffn_down_exps.weight": (experts, hidden, intermediate), + } + ) + if layer_type == "lightning_attention": + required.update( + { + prefix + "attn_qkv.weight": (3 * q_width, hidden), + prefix + "attn_gate.weight": (q_width, hidden), + prefix + "attn_norm_2.weight": (q_width,), + } + ) + else: + required.update( + { + prefix + "attn_q.weight": (q_width, hidden), + prefix + "attn_k.weight": (kv_width, hidden), + prefix + "attn_v.weight": (kv_width, hidden), + } + ) + + actual = set(gguf_model.tensor_names) + allowed = set(required) | set(optional) + out_of_range = sorted( + name + for name in actual + if (match := re.match(r"^blk\.(\d+)\.", name)) and int(match.group(1)) >= layers + ) + missing = sorted(set(required) - actual) + unexpected = sorted(actual - allowed) + if missing or unexpected or out_of_range: + raise ValueError( + "Invalid MiniMax-01 GGUF tensor closure: " + f"missing={missing}, unexpected={unexpected}, out_of_range={out_of_range}" + ) + if not hasattr(gguf_model, "tensor_items_raw"): + return + shapes = { + name: tuple(int(dimension) for dimension in shape) + for name, _raw, _qtype, shape in gguf_model.tensor_items_raw() + } + malformed = { + name: (shape, shapes.get(name)) + for name, shape in {**required, **optional}.items() + if name in shapes and shapes[name] != shape + } + if malformed: + raise ValueError(f"MiniMax-01 GGUF has invalid tensor shape(s): {malformed}") + + def _raise_for_invalid_dense_c01_tensor_contract(gguf_model) -> None: """Validate the exact pinned C01 dense profiles before config extraction.""" import numpy as np diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index 66f94cd40..f1acdd8b6 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -1271,6 +1271,137 @@ def add_q4(name: str, shape: tuple[int, ...]) -> None: writer.close() +def _write_minimax_gguf( + path: Path, + *, + quantized: bool, + quantized_embedding: bool = False, + omit: str | None = None, + extra: str | None = None, + malformed_shape: str | None = None, + recurrent_layers: list[bool] | None = None, + norm_eps: float = 1e-5, + rope_freq_base: float = 10_000_000.0, +) -> None: + """Write a tiny MiniMax-01 GGUF with one Lightning and one full-attention layer.""" + from gguf import GGMLQuantizationType, GGUFWriter + + hidden = 64 + intermediate = 32 + vocab = 64 + heads = 4 + kv_heads = 2 + head_dim = 16 + experts = 2 + rng = np.random.default_rng(601) + + writer = GGUFWriter(str(path), "minimax-01") + writer.add_context_length(64) + writer.add_embedding_length(hidden) + writer.add_feed_forward_length(intermediate) + writer.add_block_count(2) + writer.add_head_count(heads) + writer.add_head_count_kv(kv_heads) + writer.add_key_length(head_dim) + writer.add_value_length(head_dim) + writer.add_layer_norm_rms_eps(norm_eps) + writer.add_rope_freq_base(rope_freq_base) + writer.add_rope_dimension_count(8) + writer.add_expert_count(experts) + writer.add_expert_used_count(1) + writer.add_residual_scale(3.5565588200778455) + writer.add_vocab_size(vocab) + writer.add_array( + "minimax-01.attention.recurrent_layers", + recurrent_layers if recurrent_layers is not None else [True, False], + ) + + def shape_for(name: str, shape: tuple[int, ...]) -> tuple[int, ...]: + if name == malformed_shape: + return (*shape[:-1], shape[-1] + 1) + return shape + + def add_float( + name: str, + shape: tuple[int, ...], + *, + expert_base: float | None = None, + ) -> None: + if name == omit: + return + shape = shape_for(name, shape) + values = rng.normal(0.0, 0.02, shape).astype(np.float32) + if expert_base is not None: + for expert in range(shape[0]): + values[expert].fill(expert_base + expert) + writer.add_tensor(name, values) + + def add_q4(name: str, shape: tuple[int, ...]) -> None: + if name == omit: + return + shape = shape_for(name, shape) + assert shape[-1] % 32 == 0 + raw = np.zeros((*shape[:-1], shape[-1] // 32 * 18), dtype=np.uint8) + for index in np.ndindex(shape[:-1]): + for block in range(shape[-1] // 32): + offset = block * 18 + raw[(*index, slice(offset, offset + 2))] = np.array( + [rng.uniform(0.01, 0.05)], dtype=np.float16 + ).view(np.uint8) + raw[(*index, slice(offset + 2, offset + 18))] = rng.integers( + 0, 256, 16, dtype=np.uint8 + ) + writer.add_tensor(name, raw, raw_dtype=GGMLQuantizationType.Q4_0) + + (add_q4 if quantized_embedding else add_float)("token_embd.weight", (vocab, hidden)) + add_float("output_norm.weight", (hidden,)) + add_float("output.weight", (vocab, hidden)) + projection = add_q4 if quantized else add_float + q_width = heads * head_dim + kv_width = kv_heads * head_dim + for layer in range(2): + prefix = f"blk.{layer}." + add_float(prefix + "attn_norm.weight", (hidden,)) + add_float(prefix + "ffn_norm.weight", (hidden,)) + add_float(prefix + "ffn_gate_inp.weight", (experts, hidden)) + if quantized: + projection(prefix + "ffn_gate_exps.weight", (experts, intermediate, hidden)) + projection(prefix + "ffn_up_exps.weight", (experts, intermediate, hidden)) + projection(prefix + "ffn_down_exps.weight", (experts, hidden, intermediate)) + else: + add_float( + prefix + "ffn_gate_exps.weight", + (experts, intermediate, hidden), + expert_base=11.0, + ) + add_float( + prefix + "ffn_up_exps.weight", + (experts, intermediate, hidden), + expert_base=21.0, + ) + add_float( + prefix + "ffn_down_exps.weight", + (experts, hidden, intermediate), + expert_base=31.0, + ) + projection(prefix + "attn_output.weight", (hidden, q_width)) + if layer == 0: + projection(prefix + "attn_qkv.weight", (3 * q_width, hidden)) + projection(prefix + "attn_gate.weight", (q_width, hidden)) + add_float(prefix + "attn_norm_2.weight", (q_width,)) + else: + projection(prefix + "attn_q.weight", (q_width, hidden)) + projection(prefix + "attn_k.weight", (kv_width, hidden)) + projection(prefix + "attn_v.weight", (kv_width, hidden)) + if extra is not None: + add_float(extra, (1,)) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + def _write_nemotron_h_moe_gguf( path: Path, *, @@ -5303,6 +5434,234 @@ def unexpected_graph_build(*args, **kwargs): assert not graph_build_started +class TestMiniMaxGGUFBuild: + """MiniMax-01 GGUF import preserves hybrid state and expert tensor order.""" + + @staticmethod + def _inputs(tokens: np.ndarray, states: dict[str, np.ndarray] | None = None): + batch, sequence = tokens.shape + if states is None: + states = { + "past_key_values.0.recurrent_state": np.zeros( + (batch, 4, 16, 16), dtype=np.float32 + ), + "past_key_values.1.key": np.zeros((batch, 2, 0, 16), dtype=np.float32), + "past_key_values.1.value": np.zeros((batch, 2, 0, 16), dtype=np.float32), + } + past = states["past_key_values.1.key"].shape[2] + return { + "input_ids": tokens, + "attention_mask": np.ones((batch, past + sequence), dtype=np.int64), + "position_ids": np.broadcast_to( + np.arange(past, past + sequence, dtype=np.int64), (batch, sequence) + ).copy(), + **states, + } + + def test_float_import_runtime_save_reload_and_expert_order(self, tmp_path: Path) -> None: + from mobius._model_package import ModelPackage + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "minimax-f32.gguf" + _write_minimax_gguf(path, quantized=False) + package = build_from_gguf(path) + model = package["model"] + assert [value.name for value in model.graph.inputs if "past_" in value.name] == [ + "past_key_values.0.recurrent_state", + "past_key_values.1.key", + "past_key_values.1.value", + ] + for layer in range(2): + for expert in range(2): + prefix = f"model.layers.{layer}.mlp.experts.{expert}" + np.testing.assert_array_equal( + model.graph.initializers[ + f"{prefix}.gate_proj.weight_t" + ].const_value.numpy(), + 11.0 + expert, + ) + np.testing.assert_array_equal( + model.graph.initializers[f"{prefix}.up_proj.weight_t"].const_value.numpy(), + 21.0 + expert, + ) + np.testing.assert_array_equal( + model.graph.initializers[ + f"{prefix}.down_proj.weight_t" + ].const_value.numpy(), + 31.0 + expert, + ) + + session = OnnxModelSession(model) + try: + prefill = session.run(self._inputs(np.asarray([[1, 2], [3, 4]], np.int64))) + snapshot = { + "past_key_values.0.recurrent_state": prefill["present.0.recurrent_state"], + "past_key_values.1.key": prefill["present.1.key"], + "past_key_values.1.value": prefill["present.1.value"], + } + first = session.run(self._inputs(np.asarray([[5], [6]], np.int64), snapshot)) + replay = session.run(self._inputs(np.asarray([[5], [6]], np.int64), snapshot)) + reordered = session.run( + self._inputs( + np.asarray([[6], [5]], np.int64), + {name: value[[1, 0]] for name, value in snapshot.items()}, + ) + ) + finally: + session.close() + for name in first: + np.testing.assert_allclose(replay[name], first[name], rtol=0, atol=0) + np.testing.assert_allclose(reordered[name], first[name][[1, 0]], rtol=0, atol=0) + + saved = tmp_path / "saved-minimax" + package.save(saved, progress_bar=False, check_weights=True) + reloaded = ModelPackage.load(saved)["model"] + assert [value.name for value in reloaded.graph.outputs] == [ + value.name for value in model.graph.outputs + ] + + def test_quantized_source_preserves_exact_projection_roles(self, tmp_path: Path) -> None: + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "minimax-q4.gguf" + _write_minimax_gguf(path, quantized=True, quantized_embedding=True) + model = build_from_gguf(path, keep_quantized=True)["model"] + assert sum(node.op_type == "GatherBlockQuantized" for node in model.graph) == 1 + expected_projection_weights = { + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.output_gate.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.layers.1.self_attn.q_proj.weight", + "model.layers.1.self_attn.k_proj.weight", + "model.layers.1.self_attn.v_proj.weight", + "model.layers.1.self_attn.o_proj.weight", + *{ + f"model.layers.{layer}.mlp.experts.{expert}.{projection}_proj.weight" + for layer in range(2) + for expert in range(2) + for projection in ("gate", "up", "down") + }, + } + assert { + node.inputs[1].name for node in model.graph if node.op_type == "MatMulNBits" + } == expected_projection_weights + assert all( + "norm" not in node.inputs[1].name + for node in model.graph + if node.op_type == "MatMulNBits" + ) + session = OnnxModelSession(model) + try: + outputs = session.run(self._inputs(np.asarray([[1, 2]], np.int64))) + finally: + session.close() + assert np.isfinite(outputs["logits"]).all() + + def test_tied_output_uses_embedding_storage(self, tmp_path: Path) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "minimax-tied.gguf" + _write_minimax_gguf(path, quantized=False, omit="output.weight") + model = build_from_gguf(path)["model"] + assert "lm_head.weight" not in model.graph.initializers + assert "model.embed_tokens.weight" in model.graph.initializers + + def test_quantized_tied_output_uses_embedding_storage(self, tmp_path: Path) -> None: + from mobius._testing.ort_inference import OnnxModelSession + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "minimax-q4-tied.gguf" + _write_minimax_gguf( + path, + quantized=True, + quantized_embedding=True, + omit="output.weight", + ) + model = build_from_gguf(path, keep_quantized=True)["model"] + op_types = [node.op_type for node in model.graph] + assert op_types.count("GatherBlockQuantized") == 1 + assert op_types.count("MatMulNBits") == 20 + assert ( + sum(name.endswith("embed_tokens.qweight") for name in model.graph.initializers) + == 1 + ) + assert not any(name.startswith("lm_head.") for name in model.graph.initializers) + tied_head = next( + node for node in reversed(model.graph) if node.op_type == "MatMulNBits" + ) + assert tied_head.inputs[2].name == "model.embed_tokens.scales" + assert tied_head.inputs[3].name == "model.embed_tokens.zero_points" + session = OnnxModelSession(model) + try: + outputs = session.run(self._inputs(np.asarray([[1, 2]], np.int64))) + finally: + session.close() + assert np.isfinite(outputs["logits"]).all() + + def test_cli_build(self, tmp_path: Path) -> None: + from mobius.__main__ import main + from mobius._model_package import ModelPackage + + path = tmp_path / "minimax-cli.gguf" + output = tmp_path / "minimax-cli-output" + _write_minimax_gguf(path, quantized=False) + main(["build-gguf", str(path), "--output", str(output), "--dequantize"]) + assert "model" in ModelPackage.load(output) + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ( + {"omit": "blk.0.attn_gate.weight"}, + "missing=.*blk.0.attn_gate.weight", + ), + ( + {"extra": "blk.0.attn_q.weight"}, + "unexpected=.*blk.0.attn_q.weight", + ), + ( + {"extra": "blk.2.attn_q.weight"}, + "out_of_range=.*blk.2.attn_q.weight", + ), + ( + {"malformed_shape": "blk.1.attn_k.weight"}, + "invalid tensor shape", + ), + ( + {"recurrent_layers": [True]}, + "must contain exactly 2 entries", + ), + ( + {"norm_eps": 0.0}, + "inconsistent architecture metadata", + ), + ( + {"rope_freq_base": float("nan")}, + "inconsistent architecture metadata", + ), + ], + ) + def test_invalid_contract_rejected_before_graph( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + kwargs: dict, + match: str, + ) -> None: + from mobius import _builder as core_builder + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "minimax-invalid.gguf" + _write_minimax_gguf(path, quantized=False, **kwargs) + graph_build = mock.Mock(side_effect=AssertionError("graph construction reached")) + monkeypatch.setattr(core_builder, "build_from_module", graph_build) + with pytest.raises(ValueError, match=match): + build_from_gguf(path) + graph_build.assert_not_called() + + class TestGraniteHybridMoEGGUFBuild: """GraniteHybrid GGUF import preserves mixed state and routed expert order.""" @@ -5975,7 +6334,6 @@ def test_invalid_decay_fails_before_graph(self, tmp_path: Path, monkeypatch) -> @pytest.mark.parametrize( "architecture", [ - "minimax-01", "bailingmoe3", "deepseek4", "kimi-k3", diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 50e29e44c..51130019b 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -42,6 +42,7 @@ Lfm2MoeConfig, Mamba2Config, MambaConfig, + MiniMaxConfig, MuseGlimmerConfig, NemotronHConfig, Plamo2Config, @@ -207,6 +208,12 @@ "ssm.time_step_rank": "mamba_num_heads", } +_MINIMAX_KEY_MAP = { + "attention.key_length": "head_dim", + "attention.value_length": "value_head_dim", + "residual_scale": "residual_scale", +} + _T5_KEY_MAP = { "attention.key_length": "head_dim", "attention.relative_buckets_count": "relative_attention_num_buckets", @@ -229,6 +236,7 @@ "jamba": _JAMBA_KEY_MAP, "nemotron_h": _NEMOTRON_H_KEY_MAP, "granitehybrid": _GRANITEHYBRID_KEY_MAP, + "minimax": _MINIMAX_KEY_MAP, "t5": _T5_KEY_MAP, } ) @@ -437,6 +445,51 @@ def _derive_hybrid_layout( layer_types.append("full_attention") return trunk_layers, layer_types, mtp_count + if gguf_arch == "minimax-01": + recurrent_key = f"{gguf_arch}.attention.recurrent_layers" + raw_recurrent = metadata.get(recurrent_key) + if raw_recurrent is None: + interval = int(metadata.get(f"{gguf_arch}.full_attention_interval", 8)) + if interval <= 0: + raise ValueError( + f"{gguf_arch}.full_attention_interval must be positive, got {interval}" + ) + recurrent = [(layer + 1) % interval != 0 for layer in range(total_layers)] + else: + if not isinstance(raw_recurrent, (list, tuple, np.ndarray)): + raise ValueError(f"{recurrent_key} must be a boolean array") + if len(raw_recurrent) != total_layers: + raise ValueError( + f"{recurrent_key} must contain exactly {total_layers} entries, " + f"got {len(raw_recurrent)}" + ) + if any( + not isinstance(value, (bool, np.bool_, int, np.integer)) + or int(value) not in (0, 1) + for value in raw_recurrent + ): + raise ValueError(f"{recurrent_key} entries must be booleans or 0/1") + recurrent = [bool(value) for value in raw_recurrent] + if f"{gguf_arch}.full_attention_interval" in metadata: + interval = int(metadata[f"{gguf_arch}.full_attention_interval"]) + if interval <= 0: + raise ValueError( + f"{gguf_arch}.full_attention_interval must be positive, got {interval}" + ) + periodic = [(layer + 1) % interval != 0 for layer in range(total_layers)] + if recurrent != periodic: + raise ValueError( + "MiniMax-01 recurrent_layers contradicts full_attention_interval" + ) + return ( + trunk_layers, + [ + "lightning_attention" if value else "full_attention" + for value in recurrent[:trunk_layers] + ], + mtp_count, + ) + if gguf_arch not in _DELTA_NET_ARCHITECTURES: return trunk_layers, None, mtp_count @@ -625,6 +678,7 @@ def gguf_to_config( "nemotron_h_moe", "granitehybrid", "plamo2", + "minimax-01", }: nonzero = {value for value in values if value} if len(nonzero) != 1: @@ -2538,6 +2592,65 @@ def _t5_postprocess( ) +def _minimax_postprocess( + config: ArchitectureConfig, + metadata: dict[str, Any], + model: Any, +) -> MiniMaxConfig: + """Restore the exact pinned MiniMax-01 GGUF execution contract.""" + arch = model.architecture + head_dim = int(metadata[f"{arch}.attention.key_length"]) + value_dim = int(metadata[f"{arch}.attention.value_length"]) + rope_dim = int(metadata[f"{arch}.rope.dimension_count"]) + residual_scale = float(metadata[f"{arch}.residual_scale"]) + experts = int(metadata[f"{arch}.expert_count"]) + top_k = int(metadata[f"{arch}.expert_used_count"]) + if head_dim <= 0 or value_dim != head_dim: + raise ValueError( + f"MiniMax-01 requires equal positive key/value lengths, got {head_dim}/{value_dim}" + ) + if rope_dim <= 0 or rope_dim > head_dim or rope_dim % 2: + raise ValueError( + f"MiniMax-01 rope.dimension_count must be positive, even, and <= {head_dim}" + ) + if not math.isfinite(residual_scale) or residual_scale <= 0: + raise ValueError("MiniMax-01 residual_scale must be finite and positive") + if experts <= 1 or not 1 <= top_k <= experts: + raise ValueError( + f"MiniMax-01 expert counts are invalid: expert_count={experts}, " + f"expert_used_count={top_k}" + ) + if any( + key in metadata + for key in ( + f"{arch}.expert_shared_count", + f"{arch}.expert_shared_feed_forward_length", + ) + ): + raise ValueError("MiniMax-01 pinned GGUF does not support shared experts") + + _, layer_types, _ = _derive_hybrid_layout(arch, metadata, model.tensor_names) + assert layer_types is not None + fields = _shallow_fields(config) + fields.update( + model_type="minimax", + head_dim=head_dim, + partial_rotary_factor=rope_dim / head_dim, + layer_types=layer_types, + hidden_act="silu", + norm_topk_prob=True, + disable_qmoe=True, + lightning_norm_eps=config.rms_norm_eps, + full_attn_alpha_factor=residual_scale, + full_attn_beta_factor=1.0, + linear_attn_alpha_factor=residual_scale, + linear_attn_beta_factor=1.0, + mlp_alpha_factor=residual_scale, + mlp_beta_factor=1.0, + ) + return MiniMaxConfig(**fields) + + def _bert_encoder_postprocess( config: ArchitectureConfig, metadata: dict[str, Any], @@ -2677,6 +2790,7 @@ def _eagle3_postprocess( "bert_encoder": _bert_encoder_postprocess, "modern_bert_encoder": _modern_bert_encoder_postprocess, "t5": _t5_postprocess, + "minimax": _minimax_postprocess, } diff --git a/src/mobius/integrations/gguf/_config_mapping_test.py b/src/mobius/integrations/gguf/_config_mapping_test.py index d965cd8e8..50a2e68ef 100644 --- a/src/mobius/integrations/gguf/_config_mapping_test.py +++ b/src/mobius/integrations/gguf/_config_mapping_test.py @@ -1299,6 +1299,30 @@ def test_default_interval_is_four_and_mtp_is_full_attention(self) -> None: "full_attention", ] + def test_minimax_default_interval_is_eight(self) -> None: + from mobius.integrations.gguf._config_mapping import _derive_hybrid_layout + + layers, schedule, mtp = _derive_hybrid_layout( + "minimax-01", {"minimax-01.block_count": 9} + ) + + assert (layers, mtp) == (9, 0) + assert schedule == ["lightning_attention"] * 7 + [ + "full_attention", + "lightning_attention", + ] + + def test_minimax_explicit_schedule_rejects_invalid_interval(self) -> None: + from mobius.integrations.gguf._config_mapping import _derive_hybrid_layout + + metadata = { + "minimax-01.block_count": 2, + "minimax-01.attention.recurrent_layers": [True, False], + "minimax-01.full_attention_interval": 0, + } + with pytest.raises(ValueError, match="must be positive"): + _derive_hybrid_layout("minimax-01", metadata) + def test_dotted_nextn_metadata_is_rejected(self) -> None: from mobius.integrations.gguf._config_mapping import _derive_hybrid_layout diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index e72d198c6..8f771b2a9 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -726,6 +726,25 @@ "blk.{bid}.exp_probs_b": "model.layers.{bid}.feed_forward.expert_bias@", } +_MINIMAX_MAPPING: dict[str, str] = { + "token_embd": "model.embed_tokens", + "output": "lm_head", + "output_norm": "model.norm", + "blk.{bid}.attn_norm": "model.layers.{bid}.input_layernorm", + "blk.{bid}.attn_norm_2": "model.layers.{bid}.self_attn.norm", + "blk.{bid}.attn_qkv": "model.layers.{bid}.self_attn.qkv_proj", + "blk.{bid}.attn_q": "model.layers.{bid}.self_attn.q_proj", + "blk.{bid}.attn_k": "model.layers.{bid}.self_attn.k_proj", + "blk.{bid}.attn_v": "model.layers.{bid}.self_attn.v_proj", + "blk.{bid}.attn_gate": "model.layers.{bid}.self_attn.output_gate", + "blk.{bid}.attn_output": "model.layers.{bid}.self_attn.o_proj", + "blk.{bid}.ffn_norm": "model.layers.{bid}.post_attention_layernorm", + "blk.{bid}.ffn_gate_inp": "model.layers.{bid}.mlp.gate", + "blk.{bid}.ffn_gate_exps": "model.layers.{bid}.mlp.experts.gate_proj", + "blk.{bid}.ffn_up_exps": "model.layers.{bid}.mlp.experts.up_proj", + "blk.{bid}.ffn_down_exps": "model.layers.{bid}.mlp.experts.down_proj", +} + # Architectures sharing the llama HF naming convention are declared in # ``_arch_registry`` via ``tensor_map_recipe=("llama", ...)`` rather than by a # frozenset here, so the "which architectures does this cover?" question has one @@ -809,6 +828,7 @@ "qwen3next_hybrid_extras": _QWEN3NEXT_HYBRID_EXTRAS, "hunyuan_extras": _HUNYUAN_EXTRAS, "muse_glimmer_extras": _MUSE_GLIMMER_EXTRAS, + "minimax": _MINIMAX_MAPPING, } ) diff --git a/src/mobius/models/minimax.py b/src/mobius/models/minimax.py index 5b279060d..4008d2e99 100644 --- a/src/mobius/models/minimax.py +++ b/src/mobius/models/minimax.py @@ -19,27 +19,45 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING +import onnx_ir as ir import torch from onnxscript import OpBuilder, nn -from mobius._configs import ArchitectureConfig +from mobius._configs import ArchitectureConfig, MiniMaxConfig from mobius.components import ( Attention, - Embedding, MoELayer, RMSNorm, - SoftmaxTopKGate, create_attention_bias, initialize_rope, ) from mobius.components._lightning_attention import LightningAttention -from mobius.models.base import CausalLMModel -from mobius.models.moe import _rename_moe_expert_weights +from mobius.models.base import CausalLMModel, embedding_for_config +from mobius.models.moe import _quantized_linear_class, _rename_moe_expert_weights -if TYPE_CHECKING: - import onnx_ir as ir + +class MiniMaxTopKGate(nn.Module): + """MiniMax router: FP32 softmax, top-k selection, then selected-weight normalization.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + assert config.num_local_experts is not None + assert config.num_experts_per_tok is not None + self.top_k = config.num_experts_per_tok + self.weight = nn.Parameter([config.num_local_experts, config.hidden_size]) + + def forward(self, op: OpBuilder, hidden_states: ir.Value): + logits = op.MatMul(hidden_states, op.Transpose(self.weight, perm=[1, 0])) + probs = op.Softmax(op.Cast(logits, to=ir.DataType.FLOAT), axis=-1) + weights, experts = op.TopK( + probs, + op.Constant(value_ints=[self.top_k]), + axis=-1, + _outputs=2, + ) + weights = op.Div(weights, op.ReduceSum(weights, [-1], keepdims=True)) + return op.CastLike(weights, hidden_states), experts class MiniMaxDecoderLayer(nn.Module): @@ -59,30 +77,27 @@ class MiniMaxDecoderLayer(nn.Module): def __init__(self, config: ArchitectureConfig, layer_idx: int): super().__init__() - layer_types = config.layer_types or [] - self.layer_type: str = ( - layer_types[layer_idx] if layer_idx < len(layer_types) else "full_attention" - ) + assert config.layer_types is not None + self.layer_type = config.layer_types[layer_idx] + linear_class = _quantized_linear_class(config) if self.layer_type == "lightning_attention": - self.self_attn = LightningAttention(config, layer_idx) + self.self_attn = LightningAttention(config, layer_idx, linear_class=linear_class) self._attn_alpha: float = getattr(config, "linear_attn_alpha_factor", 1.0) self._attn_beta: float = getattr(config, "linear_attn_beta_factor", 1.0) else: - self.self_attn = Attention(config) + self.self_attn = Attention(config, linear_class=linear_class) self._attn_alpha = getattr(config, "full_attn_alpha_factor", 1.0) self._attn_beta = getattr(config, "full_attn_beta_factor", 1.0) self._mlp_alpha: float = getattr(config, "mlp_alpha_factor", 1.0) self._mlp_beta: float = getattr(config, "mlp_beta_factor", 1.0) - gate = SoftmaxTopKGate( - config.hidden_size, - config.num_local_experts, - config.num_experts_per_tok, - norm_topk_prob=config.norm_topk_prob, + self.mlp = MoELayer( + config, + gate=MiniMaxTopKGate(config), + linear_class=linear_class, ) - self.mlp = MoELayer(config, gate=gate) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -93,6 +108,7 @@ def forward( attention_bias: ir.Value | None, position_embeddings: tuple[ir.Value, ir.Value], past_key_value, + attention_mask: ir.Value | None = None, ): # MiniMax pre-norm: apply layernorm first, then take residual from # the normalized value (not from the original hidden_states). @@ -102,7 +118,9 @@ def forward( if self.layer_type == "lightning_attention": # Lightning Attention: single recurrent state (no conv_state) (recurrent_state,) = past_key_value - attn_out, new_state = self.self_attn(op, hidden_states, recurrent_state) + attn_out, new_state = self.self_attn( + op, hidden_states, recurrent_state, attention_mask + ) present_key_value = (new_state,) else: attn_out, present_key_value = self.self_attn( @@ -132,10 +150,17 @@ class MiniMaxTextModel(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() + if config.layer_types is None or len(config.layer_types) != config.num_hidden_layers: + raise ValueError( + "MiniMax layer_types must contain exactly num_hidden_layers entries" + ) + unknown = set(config.layer_types) - {"full_attention", "lightning_attention"} + if unknown: + raise ValueError(f"Unsupported MiniMax layer type(s): {sorted(unknown)}") + if config.head_dim <= 0: + raise ValueError("MiniMax head_dim must be explicit and positive") self._dtype = config.dtype - self.embed_tokens = Embedding( - config.vocab_size, config.hidden_size, config.pad_token_id - ) + self.embed_tokens = embedding_for_config(config) self.layers = nn.ModuleList( [MiniMaxDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)] ) @@ -169,6 +194,7 @@ def forward( attention_bias=attention_bias, position_embeddings=position_embeddings, past_key_value=past_kv, + attention_mask=attention_mask, ) present_key_values.append(present_kv) @@ -180,9 +206,9 @@ class MiniMaxCausalLMModel(CausalLMModel): """MiniMax causal language model with hybrid Lightning + GQA + MoE. Architecture: - - Even layers: full GQA attention + Sparse MoE FFN - - Odd layers: Lightning Attention + Sparse MoE FFN (no KV cache, fixed-size recurrent state) - - ``config.layer_types`` selects per-layer attention type + - ``config.layer_types`` preserves the checkpoint's explicit per-layer schedule + - Full-attention layers use partial-RoPE GQA + - Lightning layers use a fixed-size recurrent state and no KV cache Lightning Attention layers carry a single recurrent_state tensor of shape (B, num_heads, head_dim, head_dim) per layer. Full attention layers @@ -190,20 +216,35 @@ class MiniMaxCausalLMModel(CausalLMModel): Task: ``hybrid-text-generation`` (HybridCausalLMTask). - HuggingFace model_type: ``"minimax"`` + HuggingFace model types: ``"MiniMaxText01"`` and ``"minimax"`` """ default_task: str = "hybrid-text-generation" + config_class = MiniMaxConfig def __init__(self, config: ArchitectureConfig): super().__init__(config) - self.model = MiniMaxTextModel(config) + self._replace_text_model(MiniMaxTextModel(config)) def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: # Split fused MoE expert weights (gate_up_proj/down_proj → per-expert). - state_dict = _rename_moe_expert_weights(state_dict) + renamed: dict[str, torch.Tensor] = {} + for name, tensor in state_dict.items(): + name = name.replace(".self_attn.out_proj.", ".self_attn.o_proj.") + if ".experts." in name and tensor.dim() == 3: + prefix, projection = name.rsplit(".experts.", 1) + if projection in { + "gate_proj.weight", + "up_proj.weight", + "down_proj.weight", + }: + for expert, expert_tensor in enumerate(tensor): + renamed[f"{prefix}.experts.{expert}.{projection}"] = expert_tensor + continue + renamed[name] = tensor + state_dict = _rename_moe_expert_weights(renamed) return super().preprocess_weights(state_dict) diff --git a/testdata/cases/causal-lm/minimax-text-01.yaml b/testdata/cases/causal-lm/minimax-text-01.yaml index f17c34866..b20584180 100644 --- a/testdata/cases/causal-lm/minimax-text-01.yaml +++ b/testdata/cases/causal-lm/minimax-text-01.yaml @@ -1,6 +1,6 @@ model_id: "MiniMaxAI/MiniMax-Text-01" model_type: "minimax" -revision: "main" +revision: "a7351bf2bee0e1253919d349f1ad304e6dac13e9" task_type: "text-generation" dtype: "float32" @@ -15,4 +15,4 @@ generation: do_sample: false skip_reason: "MiniMax-Text-01 is 456B total — too large for CI golden data generation." -notes: "MiniMax Text-01. Hybrid LightningAttention (linear, gated) + full GQA + MoE. FP accumulation differences (atol=0.05)." +notes: "Pinned MiniMax Text-01 reference. Exact GGUF graph import covers interval-8 LightningAttention/full-GQA scheduling, normalized residual scaling, partial RoPE, routed MoE, and heterogeneous explicit state. Released ORT GenAI packaging is deferred to #605 because rollback snapshots are not representable." diff --git a/tests/_test_configs.py b/tests/_test_configs.py index f07350ea0..0501fa2bf 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -40,6 +40,7 @@ LongcatFlashConfig, Mamba2Config, MambaConfig, + MiniMaxConfig, MllamaConfig, MoonshineConfig, MuseGlimmerConfig, @@ -1245,6 +1246,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ( "minimax", { + "_config_cls": MiniMaxConfig, "layer_types": ["full_attention", "lightning_attention"], "num_local_experts": 4, "num_experts_per_tok": 2, @@ -1252,6 +1254,28 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + ( + "MiniMaxText01", + { + "_config_cls": MiniMaxConfig, + "layer_types": ["full_attention", "lightning_attention"], + "num_local_experts": 4, + "num_experts_per_tok": 2, + "head_dim": TINY_HIDDEN // TINY_HEADS, + }, + False, + ), + ( + "minimax_text_01", + { + "_config_cls": MiniMaxConfig, + "layer_types": ["full_attention", "lightning_attention"], + "num_local_experts": 4, + "num_experts_per_tok": 2, + "head_dim": TINY_HIDDEN // TINY_HEADS, + }, + False, + ), ( "gpt_oss", {