[Examples]: MiniMax-M3 DSpark

examples/speculative_decoding/requirements.txt

@@ -1,3 +1,3 @@
 accelerate>=1.12.0
 peft==0.18.1
-transformers>=5.0,<5.4
+transformers>=5.0,<5.13

modelopt/torch/speculative/plugins/hf_streaming_dataset.py

@@ -41,6 +41,7 @@
 
 import base64
 import os
+import re
 import time
 from typing import TypedDict
 
@@ -105,6 +106,22 @@ def _tokenize_with_loss_mask(
     recovery = None
     if answer_only_loss and not getattr(tokenizer, "is_fast", False):
         recovery = get_loss_mask_recovery(tokenizer)
+    if answer_only_loss and recovery is None:
+        # Fail loudly on a template without {% generation %} tags: transformers only
+        # warns and returns an ALL-ZERO assistant mask, which trains every sample at
+        # zero loss with no other symptom. (The regex matches the tag itself, not the
+        # unrelated `add_generation_prompt` identifier most templates contain.)
+        template = getattr(tokenizer, "chat_template", None) or ""
+        if not re.search(r"\{%-?\s*generation\b", template):
+            raise RuntimeError(
+                "answer_only_loss=True needs assistant masks, but the tokenizer's chat "
+                "template has no {% generation %} tags, so apply_chat_template would "
+                "return an all-zero mask and training would silently run at zero loss. "
+                "Pass a tagged template copy via data.chat_template (see "
+                "examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja for a "
+                "worked example), register a loss-mask recovery "
+                "(modelopt.torch.utils.loss_mask), or set answer_only_loss=false."
+            )
     out = tokenizer.apply_chat_template(
         conversations,
         tokenize=True,

modelopt/torch/speculative/plugins/modeling_fakebase.py

@@ -204,7 +204,9 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM
             intermediate_size=getattr(base_cfg, "intermediate_size", None),
             rms_norm_eps=getattr(base_cfg, "rms_norm_eps", 1e-6),
             rope_theta=getattr(base_cfg, "rope_theta", None),
-            final_norm_type=_select_final_norm_type(getattr(base_cfg, "model_type", None)),
+            final_norm_type=_select_final_norm_type(
+                getattr(base_cfg, "model_type", None), base_cfg
+            ),
         )
         model = cls(config)
         # Load lm_head, embed_tokens, and (for known models) the final norm into the model.

modelopt/torch/speculative/plugins/modeling_final_norm.py

@@ -22,6 +22,8 @@
 only when we know which one the model uses.
 """
 
+from collections.abc import Callable
+
 import torch
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
 
@@ -41,11 +43,37 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16):
         self.to(dtype)
 
 
+class _FinalGemmaRMSNorm(torch.nn.Module):
+    """Gemma-style RMSNorm: fp32 normalize, scale by ``(1 + weight)``, multiply-then-cast.
+
+    Mirrors MiniMax M3's ``MiniMaxM3VLRMSNorm`` (``use_gemma_norm=True`` configs): the stored
+    ``weight`` is a zero-centered delta, the effective scale is ``1 + weight``, and the multiply
+    happens in fp32 BEFORE casting back (``(x_hat * (1 + w)).type_as(x)``), unlike LlamaRMSNorm's
+    cast-then-multiply. Reusing ``_FinalRMSNorm`` here would silently drop the ``+1`` and corrupt
+    the reconstructed logits. ``weight`` is loaded from the base checkpoint.
+    """
+
+    def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16):
+        super().__init__()
+        self.eps = eps
+        self.weight = torch.nn.Parameter(torch.zeros(hidden_size, dtype=dtype))
+
+    def forward(self, x):
+        output = x.float()
+        output = output * torch.rsqrt(output.pow(2).mean(-1, keepdim=True) + self.eps)
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
 # Registry of self-implemented final-norm variants. We deliberately reimplement these
 # (rather than importing the base model's actual module) to keep FakeBaseModel lightweight.
 # Only a small, explicit set is supported; add a class here when a new type is needed.
-_FINAL_NORM_CLASSES = {
+_FINAL_NORM_CLASSES: dict[str, Callable[..., torch.nn.Module]] = {
     "rmsnorm": _FinalRMSNorm,
+    "gemma_rmsnorm": _FinalGemmaRMSNorm,
 }
 
 # Base ``model_type`` → final-norm type. ONLY listed models get a norm — applying the wrong or
@@ -62,17 +90,26 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16):
     "deepseek_v3": "rmsnorm",
     "kimi_k2": "rmsnorm",  # Kimi-K2 / K2-Thinking (DeepSeek-V3 arch) report model_type "kimi_k2"
     "kimi_k25": "rmsnorm",  # Kimi-K2.5 / K2.6 / K2.7 all report model_type "kimi_k25"
+    "minimax_m3_vl_text": "rmsnorm",  # overridden to gemma_rmsnorm via use_gemma_norm below
     # gpt_oss intentionally DISABLED: GptOssRMSNorm uses an fp32 weight + multiply-then-cast,
     # unlike _FinalRMSNorm's bf16 weight, so reusing it would silently bias reconstructed logits.
     # Re-enable once a gpt_oss-style class (fp32 weight, multiply-then-cast) is in _FINAL_NORM_CLASSES.
 }
 
 
-def _select_final_norm_type(model_type: str | None) -> str | None:
+def _select_final_norm_type(model_type: str | None, base_cfg=None) -> str | None:
     """Return the final-norm type for a base ``model_type``, or ``None`` if unknown.
 
     ``None`` means we don't know the model's final norm, so FakeBaseModel builds no norm.
+
+    ``base_cfg`` (the resolved text config, optional) takes precedence over the model_type
+    table when it carries an explicit ``use_gemma_norm=True`` flag: only MiniMax M2.x/M3 set
+    it, and their ``model_type`` is unreliable — MiniMax's VL remote code coerces a
+    model_type-less ``text_config`` to Mixtral (whose table entry is plain ``rmsnorm``), so
+    keying off model_type alone would silently apply the wrong norm flavor.
     """
+    if base_cfg is not None and getattr(base_cfg, "use_gemma_norm", False):
+        return "gemma_rmsnorm"
     return _FINAL_NORM_TYPE_BY_MODEL_TYPE.get(model_type or "")
 
 

tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py

@@ -484,3 +484,49 @@ def handler(request: httpx.Request) -> httpx.Response:
     )
     ds[0]
     assert seen_ports == {advertised_port}
+
+
+# ---------------------------------------------------------------------------
+# answer_only_loss template guard
+# ---------------------------------------------------------------------------
+
+
+def _fast_tokenizer_with_template(template: str, seq: int = 8) -> MagicMock:
+    """Fast-tokenizer mock with a given chat template; returns ids + assistant_masks."""
+    tok = MagicMock()
+    tok.is_fast = True
+    tok.chat_template = template
+    tok.apply_chat_template.return_value = {
+        "input_ids": torch.arange(seq, dtype=torch.long).unsqueeze(0),
+        "assistant_masks": torch.ones(1, seq, dtype=torch.long),
+    }
+    return tok
+
+
+_CONV = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]
+
+
+def test_answer_only_loss_rejects_template_without_generation_tags():
+    """A fast tokenizer whose template lacks {% generation %} tags fails loudly.
+
+    Without the guard transformers only warns and returns an ALL-ZERO assistant
+    mask -- every sample then trains at zero loss with no other symptom.
+    """
+    tok = _fast_tokenizer_with_template("{{ bos }}{% if add_generation_prompt %}x{% endif %}")
+    with pytest.raises(RuntimeError, match="generation"):
+        hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=True)
+
+
+def test_answer_only_loss_accepts_tagged_template():
+    """Templates carrying {% generation %} (either whitespace-control form) pass."""
+    for tag in ("{% generation %}", "{%- generation -%}"):
+        tok = _fast_tokenizer_with_template("{{ bos }}" + tag + "{{ c }}")
+        ids, mask = hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=True)
+        assert mask.sum() == ids.shape[-1]
+
+
+def test_full_loss_skips_template_guard():
+    """answer_only_loss=False never consults the template (mask is all ones)."""
+    tok = _fast_tokenizer_with_template("{{ bos }}")
+    ids, mask = hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=False)
+    assert mask.sum() == ids.shape[-1]

tools/launcher/common/eagle3/train_eagle_streaming.sh

@@ -47,6 +47,9 @@
 #   SERVE_CPU_OFFLOAD_GB  GB/GPU offloaded to host RAM (fits big models on too-few GPUs; slower)
 #   SERVE_MAX_MODEL_LEN   cap context length (trims KV/activation)
 #   SERVE_MAX_NUM_SEQS    cap concurrent sequences (trims KV/activation)
+#   SERVE_BLOCK_SIZE      KV-cache block size (e.g. 128 for MiniMax-M3 MSA sparse attention).
+#                         Needs its own knob: multi-token SERVE_EXTRA_ARGS values are mangled
+#                         by nemo_run's unquoted env export, so "--block-size 128" cannot ride it.
 #   SERVE_HOST          single-node: bind/connect host. default 127.0.0.1
 #   SERVE_GPU           single-node: CUDA_VISIBLE_DEVICES for vllm. default "0"
 #   SERVE_TP            tensor-parallel size. default 1 single-node / all serve-node GPUs
@@ -153,6 +156,7 @@ launch_vllm() {
     [ -n "${SERVE_CPU_OFFLOAD_GB:-}" ] && opt_args+=(--cpu-offload-gb "$SERVE_CPU_OFFLOAD_GB")
     [ -n "${SERVE_MAX_MODEL_LEN:-}" ]  && opt_args+=(--max-model-len "$SERVE_MAX_MODEL_LEN")
     [ -n "${SERVE_MAX_NUM_SEQS:-}" ]   && opt_args+=(--max-num-seqs "$SERVE_MAX_NUM_SEQS")
+    [ -n "${SERVE_BLOCK_SIZE:-}" ]     && opt_args+=(--block-size "$SERVE_BLOCK_SIZE")
     # --no-enable-chunked-prefill / --no-enable-prefix-caching: connector captures hidden states during prefill; both skip recomputing cached/partial prefixes, yielding short/empty hidden_states. Required.
     # --no-enable-flashinfer-autotune: on NVFP4 MoE the autotuner re-tunes on the first serving step and stalls a worker past vLLM's execute-model timeout, killing EngineCore.
     # Hidden states move serve -> trainer over NIXL RDMA (no disk round-trip): one

tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml

@@ -0,0 +1,134 @@
+# DSpark streaming speculative-decoding training for MiniMax-M3 (multi-node).
+# DSpark = the DFlash backbone + a lightweight Markov head + a confidence head,
+# generating a causal block semi-autoregressively; see dspark.yaml for the head
+# and loss config. Runs the shared streaming pipeline
+# (common/eagle3/train_eagle_streaming.sh) with the M3-specific base, draft dims,
+# mask token and chat template; trained from scratch. A starting point for
+# reproduction — tune node counts, batch, steps and serve limits for your cluster.
+#
+# MiniMax-M3 specifics this yaml encodes (each was a silent failure mode):
+#   * SERVE_BLOCK_SIZE=128: M3's MSA sparse attention (sparse_block_size=128)
+#     requires KV block 128 or vLLM dies at engine init ("No common block size").
+#   * data.chat_template: M3 ships a FAST tokenizer whose template has no
+#     {% generation %} tags -> assistant_masks come back ALL-ZERO and
+#     answer_only_loss training silently runs at zero loss. The tagged template
+#     copy next to this yaml wraps the assistant turn (think prefix + content +
+#     tool calls + eos) in {% generation %} tags.
+#   * The draft does NOT inherit base GQA/FFN dims (set explicitly below), and
+#     M3's base rope_theta (5e6) is pinned onto the draft.
+#   * EAGLE_CAPTURE_IDS = draft default target_layer_ids+1 (6 aux) + final (60).
+#   * M3 is a VLM wrapper (text_config nested): the base's gemma-style final
+#     norm is selected via the config's use_gemma_norm flag (see
+#     modeling_final_norm.py) — its text_config coerces to a mixtral model_type,
+#     so the model_type table alone would pick the wrong norm.
+#
+# Run ON the cluster login node (paramiko can't reach it through the login proxy):
+#   export SLURM_HOST=localhost SLURM_ACCOUNT=<your_account> \
+#          SLURM_PARTITION=<multi_node_partition> \
+#          SLURM_HF_LOCAL=<hf_models_dir> \
+#          SLURM_JOB_DIR=<experiments_dir> \
+#          NEMORUN_HOME=$PWD
+#   uv run launch.py --yaml examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml \
+#          identity=$HOME/.ssh/id_ecdsa detach=True --yes
+#
+# The export lands in /scratchspace/export.
+
+job_name: MiniMax-M3_DSpark_streaming_multi_node
+pipeline:
+  allow_to_fail: false
+  skip: false
+  note:
+
+  global_vars:
+    hf_model: /hf-local/MiniMaxAI/MiniMax-M3
+
+  # Build /scratchspace/data/train.jsonl. Point data.data_path at the full
+  # Spec-Decoding-Dataset-v2 corpus to reproduce; eagle_utils also accepts a
+  # directory of *.jsonl shards directly.
+  task_0:
+    script: common/eagle3/make_dataset.sh
+    args:
+      - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml
+      - --full-conversations
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 8
+      container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10
+
+  task_1:
+    script: common/eagle3/train_eagle_streaming.sh
+    args:
+      - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml
+      - model.model_name_or_path=<<global_vars.hf_model>>
+      - model.use_fake_base_for_offline=true
+      - model.trust_remote_code=true
+      - data.mode=streaming
+      - data.data_path=/scratchspace/data/train.jsonl
+      # M3's own template has no {% generation %} tags; without this tagged copy
+      # answer_only_loss trains on an all-zero mask (see header).
+      - data.chat_template=examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja
+      - training.output_dir=/scratchspace/dspark
+      - training.training_seq_len=4096
+      - training.disable_tqdm=true
+      - training.ar_validate_steps=500000
+      - training.num_train_epochs=1
+      - training.per_device_train_batch_size=4
+      - training.gradient_accumulation_steps=1
+      - training.save_steps=1000
+      - training.logging_steps=20
+      - training.learning_rate=1.0e-4
+      - training.warmup_steps=2000
+      - training.answer_only_loss=true
+      # The vLLM serve container has no tensorboard -> trainer init crash.
+      - training.report_to=none
+      # The DSpark draft does NOT inherit the base GQA/FFN dims, so set them
+      # explicitly to match the M3 backbone (else a silently wrong-shape draft).
+      # intermediate_size matches M3's dense FFN (dense_intermediate_size).
+      - dflash.dflash_architecture_config.num_hidden_layers=6
+      - dflash.dflash_architecture_config.num_key_value_heads=8
+      - dflash.dflash_architecture_config.intermediate_size=12288
+      # Pin the base's rope_theta onto the draft (M3 uses 5e6, not the Qwen3
+      # default 1e6; a mismatch trains rope into the weights and caps AL).
+      - dflash.dflash_architecture_config.rope_theta=5000000
+      # Semi-AR generation block (dspark.yaml ships 16; Kimi/M3 runs use 8).
+      - dflash.dflash_block_size=8
+      # M3 has no mask token; vocab 200064, added tokens end at 200060 -> 200063 free.
+      - dflash.dflash_mask_token_id=200063
+    environment:
+      - HF_MODEL_CKPT: <<global_vars.hf_model>>
+      # 6 aux capture ids = the draft's default target_layer_ids+1, plus the true
+      # final hidden (60). Requires the aux-capture fix vllm#46788 (in-tree in
+      # recent nightlies); mismatched ids silently skew train vs inference.
+      - EAGLE_CAPTURE_IDS: "[2,13,24,36,47,58,60]"
+      - SERVE_NODES: "4"
+      - SERVE_TP: "8"
+      - STREAMING_NUM_WORKERS: "4"
+      # M3's custom-modeling base needs trust_remote_code at export and serve.
+      - EXPORT_EXTRA_ARGS: "--trust_remote_code"
+      - SERVE_EXTRA_ARGS: "--trust-remote-code"
+      # REQUIRED for M3: MSA sparse attention needs KV block 128 (dedicated knob —
+      # multi-token SERVE_EXTRA_ARGS values are mangled by nemo_run's unquoted
+      # env export, so "--block-size 128" cannot ride it).
+      - SERVE_BLOCK_SIZE: "128"
+      - SERVE_MAX_MODEL_LEN: "4160"
+      - SERVE_MAX_NUM_SEQS: "32"
+      - SERVE_GPU_MEM_UTIL: "0.9"
+      - SERVE_READY_TIMEOUT: "3600"
+      - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200"
+      - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200"
+      # RDMA transport is UCX (InfiniBand) by default. On AWS EFA, uncomment —
+      # and note UCX SEGFAULTS at agent init on EFA nodes (it detects the EFA
+      # devices), so LIBFABRIC is required there even for single-node runs:
+      # - NIXL_BACKENDS: "LIBFABRIC"
+      # - FI_PROVIDER: "efa"
+      # - NCCL_IB_DISABLE: "1"
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 6
+      ntasks_per_node: 1
+      gpus_per_node: 8
+      # vLLM x86_64 build with native MiniMax-M3 support (vllm/models/minimax_m3)
+      # and the aux-capture fix (vllm#46788).
+      container: <vllm-image-with-native-minimax-m3>