[1/4] Add MiniMax-M3 support

miles/backends/megatron_utils/megatron_to_hf/__init__.py

@@ -9,6 +9,7 @@
 from .kimi_vl import convert_kimi_k25_to_hf, convert_kimivl_to_hf
 from .llama import convert_llama_to_hf
 from .mimo import convert_mimo_to_hf
+from .minimax_m3 import convert_minimax_m3_to_hf
 from .processors import quantize_params, remove_padding
 from .qwen2 import convert_qwen2_to_hf
 from .qwen3_5 import convert_qwen3_5_to_hf
@@ -77,6 +78,8 @@ def _convert_to_hf_core(args, model_name, name, param):
         converted_named_tensors = convert_inkling_to_hf(args, name, param)
     elif "llama" in model_name:
         converted_named_tensors = convert_llama_to_hf(args.backend, name, param)
+    elif "minimax_m3" in model_name:
+        converted_named_tensors = convert_minimax_m3_to_hf(args.backend, name, param)
     elif "mimo" in model_name:
         converted_named_tensors = convert_mimo_to_hf(args.backend, name, param)
     elif "kimi_k3" in model_name:

miles/backends/megatron_utils/megatron_to_hf/minimax_m3.py

@@ -0,0 +1,71 @@
+import re
+
+import torch
+
+_TEXT_MODEL = "language_model.model"
+
+_DIRECT = {
+    "module.module.embedding.word_embeddings.weight": f"{_TEXT_MODEL}.embed_tokens.weight",
+    "module.module.output_layer.weight": "language_model.lm_head.weight",
+    "module.module.decoder.final_layernorm.weight": f"{_TEXT_MODEL}.norm.weight",
+}
+
+_LAYER_RENAMES = {
+    "input_layernorm.weight": "input_layernorm.weight",
+    "pre_mlp_layernorm.weight": "post_attention_layernorm.weight",
+    "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight",
+    "self_attention.linear_proj.weight": "self_attn.o_proj.weight",
+    "self_attention.q_layernorm.weight": "self_attn.q_norm.weight",
+    "self_attention.k_layernorm.weight": "self_attn.k_norm.weight",
+    "self_attention.indexer.index_q_proj.weight": "self_attn.index_q_proj.weight",
+    "self_attention.indexer.index_k_proj.weight": "self_attn.index_k_proj.weight",
+    "self_attention.indexer.index_q_norm": "self_attn.index_q_norm.weight",
+    "self_attention.indexer.index_k_norm": "self_attn.index_k_norm.weight",
+    "mlp.linear_fc2.weight": "mlp.down_proj.weight",
+    "mlp.router.weight": "block_sparse_moe.gate.weight",
+    "mlp.router.expert_bias": "block_sparse_moe.e_score_correction_bias",
+    "mlp.shared_experts.linear_fc2.weight": "block_sparse_moe.shared_experts.down_proj.weight",
+}
+
+
+def convert_minimax_m3_to_hf(args, name, param):
+    if name in _DIRECT:
+        return [(_DIRECT[name], param)]
+
+    match = re.match(r"module\.module\.decoder\.layers\.(\d+)\.(.+)", name)
+    if match is None:
+        raise ValueError(f"Unknown parameter name: {name}")
+    layer_idx, rest = match.groups()
+    prefix = f"{_TEXT_MODEL}.layers.{layer_idx}"
+
+    expert = re.match(r"mlp\.experts\.(linear_fc1|linear_fc2)\.weight(\d+)", rest)
+    if expert:
+        kind, expert_idx = expert.groups()
+        base = f"{prefix}.block_sparse_moe.experts.{expert_idx}"
+        if kind == "linear_fc1":
+            gate, up = param.chunk(2, dim=0)
+            return [(f"{base}.w1.weight", gate), (f"{base}.w3.weight", up)]
+        return [(f"{base}.w2.weight", param)]
+
+    if rest == "mlp.shared_experts.linear_fc1.weight":
+        gate, up = param.chunk(2, dim=0)
+        return [
+            (f"{prefix}.block_sparse_moe.shared_experts.gate_proj.weight", gate),
+            (f"{prefix}.block_sparse_moe.shared_experts.up_proj.weight", up),
+        ]
+    if rest == "mlp.linear_fc1.weight":
+        gate, up = param.chunk(2, dim=0)
+        return [(f"{prefix}.mlp.gate_proj.weight", gate), (f"{prefix}.mlp.up_proj.weight", up)]
+    if rest == "self_attention.linear_qkv.weight":
+        head_dim = args.kv_channels
+        per_group = args.num_attention_heads // args.num_query_groups
+        param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
+        q, k, v = torch.split(param, [per_group, 1, 1], dim=1)
+        return [
+            (f"{prefix}.self_attn.q_proj.weight", q.reshape(-1, args.hidden_size)),
+            (f"{prefix}.self_attn.k_proj.weight", k.reshape(-1, args.hidden_size)),
+            (f"{prefix}.self_attn.v_proj.weight", v.reshape(-1, args.hidden_size)),
+        ]
+    if rest in _LAYER_RENAMES:
+        return [(f"{prefix}.{_LAYER_RENAMES[rest]}", param)]
+    raise ValueError(f"Unknown parameter name: {name}")

miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py

@@ -402,7 +402,7 @@ def _check_and_fix_partition(args: Namespace, name: str, partition_stride: int,
     (GLU/SwiGLU interleaved [gate, up]), so assert partition_stride==2 is removed.
     But TEGroupedLinear still does not set partition_stride/partition_dim correctly for grouped moe gemm
     """
-    if "linear_fc1.weight" in name and args.backend.swiglu:
+    if "linear_fc1.weight" in name and (args.backend.swiglu or args.backend.quick_geglu):
         partition_stride = 2
         if partition_dim < 0:
             partition_dim = 0

miles/kernels/attention/qsa/block_sparse_attention.py

@@ -347,7 +347,7 @@ def selection_to_block_bitmap(indices: Tensor, num_tokens: int, block_size: int)
     return flags
 
 
-def _tile_mask(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int) -> Tensor:
+def _tile_mask(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int, block_size: int) -> Tensor:
     """``[NQT, NKT]`` bool: some query in query tile ``i`` selected a block with a token in key tile ``j``.
 
     Key tiles are absolute token ranges ``[j * bk, (j + 1) * bk)``, as the kernels walk them.
@@ -362,9 +362,11 @@ def _tile_mask(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk
         sel = torch.nn.functional.pad(sel, (0, 0, 0, pad_q))
     picked = sel.view(nqt, bq, nb).amax(dim=1).to(torch.int32)
     tile = torch.zeros(nqt, nkt, dtype=torch.int32, device=sel.device)
-    # block_size <= bk, so a block touches at most the tiles of its first and last token
-    tile.index_add_(1, block_first.long() // bk, picked)
-    tile.index_add_(1, block_last.long() // bk, picked)
+    # a block wider than a tile, or off the tile grid, touches every tile between its first and last token
+    first_tile = block_first.long() // bk
+    last_tile = block_last.long() // bk
+    for off in range(-(-block_size // bk) + 1):
+        tile.index_add_(1, torch.minimum(first_tile + off, last_tile), picked)
     return tile > 0
 
 
@@ -376,22 +378,24 @@ def _compact_tile_mask(mat: Tensor) -> tuple[Tensor, Tensor]:
     return order[:, :maxc].contiguous().to(torch.int32), cnt
 
 
-def build_tile_index(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int) -> tuple[Tensor, Tensor]:
+def build_tile_index(
+    sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int, block_size: int
+) -> tuple[Tensor, Tensor]:
     """``[T, NB]`` block flags -> (``klist`` [NQT, maxc] int32, ``kcnt`` [NQT] int32).
 
     ``klist[i]`` lists, ascending, the key tiles that at least one query in query-tile
     ``i`` selected. Per-query exactness still comes from the in-kernel mask; this only
     decides which tiles are worth visiting.
     """
-    return _compact_tile_mask(_tile_mask(sel, block_first, block_last, bq, bk))
+    return _compact_tile_mask(_tile_mask(sel, block_first, block_last, bq, bk, block_size))
 
 
-def build_tile_index_pair(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int):
+def build_tile_index_pair(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int, block_size: int):
     """Both directions of the tile map: (klist, kcnt) per query tile, (qlist, qcnt) per key tile.
 
     The transposed half is what lets dK/dV be keyed on the key tile and so avoid atomics.
     """
-    tile = _tile_mask(sel, block_first, block_last, bq, bk)
+    tile = _tile_mask(sel, block_first, block_last, bq, bk, block_size)
     klist, kcnt = _compact_tile_mask(tile)
     qlist, qcnt = _compact_tile_mask(tile.t().contiguous())
     return klist, kcnt, qlist, qcnt
@@ -407,7 +411,7 @@ def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, block_first, block_la
         qc, kc, vc = q.contiguous(), k.contiguous(), v.contiguous()
         selc = sel.contiguous()
         BQ, BK = 64, 64
-        klist, kcnt = build_tile_index(selc, block_first, block_last, BQ, BK)
+        klist, kcnt = build_tile_index(selc, block_first, block_last, BQ, BK, block_size)
         o = torch.empty(T, Hq, D, device=q.device, dtype=torch.float32)
         lse = torch.empty(Hq, T, device=q.device, dtype=torch.float32)
         grid = (triton.cdiv(T, BQ), Hq)
@@ -463,7 +467,7 @@ def backward(ctx, grad_out):
         dv = torch.zeros(vc.shape, device=vc.device, dtype=torch.float32)
 
         BQ, BK = 64, 32
-        klist, kcnt, qlist, qcnt = build_tile_index_pair(selc, block_first, block_last, BQ, BK)
+        klist, kcnt, qlist, qcnt = build_tile_index_pair(selc, block_first, block_last, BQ, BK, ctx.block_size)
         common = (
             qc.stride(0),
             qc.stride(1),

miles/utils/arguments.py

@@ -1855,12 +1855,17 @@ def equal(x, y):
     if (hf_config.model_type == "deepseek_v4" or "deepseekv4" in model_name) and args.context_parallel_size > 1:
         assert args.allgather_cp, "zigzag CP is not supported for DeepSeek V4."
 
+    # MiniMax-M3 names the dense width dense_intermediate_size and the expert width intermediate_size
+    if hasattr(hf_config, "dense_intermediate_size"):  # config-access-exempt: model-specific key
+        dense_ffn_key, moe_ffn_key = "dense_intermediate_size", "intermediate_size"
+    else:
+        dense_ffn_key, moe_ffn_key = "intermediate_size", "moe_intermediate_size"
     for hf_config_name, megatron_config_name, compare_fn in [
         ("hidden_size", "hidden_size", equal),
         ("num_attention_heads", "num_attention_heads", equal),
         ("num_hidden_layers", "num_layers", equal),
-        ("intermediate_size", "ffn_hidden_size", equal),
-        ("moe_intermediate_size", "moe_ffn_hidden_size", equal),
+        (dense_ffn_key, "ffn_hidden_size", equal),
+        (moe_ffn_key, "moe_ffn_hidden_size", equal),
         ("tie_word_embeddings", "untie_embeddings_and_output_weights", lambda x, y: not x == y),
         (
             "rms_norm_eps",

miles/utils/hf_utils/config.py

@@ -15,7 +15,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-from transformers import AutoConfig, AutoModelForCausalLM
+from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES
 
 
@@ -130,6 +130,9 @@ def load_hf_config(
     """
     register_hf_config_aliases()
     config = AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=trust_remote_code, **autoconfig_kwargs)
+    # MiniMax-M3's remote config class keeps the nested text config as a plain dict
+    if isinstance(getattr(config, "text_config", None), dict):  # config-access-exempt: optional key
+        config.text_config = PretrainedConfig.from_dict(config.text_config)
 
     if overrides:
         for key, value in overrides.items():

miles_plugins/mbridge/__init__.py

@@ -11,6 +11,7 @@
 from .kimi_k3 import KimiK3Bridge
 from .kimi_k25 import KimiK25Bridge
 from .mimo import MimoBridge
+from .minimax_m3 import MiniMaxM3Bridge
 from .qwen3_5 import Qwen3_5Bridge
 from .qwen3_8_next import Qwen38NextBridge
 from .qwen3_next import Qwen3NextBridge
@@ -23,6 +24,7 @@
     "Qwen3_5Bridge",
     "Qwen38NextBridge",
     "MimoBridge",
+    "MiniMaxM3Bridge",
     "DeepseekV32Bridge",
     "Glm5NextBridge",
     "DeepseekV4Bridge",

miles_plugins/mbridge/minimax_m3.py

@@ -0,0 +1,111 @@
+"""Bridge for MiniMax-M3 (HF ``MiniMaxM3SparseForConditionalGeneration``, text part only).
+
+Text weights live under ``language_model.``; the vision tower and projector are never
+created on the Megatron side.
+"""
+
+import torch
+from transformers import PretrainedConfig
+
+from mbridge.core import LLMBridge, register_model
+from miles_plugins.models.minimax_m3.minimax_m3 import apply_minimax_m3_config
+
+_LAYER = "language_model.model.layers.{layer_number}"
+
+
+@register_model(["minimax_m3_vl", "minimax_m3"])
+class MiniMaxM3Bridge(LLMBridge):
+    _CONFIG_MAPPING = {**LLMBridge._CONFIG_MAPPING, "attention_dropout": ("attention_dropout", 0.0)}
+
+    _DIRECT_MAPPING = {
+        "embedding.word_embeddings.weight": "language_model.model.embed_tokens.weight",
+        "decoder.final_layernorm.weight": "language_model.model.norm.weight",
+        "output_layer.weight": "language_model.lm_head.weight",
+    }
+
+    _ATTENTION_MAPPING = {
+        "self_attention.linear_proj.weight": [f"{_LAYER}.self_attn.o_proj.weight"],
+        "self_attention.q_layernorm.weight": [f"{_LAYER}.self_attn.q_norm.weight"],
+        "self_attention.k_layernorm.weight": [f"{_LAYER}.self_attn.k_norm.weight"],
+        "self_attention.linear_qkv.weight": [
+            f"{_LAYER}.self_attn.q_proj.weight",
+            f"{_LAYER}.self_attn.k_proj.weight",
+            f"{_LAYER}.self_attn.v_proj.weight",
+        ],
+        "self_attention.indexer.index_q_proj.weight": [f"{_LAYER}.self_attn.index_q_proj.weight"],
+        "self_attention.indexer.index_k_proj.weight": [f"{_LAYER}.self_attn.index_k_proj.weight"],
+        "self_attention.indexer.index_q_norm": [f"{_LAYER}.self_attn.index_q_norm.weight"],
+        "self_attention.indexer.index_k_norm": [f"{_LAYER}.self_attn.index_k_norm.weight"],
+    }
+
+    # first match wins, so the per-expert and shared-expert keys precede the dense ones
+    _MLP_MAPPING = {
+        "mlp.experts.linear_fc1": [
+            f"{_LAYER}.block_sparse_moe.experts.{{expert_id}}.w1.weight",
+            f"{_LAYER}.block_sparse_moe.experts.{{expert_id}}.w3.weight",
+        ],
+        "mlp.experts.linear_fc2": [f"{_LAYER}.block_sparse_moe.experts.{{expert_id}}.w2.weight"],
+        "mlp.shared_experts.linear_fc1.weight": [
+            f"{_LAYER}.block_sparse_moe.shared_experts.gate_proj.weight",
+            f"{_LAYER}.block_sparse_moe.shared_experts.up_proj.weight",
+        ],
+        "mlp.shared_experts.linear_fc2.weight": [f"{_LAYER}.block_sparse_moe.shared_experts.down_proj.weight"],
+        "mlp.router.expert_bias": [f"{_LAYER}.block_sparse_moe.e_score_correction_bias"],
+        "mlp.router.weight": [f"{_LAYER}.block_sparse_moe.gate.weight"],
+        "pre_mlp_layernorm.weight": [f"{_LAYER}.post_attention_layernorm.weight"],
+        "mlp.linear_fc1.layer_norm_weight": [f"{_LAYER}.post_attention_layernorm.weight"],
+        "mlp.linear_fc1.weight": [f"{_LAYER}.mlp.gate_proj.weight", f"{_LAYER}.mlp.up_proj.weight"],
+        "mlp.linear_fc2.weight": [f"{_LAYER}.mlp.down_proj.weight"],
+    }
+
+    _OTHER_MAPPING = {
+        "input_layernorm.weight": [f"{_LAYER}.input_layernorm.weight"],
+    }
+
+    def __init__(self, hf_config, *args, **kwargs):
+        text_config = getattr(hf_config, "text_config", None) or hf_config  # config-access-exempt: optional key
+        # the remote config class keeps the nested text config as a plain dict
+        if isinstance(text_config, dict):
+            text_config = PretrainedConfig.from_dict(text_config)
+        super().__init__(text_config, *args, **kwargs)
+
+    def _weight_name_mapping_mlp(self, name: str) -> list[str]:
+        layer_number = name.split(".")[2]
+        for keyword, targets in self._MLP_MAPPING.items():
+            if keyword in name:
+                expert_id = name.split("weight")[-1] if "{expert_id}" in targets[0] else None
+                return [t.format(layer_number=layer_number, expert_id=expert_id) for t in targets]
+        raise NotImplementedError(f"Unsupported parameter name: {name}")
+
+    def _weight_to_mcore_format(self, mcore_weights_name: str, hf_weights: list[torch.Tensor]) -> torch.Tensor:
+        if mcore_weights_name.endswith("mlp.router.expert_bias"):
+            return hf_weights[0].to(torch.float32)
+        return super()._weight_to_mcore_format(mcore_weights_name, hf_weights)
+
+    def _build_config(self):
+        cfg = self.hf_config
+        config = self._build_base_config(
+            use_cpu_initialization=False,
+            num_moe_experts=cfg.num_local_experts,
+            moe_ffn_hidden_size=cfg.intermediate_size,
+            ffn_hidden_size=cfg.dense_intermediate_size,
+            moe_shared_expert_intermediate_size=cfg.shared_intermediate_size * cfg.n_shared_experts,
+            moe_router_topk=cfg.num_experts_per_tok,
+            moe_router_score_function="sigmoid",
+            moe_router_enable_expert_bias=True,
+            moe_router_bias_update_rate=0.0,
+            moe_router_topk_scaling_factor=cfg.routed_scaling_factor,
+            moe_router_load_balancing_type="none",
+            moe_router_dtype="fp32",
+            moe_grouped_gemm=True,
+            moe_layer_freq=list(cfg.moe_layer_freq),
+            qk_layernorm=True,
+            layernorm_zero_centered_gamma=True,
+            add_qkv_bias=False,
+            add_bias_linear=False,
+            kv_channels=cfg.head_dim,
+            rotary_interleaved=False,
+        )
+        apply_minimax_m3_config(config, cfg)
+        config.activation_func_clamp_value = cfg.swiglu_limit
+        return config