[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/msa/__init__.py

@@ -0,0 +1,3 @@
+from miles.kernels.attention.msa.block_max_scores import msa_block_max_scores
+
+__all__ = ["msa_block_max_scores"]

miles/kernels/attention/msa/block_max_scores.py

@@ -0,0 +1,83 @@
+"""MSA block scores for training: each query's maximum logit over every causal key block of its sequence."""
+
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+
+_BQ = 64
+
+
+@triton.jit
+def _block_max_scores_kernel(
+    Q,
+    K,
+    TILES,
+    OUT,
+    stride_qt,
+    stride_qh,
+    stride_kt,
+    stride_oh,
+    stride_ot,
+    scale,
+    D: tl.constexpr,
+    BQ: tl.constexpr,
+    BLK: tl.constexpr,
+):
+    tile = tl.program_id(0)
+    head = tl.program_id(1)
+    q_start = tl.load(TILES + tile * 3)
+    seq_start = tl.load(TILES + tile * 3 + 1)
+    seq_end = tl.load(TILES + tile * 3 + 2)
+
+    offs_q = q_start + tl.arange(0, BQ)
+    offs_d = tl.arange(0, D)
+    q_in = offs_q < seq_end
+    q = tl.load(Q + offs_q[:, None] * stride_qt + head * stride_qh + offs_d[None, :], mask=q_in[:, None], other=0.0)
+
+    # a tile never straddles sequences, so every query shares the key blocks; the last query bounds them
+    last_block = (tl.minimum(q_start + BQ, seq_end) - 1 - seq_start) // BLK
+    for block in range(0, last_block + 1):
+        offs_k = seq_start + block * BLK + tl.arange(0, BLK)
+        k_in = offs_k < seq_end
+        k = tl.load(K + offs_k[:, None] * stride_kt + offs_d[None, :], mask=k_in[:, None], other=0.0)
+        logits = tl.dot(q, tl.trans(k))
+        causal = (offs_k[None, :] <= offs_q[:, None]) & k_in[None, :]
+        best = tl.max(tl.where(causal, logits, float("-inf")), axis=1) * scale
+        tl.store(OUT + head * stride_oh + offs_q * stride_ot + block, best, mask=q_in)
+
+
+def msa_block_max_scores(q: Tensor, k: Tensor, cu_seqlens: Tensor, block_size: int, scale: float) -> Tensor:
+    """``q`` [T, H, D], ``k`` [T, D] -> ``[H, T, NB]`` fp32 block scores, ``-inf`` where no key is causal.
+
+    ``NB`` is the block count of the longest sequence; blocks start at each sequence's first token.
+    """
+    T, H, D = q.shape
+    bounds = cu_seqlens.tolist()
+    tiles = [
+        (q_start, start, end)
+        for start, end in zip(bounds[:-1], bounds[1:], strict=True)
+        for q_start in range(start, end, _BQ)
+    ]
+    num_blocks = max(-(-(end - start) // block_size) for start, end in zip(bounds[:-1], bounds[1:], strict=True))
+    tiles = torch.tensor(tiles, dtype=torch.int32, device=q.device)
+    out = torch.full((H, T, num_blocks), float("-inf"), dtype=torch.float32, device=q.device)
+    qc, kc = q.contiguous(), k.contiguous()
+    _block_max_scores_kernel[(tiles.shape[0], H)](
+        qc,
+        kc,
+        tiles,
+        out,
+        qc.stride(0),
+        qc.stride(1),
+        kc.stride(0),
+        out.stride(0),
+        out.stride(1),
+        scale,
+        D=D,
+        BQ=_BQ,
+        BLK=block_size,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out