MiniMax-M3.1-Flash support

docs/docs.json

@@ -208,6 +208,7 @@
        },
        "models/gpt-oss/gpt-oss",
        "models/mimo/mimo-v2-6-flash",
+       "models/minimax/minimax-m3-1-flash",
        "models/stepfun/step5",
        "models/muse-spark/muse-spark-open",
        "models/reflection/beam"

docs/models/index.md

@@ -23,6 +23,7 @@ Each model name links to its recipe page.
 | **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) |
 | **GPT-OSS** | [gpt-oss-20b](/models/gpt-oss/gpt-oss) |
 | **MiMo** | [MiMo-V2.6-Flash](/models/mimo/mimo-v2-6-flash) |
+| **MiniMax** | [MiniMax-M3.1-Flash](/models/minimax/minimax-m3-1-flash) |
 | **StepFun** | [Step-5-Preview](/models/stepfun/step5) |
 | **Muse Spark** | [Muse Spark Open](/models/muse-spark/muse-spark-open) |
 | **Reflection** | [Beam-1.0-501B-A23B](/models/reflection/beam) |

docs/models/minimax/minimax-m3-1-flash.md

@@ -0,0 +1,98 @@
+---
+title: MiniMax-M3.1-Flash
+sidebarTitle: MiniMax-M3.1
+description: Launch recipe for MiniMax-M3.1-Flash (MoE with MSA block-sparse attention) — BF16 train / BF16 rollout, colocated, on GB300.
+---
+## 1. Model Introduction
+
+[MiniMax-M3.1-Flash](https://huggingface.co/MiniMaxAI/MiniMax-M3.1-Flash) is MiniMax's MoE model with block-sparse attention on every layer. miles trains its text decoder.
+
+**Key highlights:**
+
+- **MSA block-sparse attention on all 60 layers**: a 4-head indexer (128-dim, its own gemma q/k norm and RoPE) scores every causal 128-token block by its maximum logit, and each query attends to the 16 best blocks, its own block always included. Index head `h` selects the blocks for KV head `h`.
+- **Attention**: GQA with 64 query and 4 KV heads of dim 128, per-head gemma q/k norm, RoPE on the first 64 dims (base 1e7).
+- **MoE**: layers 0–2 are dense (FFN 12288); layers 3–59 route to 4 of 128 experts (FFN 3072) with sigmoid scores, a score-correction bias and a 2.0 scaling factor, plus one shared expert.
+- **Activation**: `swigluoai` — `clamp(g, max=7) * sigmoid(1.702 g) * (clamp(u, -7, 7) + 1)` in every MLP.
+- **Compressed checkpoint**: MXFP8 linears and NVFP4 routed experts (252 GB). Training and the BF16 rollout read a BF16 conversion (852 GB).
+
+## 2. Supported Variants
+
+| `--model-name` | What it is | GPUs |
+|---|---|---|
+| `MiniMax-M3.1-Flash-bf16` | text decoder, 60 layers | 32 (8 × 4 GB300) |
+| `MiniMax-M3.1-Flash-4layer-bf16` | layers 0, 1, 3, 4: two dense and two MoE layers, every layer sparse | 4 (1 × 4 GB300) |
+
+Both are converted from [MiniMaxAI/MiniMax-M3.1-Flash](https://huggingface.co/MiniMaxAI/MiniMax-M3.1-Flash). The vision tower and the `dspark/` draft model stay in the checkpoint but are not trained (text-only RL).
+
+## 3. Environment Setup
+
+### 3.1 Download + BF16 conversion
+
+`prepare` (run by the launcher before training) downloads the release to `<model-dir>/MiniMax-M3.1-Flash`, converts it unless `<model-dir>/<model-name>` already holds a finished conversion, converts that to the Megatron reference checkpoint `<ckpt-dir>/<megatron-model-type>_torch_dist`, and fetches dapo-math-17k:
+
+```bash
+hf download MiniMaxAI/MiniMax-M3.1-Flash --local-dir /root/models/MiniMax-M3.1-Flash
+python tools/convert_minimax_m3_to_bf16.py --model-dir /root/models/MiniMax-M3.1-Flash \
+    --save-dir /root/models/MiniMax-M3.1-Flash-4layer-bf16 --device cuda --layers 0,1,3,4
+```
+
+The converter dequantizes the MXFP8 linears (one ue8m0 scale per 32 inputs) and the NVFP4 experts (an e4m3 scale per 16 inputs and a per-tensor encode scale) and drops `quantization_config`. Without `--layers` it converts the full model. The full model's torch_dist conversion is sharded PP4 × EP8 over the 8 nodes, since one node runs out of host memory saving it.
+
+### 3.2 SGLang
+
+The colocated BF16 engine needs online weight updates into MiniMax-M3's fused q/k/v + index GEMM: after loading, every sparse layer fuses `qkv_proj` and `index_qkv_proj` into one GEMM and frees the originals, so without the update the first weight sync fails with `IndexError: start out of range`. `sglang-miles` has it since (sglang-miles PR link), and the `radixark/miles:dev` image built from that branch includes it.
+
+## 4. Launch
+
+### 4.1 Four-layer model on one node
+
+```bash
+python scripts/run_minimax_m3_1_flash.py --model-name MiniMax-M3.1-Flash-4layer-bf16 --num-nodes 1
+```
+
+On 4 × GB300 the BF16 rollout and the trainer agree to a mean absolute log-prob difference of 0.0052, and the bitwise weight check (`--check-weight-update-equal`, on by default for this variant) passes after every update. The model is too shallow to solve the prompts, so every reply runs to the length limit and the reward stays 0; use it to check the setup, not learning.
+
+### 4.2 Full model on eight nodes
+
+Join all eight nodes to a ray head, then launch from the head:
+
+```bash
+MILES_SCRIPT_EXTERNAL_RAY=1 python scripts/run_minimax_m3_1_flash.py \
+    --model-name MiniMax-M3.1-Flash-bf16 --num-nodes 8 --train-offload-dir /scratch/offload
+```
+
+## 5. Recipe Configuration
+
+### 5.1 Parallelism
+
+| `--model-name` | TP | PP | CP | EP | expert-TP | GPUs | rollout engine |
+|---|---|---|---|---|---|---|---|
+| `MiniMax-M3.1-Flash-bf16` | 4 | 4 | 1 | 8 | 1 | 32 (8 × 4) | TP8 / EP8 |
+| `MiniMax-M3.1-Flash-4layer-bf16` | 2 | 1 | 1 | 2 | 1 | 4 | TP4 / EP4 |
+
+TP is capped at 4 by the four KV heads. Context parallelism is not supported, and the sparse attention needs THD packing (`--qkv-format thd`).
+
+### 5.2 Algorithm
+
+GRPO with `--eps-clip 0.2 --eps-clip-high 0.28`, no KL or entropy term, and `--rm-type math` on dapo-math-17k: 16 prompts × 16 samples per step at temperature 0.8. Rollout routing replay (`--use-rollout-routing-replay`) is on, so the trainer routes every token to the experts the engine picked.
+
+### 5.3 Optimizer
+
+Adam at `--lr 3e-6`. The full model's Adam state fits neither the GPUs nor the hosts' memory, so it streams through node-local NVMe:
+
+```bash
+--stream-optimizer-state-to-disk
+--stream-optimizer-state-moment-dtype bf16
+--offload-train-target disk
+--offload-train-disk-dir <train-offload-dir>
+```
+
+### 5.4 Notable quirks
+
+- **The indexer is frozen**: its projections and norms are loaded but not trained, and the trainer recomputes the block selection from its own hidden states rather than replaying the engine's. The two selections differ only at near-ties.
+- **Only the all-sparse layout**: the plugin builds MSA attention on every layer and rejects a checkpoint with dense-attention layers (for example MiniMax-M3) with a `NotImplementedError`.
+- **Frozen towers and `--check-weight-update-equal`**: the trainer never sends the vision weights, so the launcher excludes them from the post-update check with `--check-weight-update-skip-list vision_tower. multi_modal_projector. patch_merge_mlp.`. The check snapshots every weight on the host, which the full model does not fit, so it is off there.
+
+## 6. Pairs Well With
+
+- [Low Precision RL](/advanced/low-precision)

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, name, param)
+    elif "minimax_m3" in model_name:
+        converted_named_tensors = convert_minimax_m3_to_hf(args, name, param)
     elif "mimo" in model_name:
         converted_named_tensors = convert_mimo_to_hf(args, 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/megatron_to_hf/processors/__init__.py

@@ -1,6 +1,7 @@
 from .padding_remover import remove_padding
 from .quantizer_compressed_tensors import quantize_params_compressed_tensors
 from .quantizer_fp8 import quantize_params_fp8
+from .quantizer_minimax_m3_nvfp4 import quantize_params_minimax_m3_nvfp4
 from .quantizer_mxfp8 import quantize_params_mxfp8
 from .quantizer_nvfp4 import quantize_params_nvfp4
 
@@ -17,6 +18,8 @@
 def quantize_params(args, megatron_name, converted_named_params, quantization_config, packed_weight_basenames=None):
     if quantization_config is None:
         return converted_named_params
+    elif quantization_config["quant_method"] == "minimax_m3_nvfp4":
+        return quantize_params_minimax_m3_nvfp4(converted_named_params)
     elif quantization_config["quant_method"] == "fp8":
         return quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config)
     elif quantization_config["quant_method"] == "mxfp8":

miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_minimax_m3_nvfp4.py

@@ -0,0 +1,49 @@
+import re
+
+import torch
+
+from miles.utils.mxfp8 import mxfp8_quantize
+from miles.utils.nvfp4 import nvfp4_global_encode_scale_te, nvfp4_quantize_1d, nvfp4_quantize_1d_pair
+
+# the M3.1 checkpoint: routed experts in compressed-tensors NVFP4, every other projection MXFP8
+_EXPERT = re.compile(r"\.block_sparse_moe\.experts\.\d+\.(w[123])\.weight$")
+_MXFP8 = re.compile(
+    r"\.(self_attn\.(q|k|v|o|index_q|index_k)_proj"
+    r"|mlp\.(gate|up|down)_proj"
+    r"|block_sparse_moe\.shared_experts\.(gate|up|down)_proj)\.weight$"
+)
+
+
+def _nvfp4_entries(name: str, qweight, block_scale, encode_scale):
+    base = name.removesuffix(".weight")
+    return [
+        (f"{base}.weight_packed", qweight),
+        (f"{base}.weight_scale", block_scale),
+        (f"{base}.weight_global_scale", encode_scale.reshape(()).to(torch.float32)),
+    ]
+
+
+def quantize_params_minimax_m3_nvfp4(converted_named_params):
+    experts = {}
+    quantized = []
+    for name, param in converted_named_params:
+        if (m := _EXPERT.search(name)) is not None:
+            experts[m[1]] = (name, param)
+        elif _MXFP8.search(name) is not None:
+            qweight, scale = mxfp8_quantize(param)
+            quantized += [(name, qweight), (name.replace(".weight", ".weight_scale_inv"), scale)]
+        else:
+            quantized.append((name, param))
+
+    if "w2" in experts:
+        name, param = experts.pop("w2")
+        qweight, block_scale, _ = nvfp4_quantize_1d(param)
+        quantized += _nvfp4_entries(name, qweight, block_scale, nvfp4_global_encode_scale_te(param.abs().amax()))
+    if experts:
+        assert set(experts) == {"w1", "w3"}, f"gate and up must convert together, got {sorted(experts)}"
+        (gate_name, gate), (up_name, up) = experts["w1"], experts["w3"]
+        # gate and up share one encode scale, as in the checkpoint
+        encode = nvfp4_global_encode_scale_te(torch.maximum(gate.abs().amax(), up.abs().amax()))
+        (gq, gs, _), (uq, us, _) = nvfp4_quantize_1d_pair(gate, up)
+        quantized += _nvfp4_entries(gate_name, gq, gs, encode) + _nvfp4_entries(up_name, uq, us, encode)
+    return quantized

miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py

@@ -394,7 +394,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.swiglu:
+    if "linear_fc1.weight" in name and (args.swiglu or args.quick_geglu):
         partition_stride = 2
         if partition_dim < 0:
             partition_dim = 0

miles/kernels/attention/qsa/block_sparse_attention.py

@@ -362,9 +362,12 @@ 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
+    span = int((last_tile - first_tile).max().item()) if first_tile.numel() else 0
+    for off in range(span + 1):
+        tile.index_add_(1, torch.minimum(first_tile + off, last_tile), picked)
     return tile > 0
 
 

miles/utils/arguments.py

@@ -4379,12 +4379,16 @@ 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 expert width intermediate_size and the dense width dense_intermediate_size
+    if hasattr(hf_config, "dense_intermediate_size"):
+        ffn_pairs = [("dense_intermediate_size", "ffn_hidden_size"), ("intermediate_size", "moe_ffn_hidden_size")]
+    else:
+        ffn_pairs = [("intermediate_size", "ffn_hidden_size"), ("moe_intermediate_size", "moe_ffn_hidden_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),
+        *((hf_name, mg_name, equal) for hf_name, mg_name in ffn_pairs),
         ("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
 
 
@@ -126,6 +126,9 @@ def load_hf_config(
     """
     register_hf_config_aliases()
     config = AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=trust_remote_code, **autoconfig_kwargs)
+    # some remote config classes (MiniMax-M3) keep the nested text config as a plain dict
+    if isinstance(getattr(config, "text_config", None), dict):
+        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",