DFlash for MiniMax-M3 (WIP): synthesis thinking-mode mix by yeyu-nvidia · Pull Request #1749 · NVIDIA/Model-Optimizer

yeyu-nvidia · 2026-06-16T18:04:16Z

Tracking PR for DFlash on MiniMax-M3, reusing the M2.7 DFlash work merged in #1621. Opening early with the first piece — the data-synthesis tooling — more to follow (offline recipe folder, mask-token / YaRN config for M3).

This PR so far

Per-conversation thinking-mode mix in data synthesis (server_generate.py + distributed_generate/worker.sh):

M3 has three thinking modes (enabled / disabled / adaptive). To train a DFlash draft that generalizes across all of them, synthetic conversations should span the mix.
server_generate.py gains --thinking-modes enabled,disabled,adaptive: conversation i uses modes[i % len(modes)] (even split), passed via chat_template_kwargs, and the chosen mode is recorded on each output record. Default empty → unchanged for models without thinking modes.
We intentionally do not enable the server's --reasoning-parser for synthesis, so the full <mm:think>+answer lands in content (the draft must learn to draft the entire generated sequence).
worker.sh threads THINKING_MODES through and adds VLLM_SERVE_EXTRA_ARGS / SGLANG_SERVE_EXTRA_ARGS passthroughs for model-specific serve flags.

M3 serving notes (validated on H100)

MiniMaxAI/MiniMax-M3-MXFP8 serves single-node TP8 on H100 with image vllm/vllm-openai:minimax-m3, --block-size 128 (mandatory for MSA sparse attention) and --language-model-only (text-only). KV cache must stay bf16 — M3's MSA fused kernel rejects fp8 KV.

To come

examples/.../MiniMax/MiniMax-M3-DFlash/ offline recipe (mirroring M2.7, offline path).
M3 dflash_mask_token_id=200061 (M3's 200054 is now a real special token) + YaRN export config.

🤖 Generated with Claude Code

…niMax-M3) Adds a --thinking-modes cycle to server_generate.py so synthetic conversations are generated across a mix of thinking modes — e.g. MiniMax-M3's enabled/disabled/adaptive, passed via chat_template_kwargs — so a DFlash/EAGLE draft trained on the data generalizes across modes. Conversation i uses modes[i % len(modes)] for an even split; the mode is recorded on each output record. Empty (default) sends no thinking_mode, unchanged for models without it. distributed_generate/worker.sh: pass THINKING_MODES through to server_generate.py, and add VLLM_SERVE_EXTRA_ARGS / SGLANG_SERVE_EXTRA_ARGS passthroughs for model-specific serve flags (M3 needs --block-size 128 for MSA sparse attention and --language-model-only for text-only synthesis; KV cache stays bf16 — M3's MSA fused kernel rejects fp8 KV). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

copy-pr-bot · 2026-06-16T18:04:20Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

coderabbitai · 2026-06-16T18:04:51Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e0003d34-90b0-4053-b8d3-2cca48635351

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch yeyu/dflash-minimax-m3

_{Comment @coderabbitai help to get the list of available commands.}

codecov · 2026-06-16T18:20:43Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.96%. Comparing base (e6790ef) to head (cc91cce).
⚠️ Report is 64 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1749      +/-   ##
==========================================
- Coverage   77.12%   73.96%   -3.16%     
==========================================
  Files         511      516       +5     
  Lines       56273    62769    +6496     
==========================================
+ Hits        43399    46430    +3031     
- Misses      12874    16339    +3465

Flag	Coverage Δ
unit	`54.91% <100.00%> (+0.52%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…tput While wiring MiniMax-M3 synthesis (prompt set Speculative-Decoding-Dataset-v2 is OAI-format): - Accept both 'conversations' (ShareGPT) and 'messages' (OAI) prompt datasets on input (previously KeyError'd on 'messages'). - Add --output-format {oai,sharegpt} (default oai): emit the OpenAI standard {'messages': [{role, content}, ...]} instead of the legacy {'conversations': [...]}. Pass --output-format sharegpt for the old key. Validated end-to-end on MiniMax-M3-MXFP8 (single-node TP8 H100): the 3-way thinking-mode mix renders correctly (disabled -> direct answer; adaptive -> <mm:think> reasoning captured in content since no reasoning-parser), OAI in/out flows into the vLLM hidden-state dump. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…747) Validated extract_hidden_states on MiniMax-M3-MXFP8 (single-node TP8 H100). Required M3 enablement in compute_hidden_states_vllm.py: - --block-size (M3's MSA sparse attention mandates 128; default None elsewhere). - --enforce-eager: M3's MSA Triton kernel (_gqa_sparse_fwd_kernel) JIT-recompiles per input shape; under cudagraph capture a recompile blows the executor RPC timeout and hangs the engine (sample_tokens timeout). Eager mode + a long VLLM_RPC_TIMEOUT fixes it. - --language-model-only: skip the vision encoder for text-only dumps (M3 is VL). - Read num_hidden_layers from text_config/llm_config for wrapped VL configs (MiniMaxM3VLConfig nests it; previously raised 'no num_hidden_layers attribute'). Output verified: per-conv .pt with input_ids / hidden_states (T,6144) / aux_hidden_states / loss_mask (length-matched) / conversation_id. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…e (OMNIML-4747) 2-step offline DFlash recipe for MiniMax-M3 (427B VL-MoE), mirroring MiniMax-M2.7-DFlash: - hf_offline_dflash.yaml: dump (vLLM extract_hidden_states, MXFP8 single-node TP8) + train (FakeBaseModel on bf16). M3-specific: --block-size 128 (MSA), --language-model-only, --enforce-eager + VLLM_RPC_TIMEOUT=1800000 (avoid MSA Triton-kernel JIT RPC-timeout hang), seq-len 8192 end-to-end, mask token 200061 (200054 is a real special token in M3), OVERRIDE_TRANSFORMERS 4.52.4, export-YaRN original_max_position 8192 / factor 24 (tunable; 128 for full 1M). - chat_template_train.jinja: M3 chat template with {% generation %} wrapping the assistant turn (think + content + tool_calls) for answer_only_loss; header + eos sit outside the span, matching the M2.7 convention. Thinking-mode handling preserved verbatim. Validated: generation spans cover exactly the assistant outputs across multi-turn + no-think turns. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…ctor 128) original_max_position_embeddings=8192 (training seq-len) x factor 128 = 1048576 = M3's full 1M context. Export-time tunable (factor 24 -> 196608 for the M2.7-equivalent target). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…taset Synth campaign complete: 2,761,690 clean records / 16,472 shards at /hf-local/modelopt/MiniMax-M3-synthetic-data-clean (empties + finished-markers dropped, provenance joined from source prompts, even 3-way thinking-mode mix). Replaces the pre-campaign placeholder path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…NIML-4747) Dry run surfaced that M3's MSA Triton kernel (_gqa_sparse_fwd_kernel) JIT-recompiles per prefill shape (~110s vs ~10s cached). On real 8192-token data hundreds of distinct lengths each recompile, blowing VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (default 300) -> sample_tokens RPC timeout -> EngineDeadError. (M2.7 uses dense attention + stock vLLM, so it never hits this.) The earlier VLLM_RPC_TIMEOUT 'fix' was a no-op: that env var is unrecognized in this vLLM build and silently ignored. compute_hidden_states_vllm.py: - --length-buckets: right-pad each prompt to the smallest bucket >= its length so the model sees only a handful of prefill shapes. Causal attention leaves the real prefix's hidden states identical to the unpadded forward; pad positions are sliced off on save via real_len. - --warmup: one throwaway prefill per bucket to JIT-compile shapes up front (and populate a persistent TRITON_CACHE_DIR shared across tasks). M3 recipe: add --length-buckets 1024,2048,4096,8192 + --warmup; set VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 (was bogus VLLM_RPC_TIMEOUT) + a persistent TRITON_CACHE_DIR so all parallel dump tasks reuse the warmed kernels. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

A prompt (or a length bucket) exactly equal to max_seq_len + the dump's dummy 1-token generation = max_seq_len+1, which vLLM rejects against max_model_len=max_seq_len ('decoder prompt plus requested output tokens is longer than the maximum model length'). Set max_model_len = max_seq_len + 1. Surfaced by the M3 dry run when --length-buckets padded prompts to exactly the 8192 top bucket. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

… TP deadlock (OMNIML-4747) Dry run with bucketing+warmup got further (MSA attention kernel handled, first real prompt fast) but then deadlocked on a SECOND kernel: M3's MoE expert-routing _topk_index_kernel uses a raw @triton.autotune (5 configs). Its ~150s per-shape benchmark runs UNSYNCHRONIZED across the 8 TP ranks; while some ranks benchmark, others block at the hidden-states connector's shared-memory collective -> permanent desync (22min 'no broadcast block') -> idle-GPU reaper kills it. (Serving never hits this: it's decode-dominated with warm kernels and no connector collective. Confirmed via the reaper alert showing 4/8 GPUs idle = the rank split.) --disable-triton-autotune pins every @triton.autotune to a single (last/most-conservative) config so no benchmark runs -> no cross-rank race. Applied in-process (forked workers) and via a sitecustomize injected on PYTHONPATH (spawned workers), since the autotune executes in vLLM's worker subprocesses. Autotune configs are functionally equivalent (perf-only), so pinning is correct. Wired into the M3 recipe alongside the existing bucketing/warmup/timeout/cache fixes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…sync (OMNIML-4747) Dry run confirmed --disable-triton-autotune works (0 autotune benchmarks vs the prior 153s ones), but the dump still deadlocked: with autotune gone, plain per-shape kernel COMPILATION (slot_mapping / index_block_score / topk / gqa_sparse) still runs per-rank and desyncs the TP collective at the first real prompt whose batch shape the warm-up didn't cover. Warm-up synchronizes (all ranks compile the same batch together), so the fix is to make the real loop hit only warmed shapes. Bucketing fixes the sequence-length dimension; max_num_seqs=1 fixes the batch dimension, so a batch=1 bucketed warm-up covers every (1, bucket) shape the real loop sees -> no compilation during the timed dump -> no desync. Wired --max-num-seqs 1 into the M3 recipe. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…rward (OMNIML-4747) DFlash offline training crashed in the draft forward with 'mat1 and mat2 must have the same dtype: BFloat16 != float': the dumped hidden states are bf16, but with bf16=false the draft (incl. lm_head + input projections) is fp32. The M3 dump saves no logits, so the self-logit-distillation path recomputes them via lm_head(hidden) -> bf16 input x fp32 weight -> error; target_hidden would hit the same downstream. Cast the offline target_hidden and the lm_head input up to the model dtype (lm_head.weight.dtype); bf16 -> fp32 is lossless. Only the offline path is affected (the online path's hidden states already come from the fp32 forward). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

github-actions · 2026-06-29T22:07:16Z

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1749/
Built to branch `gh-pages` at 2026-07-07 20:52 UTC. Preview will be ready when the GitHub Pages deployment is complete.

…(OMNIML-4747) Root-cause fix for the offline-training 'mat1 and mat2 must have the same dtype: BFloat16 != float' crash, replacing the earlier symptom cast (now reverted). FakeBaseModel.from_source resolved its dtype from getattr(base_cfg, 'dtype', bf16). For a VLM, base_cfg is the nested text config, which often does NOT set torch_dtype — only the top-level config does (MiniMax-M3: top-level torch_dtype=bfloat16, text_config none). base_cfg.dtype then falls back to PyTorch-default fp32, so lm_head/embed_tokens load fp32 while the dumped hidden states are bf16 -> dtype mismatch in the draft forward. (MiniMax-M2.7 is non-VL, so base_cfg is the top-level config carrying torch_dtype=bfloat16 -> bf16 -> never hit this.) Resolve torch_dtype preferring the nested config, then the top-level, then bf16 (mapping a string to torch.dtype), so the fake base loads in the checkpoint dtype and matches the dump. Reverts the hf_dflash.py offline-hidden-states cast, which masked this by upcasting to fp32. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…ted e2e) Validated the full M3 offline path end-to-end (dump -> train -> export): training must use transformers 4.57.1, not M3's 4.52.4. 4.52.4's Trainer imports apex.amp (absent in the training container) and crashes at import; 4.57.1 guards it and loads M3's config/tokenizer fine. With this + the FakeBase VL-dtype fix, training runs and exports the Qwen3 draft with the full-1M YaRN rope_scaling (factor 128) in bf16. The dump step is unaffected (vllm image). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

Asserts FakeBaseModel.from_source resolves the base dtype from the top-level config when the VLM's nested text_config carries no torch_dtype (the MiniMax-M3 case): lm_head / embed_tokens must load bf16, not fp32. Without the fix they load fp32 and mismatch the bf16 offline hidden states in the DFlash draft forward. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…-4747) Full-scale dump filled /dev/shm: the script called llm.generate(ALL prompts) then saved afterward, so every conv's hidden states (~190MB) accumulated in the KV connector's staging dir (DFLASH_HS_STAGING_DIR, default /dev/shm RAM, ~1TB/node) -> 'No space left on device' at ~21%. (The fix stack itself held at scale: all 8 tasks hit 100% at ~2.4 it/s, no deadlock.) Process in chunks of --save-chunk-size (default 256): generate a chunk, save + cleanup each output, then the next chunk — so staging holds at most one chunk (~48GB) at a time. Bonus: saves are now incremental, so a task killed at the 4h limit keeps its completed chunks and keep_conversation resume continues from there. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…bs=1) Offline training is lightweight — FakeBaseModel loads only lm_head + embed_tokens (not the 427B base), the draft is a small 5-layer Qwen3, and dflash_num_anchors=512 caps the vocab-logit memory to 512 positions rather than all 8192. So the 8-node x bs=1 config (inherited from M2.7) over-provisions GPUs. Use 2 nodes x 8 GPU x bs=4 = the same effective batch (64) at 1/4 the GPUs; bs is tunable upward given the anchor-bounded logits. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…vs (OMNIML-4747) The skip-existing filter (keep_conversation, which checks whether each conv's .pt already exists) was being served from datasets' cache: filter/map results are fingerprinted on the function + dataset, NOT on external disk state. With a persistent HF cache across requeues, the cached 'keep everything' from an early run (few/no .pt) was reused, so resume reported 'Removed 0' and re-dumped/overwrote already-done conversations (observed: 62k .pt rewritten in 2h, total flat). Pass load_from_cache_file=False so the filter re-checks the disk every run. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

Full-scale training OOM'd at bs=4 on 80GB (73.7GB in use, +6.1GB requested): M3's offline inputs are activation-heavy (6144 hidden x 5 aux layers x 8192 seq) even though the FakeBase + 5-layer draft are small and dflash_num_anchors=512 bounds the logits. bs=2 x grad_accum=2 keeps effective batch 64 (16 ranks / 2 nodes) and fits. Also set PYTORCH_CUDA_ALLOC_CONF= expandable_segments:True to reduce fragmentation. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

…s unstable) The 2-node training repeatedly crashed on cw-dfw with C10d rendezvous failures (torch.distributed DistNetworkError ~every 20-60min); slurm --requeue doesn't recover it since it's an app-level NCCL/rendezvous crash, not a node failure. Offline training is light enough for a single node anyway (FakeBaseModel + 5-layer draft), which eliminates cross-node rendezvous entirely. 1 node x 8 GPU x bs=2 x grad_accum=4 = effective batch 64 (same as before), stable + hands-off with per-window self-requeue (main.py resumes from checkpoints). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ye Yu <yeyu@nvidia.com>

yeyu-nvidia and others added 10 commits June 16, 2026 11:51

yeyu-nvidia and others added 8 commits June 29, 2026 15:12

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

DFlash for MiniMax-M3 (WIP): synthesis thinking-mode mix - #1749

DFlash for MiniMax-M3 (WIP): synthesis thinking-mode mix#1749
yeyu-nvidia wants to merge 19 commits into
mainfrom
yeyu/dflash-minimax-m3

yeyu-nvidia commented Jun 16, 2026

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

coderabbitai Bot commented Jun 16, 2026 •

edited

Loading

Review skipped

Uh oh!

codecov Bot commented Jun 16, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 29, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-07-07 20:52 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

yeyu-nvidia commented Jun 16, 2026

This PR so far

M3 serving notes (validated on H100)

To come

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

coderabbitai Bot commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

codecov Bot commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Built to branch gh-pages at 2026-07-07 20:52 UTC. Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

coderabbitai Bot commented Jun 16, 2026 •

edited

Loading

codecov Bot commented Jun 16, 2026 •

edited

Loading

github-actions Bot commented Jun 29, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-07-07 20:52 UTC.
Preview will be ready when the GitHub Pages deployment is complete.