[Examples]: MiniMax-M3 DSpark#1965
Conversation
📝 WalkthroughWalkthroughThe changes add Gemma final-norm support, validate answer-only loss masking prerequisites, add a MiniMax-M3 chat template, and provide a multi-node DSpark streaming training configuration with configurable vLLM block size. ChangesSpeculative streaming support
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant DSparkPipeline
participant MakeDataset
participant TrainEagleStreaming
participant VLLMServe
DSparkPipeline->>MakeDataset: create train.jsonl
DSparkPipeline->>TrainEagleStreaming: launch multi-node training
TrainEagleStreaming->>VLLMServe: start vLLM with block size 128
VLLMServe-->>TrainEagleStreaming: serve speculative decoding
Caution Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional.
❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/speculative/plugins/hf_streaming_dataset.py`:
- Around line 116-124: Update the error message in the assistant-mask validation
guard to reference the actual template location,
tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja,
while leaving the surrounding guidance unchanged.
- Around line 106-124: The guard around answer_only_loss must also raise when
the tokenizer is non-fast and no loss-mask recovery is available, regardless of
whether the chat template contains {% generation %} tags. Update the condition
using recovery, tokenizer.is_fast, and the existing template check, and adjust
the RuntimeError message to identify the unsupported slow-tokenizer path while
preserving the current guidance.
In
`@tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml`:
- Around line 99-100: Add the required MLM_MODEL_CFG environment variable
alongside HF_MODEL_CKPT in the MiniMax-M3 configuration, setting it to the
model’s Hugging Face repository ID and preserving the existing HF_MODEL_CKPT
reference.
In
`@tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja`:
- Line 1: Add the required provenance and licensing notices at the top of the
template before the special token variables: include the MiniMax source
reference, the upstream copyright and license notice, and the NVIDIA Apache 2.0
header. Also update the repository LICENSE to include the MiniMax Community
License.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 18176f48-ae3d-49dc-b512-26580fa2bc17
📒 Files selected for processing (8)
examples/speculative_decoding/requirements.txtmodelopt/torch/speculative/plugins/hf_streaming_dataset.pymodelopt/torch/speculative/plugins/modeling_fakebase.pymodelopt/torch/speculative/plugins/modeling_final_norm.pytests/unit/torch/speculative/plugins/test_hf_streaming_dataset.pytools/launcher/common/eagle3/train_eagle_streaming.shtools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yamltools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja
| @@ -0,0 +1,250 @@ | |||
| {# ---------- special token variables ---------- #} | |||
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target file and related example templates for provenance/licensing patterns.
git ls-files 'tools/launcher/examples/MiniMaxAI/MiniMax-M3/*' | sed 's#^`#FILE`: #'
echo
wc -l tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja
echo
cat -n tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja | sed -n '1,80p'
echo
rg -n --hidden --no-messages 'SPDX-License-Identifier|copyright|license|provenance|source' tools/launcher/examples/MiniMaxAI/MiniMax-M3 tools/launcher/examples/MiniMaxAI -g '!**/*.pyc'Repository: NVIDIA/Model-Optimizer
Length of output: 5262
🌐 Web query:
MiniMaxAI MiniMax-M3 chat_template.jinja license source commit minimax-community
💡 Result:
The MiniMax-M3 model is released under the MiniMax Community License [1][2][3], which is not an unrestricted permissive open-source license [1]. The chat_template.jinja file for MiniMax-M3 is hosted in the official MiniMaxAI/MiniMax-M3 repository on Hugging Face [4][5]. As of July 12, 2026, the model repository indicates that the license governing its use is the "minimax-community" license [2][3]. This license typically includes specific terms such as attribution requirements and restrictions on certain commercial uses [1]. You can view the specific content of the chat template and track its history directly via the Hugging Face repository [4][6]. The repository for MiniMax-M3 is maintained under the MiniMaxAI organization on Hugging Face [4][3], and related code can also be found in the MiniMax-AI GitHub organization [7]. Always refer to the official LICENSE file [8] located in the repository for the most accurate and legally binding terms regarding the model's usage.
Citations:
- 1: https://minimax-ai.chat/models/minimax-m3/
- 2: https://huggingface.co/MiniMaxAI/MiniMax-M3/raw/ca22f9aa9ab883c93f0154f6f94e0f585ee60e21/README.md
- 3: https://huggingface.co/MiniMaxAI/MiniMax-M3
- 4: https://huggingface.co/MiniMaxAI/MiniMax-M3/blob/main/chat_template.jinja
- 5: https://huggingface.co/MiniMaxAI/MiniMax-M3/tree/main
- 6: https://huggingface.co/MiniMaxAI/MiniMax-M3/blob/2e2cee36b34856cb2a446d2f31677e51c97dbcbe/chat_template.jinja
- 7: https://github.com/MiniMax-AI/MiniMax-M3
- 8: https://huggingface.co/MiniMaxAI/MiniMax-M3/blob/main/LICENSE
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether the repo already has license notices or hook exclusions for this file/path.
rg -n --no-messages 'MiniMax|minimax-community|third-party|Third-Party Software Notices|insert-license|m3_chat_template_generation\.jinja' LICENSE .pre-commit-config.yaml tools/launcher/examples/MiniMaxAI -g '!**/*.pyc'
echo
echo "LICENSE lines:"
wc -l LICENSE
echo
sed -n '1,260p' LICENSE | cat -n | sed -n '1,260p'
echo
echo ".pre-commit-config.yaml lines:"
wc -l .pre-commit-config.yaml
echo
sed -n '1,260p' .pre-commit-config.yaml | cat -n | sed -n '1,260p'Repository: NVIDIA/Model-Optimizer
Length of output: 26892
Add the MiniMax provenance and license notices here.
This template needs a source reference, the upstream copyright/license notice, and the NVIDIA Apache 2.0 header, plus an update to LICENSE for the MiniMax Community License.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja`
at line 1, Add the required provenance and licensing notices at the top of the
template before the special token variables: include the MiniMax source
reference, the upstream copyright and license notice, and the NVIDIA Apache 2.0
header. Also update the repository LICENSE to include the MiniMax Community
License.
Source: Coding guidelines
h-guo18
changed the title
468013d to
b31352a
Compare
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| "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 |
There was a problem hiding this comment.
[IMPORTANT ModeState] This new minimax_m3_vl_text → "rmsnorm" table entry is a silent-corruption fallback that undermines the fail-safe this file is built around.
The selection order in _select_final_norm_type is: use_gemma_norm first, then this table. So the entry is only ever consulted when use_gemma_norm is absent/falsy but model_type == "minimax_m3_vl_text". In that case it returns plain "rmsnorm", which drops the (1 + weight) scale — exactly the silent bias into the reconstructed logits / distillation target that the rest of this module (and the gpt_oss-disabled comment above) deliberately avoids.
The file's stated contract is "unlisted → no norm → hard error" precisely because a wrong norm is worse than no norm. Mapping M3 to rmsnorm as a fallback trades a loud RuntimeError (from _maybe_apply_base_final_norm) for a silent +1 drop.
Since every M3 config that reaches here is gemma-normed, either:
- Remove this entry so a
use_gemma_norm-less M3 falls through toNone→ hard error, or - Map it to
"gemma_rmsnorm"so the correct flavor is applied even if the flag is missing.
Both are safer than the current "rmsnorm".
There was a problem hiding this comment.
Claude Review Summary
Findings — CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 0
Reviewed all 8 changed files (5 source/example, 1 test, 1 shell, 1 requirements). The four fixes are well-motivated and the per-fix rationale in the PR body / inline comments is unusually thorough. The Gemma-norm math (fp32 normalize -> *(1+weight) -> type_as) is correct and matches MiniMaxM3VLRMSNorm; the constructor signature matches _FinalRMSNorm so the registry dispatch is consistent; the transformers pin (<5.13) is consistent with pyproject.toml; and the SERVE_BLOCK_SIZE knob is a clean, opt-in addition.
Most impactful finding
[IMPORTANT ModeState] modeling_final_norm.py:93 — the new minimax_m3_vl_text -> "rmsnorm" table entry is a silent-corruption fallback. Because use_gemma_norm is checked first, this entry only fires when the gemma flag is absent — and then it applies plain RMSNorm, dropping the (1 + weight) scale and silently biasing the reconstructed distillation target. That is the exact failure mode this file's "unlisted -> hard error" design (and the disabled gpt_oss entry) exists to prevent. Map it to gemma_rmsnorm or remove the entry entirely so a flag-less M3 hard-errors rather than mis-norms.
Assessment
Low-to-moderate risk. The change is backward compatible (norm selection only affects use_gemma_norm=True models, the mask guard converts a silent zero-loss run into a loud error, and the block-size/pin changes are additive). The one IMPORTANT item is a latent fail-safe regression rather than a live bug, since M3 configs set use_gemma_norm in practice — but it is worth closing since defeating the safety net is precisely what the surrounding code warns against.
Note: CodeRabbit's pre-merge gate flags the hardcoded trust_remote_code=true in the YAML (SECURITY.md exception needed) and a slow-tokenizer edge in the mask guard — deferring those to CodeRabbit rather than duplicating.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1965 +/- ##
==========================================
+ Coverage 75.79% 77.84% +2.05%
==========================================
Files 518 518
Lines 58658 58678 +20
==========================================
+ Hits 44459 45680 +1221
+ Misses 14199 12998 -1201
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja (2)
209-212: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winValidate tool messages against the immediately preceding tool-call turn.
last_tool_call.nameremains set after a tool response and is not cleared by a subsequent user message, allowing a later standalonetoolmessage to pass this guard incorrectly. Validate the adjacent message/contiguous tool block or clear the pending state whenever a non-tool turn intervenes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja` around lines 209 - 212, Update the tool-message branch in the chat template to validate that the tool message belongs to the immediately preceding assistant tool-call turn or contiguous tool block, rather than relying on stale last_tool_call.name state. Clear or invalidate the pending tool-call state whenever a non-tool message intervenes, while preserving valid consecutive tool responses.
159-166: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winTreat an empty
reasoning_contentas absent.When
reasoning_contentis"",is stringsucceeds and the fallback parser is skipped. Ifmessage.contentstill contains<mm:think>...</mm:think>, the template emits an extra closing tag and leaves reasoning in visible content, corrupting answer-only training masks. Fall back to parsing when the field is empty.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja` around lines 159 - 166, Update the reasoning_content condition in the message template so an empty string is treated as absent and enters the existing think-token parsing fallback. Preserve direct assignment for non-empty string reasoning_content and ensure content containing think_end_token is split into reasoning and visible content as before.modelopt/torch/speculative/plugins/hf_streaming_dataset.py (1)
109-135: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winNormalize dict-valued chat templates before regex validation.
tokenizer.chat_templatecan be a dict for multi-template models, andre.search(...)will crash on it beforeapply_chat_templatecan select the default template. Normalize the template first and add a regression test for{"default": ...}.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/speculative/plugins/hf_streaming_dataset.py` around lines 109 - 135, The template validation in the answer_only_loss guard must handle dict-valued tokenizer.chat_template values before calling re.search. Normalize the chat template using the default entry selected by apply_chat_template, then validate that string for {% generation %} tags while preserving existing behavior for string templates; add a regression test covering a {"default": ...} template.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@modelopt/torch/speculative/plugins/hf_streaming_dataset.py`:
- Around line 109-135: The template validation in the answer_only_loss guard
must handle dict-valued tokenizer.chat_template values before calling re.search.
Normalize the chat template using the default entry selected by
apply_chat_template, then validate that string for {% generation %} tags while
preserving existing behavior for string templates; add a regression test
covering a {"default": ...} template.
In
`@tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja`:
- Around line 209-212: Update the tool-message branch in the chat template to
validate that the tool message belongs to the immediately preceding assistant
tool-call turn or contiguous tool block, rather than relying on stale
last_tool_call.name state. Clear or invalidate the pending tool-call state
whenever a non-tool message intervenes, while preserving valid consecutive tool
responses.
- Around line 159-166: Update the reasoning_content condition in the message
template so an empty string is treated as absent and enters the existing
think-token parsing fallback. Preserve direct assignment for non-empty string
reasoning_content and ensure content containing think_end_token is split into
reasoning and visible content as before.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ed5cf9b3-3c40-4de5-9cc1-3931fd61a122
📒 Files selected for processing (4)
modelopt/torch/speculative/plugins/hf_streaming_dataset.pymodelopt/torch/speculative/plugins/modeling_final_norm.pytests/unit/torch/speculative/plugins/test_hf_streaming_dataset.pytools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/speculative/plugins/modeling_final_norm.py
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
b14501a to
a9303ff
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/speculative/plugins/hf_streaming_dataset.py`:
- Around line 114-115: Update the template handling before the regex check in
the dataset formatting logic to support dict-valued tokenizer.chat_template by
selecting the effective named template, preferably the "default" entry, before
calling re.search. Preserve the existing empty-string fallback and intended
RuntimeError behavior when no generation marker exists, and add a regression
test covering a dictionary-valued chat template.
In
`@tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml`:
- Around line 99-100: Update the environment configuration in the MiniMax-M3
launcher recipe to include the required QUANT_CFG variable, setting it to the
appropriate quantization configuration for this job while preserving the
existing HF_MODEL_CKPT entry.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2fbb5455-915d-4f69-8a64-1e6c215dcbdc
📒 Files selected for processing (8)
examples/speculative_decoding/requirements.txtmodelopt/torch/speculative/plugins/hf_streaming_dataset.pymodelopt/torch/speculative/plugins/modeling_fakebase.pymodelopt/torch/speculative/plugins/modeling_final_norm.pytests/unit/torch/speculative/plugins/test_hf_streaming_dataset.pytools/launcher/common/eagle3/train_eagle_streaming.shtools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yamltools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja
🚧 Files skipped from review as they are similar to previous changes (4)
- examples/speculative_decoding/requirements.txt
- modelopt/torch/speculative/plugins/modeling_fakebase.py
- modelopt/torch/speculative/plugins/modeling_final_norm.py
- tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja
| template = getattr(tokenizer, "chat_template", None) or "" | ||
| if not re.search(r"\{%-?\s*generation\b", template): |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target file and surrounding logic.
sed -n '1,220p' modelopt/torch/speculative/plugins/hf_streaming_dataset.py
# Find any existing tests or usages related to chat_template handling.
rg -n "chat_template|generation\\b|hf_streaming_dataset" modelopt tests -SRepository: NVIDIA/Model-Optimizer
Length of output: 40279
🌐 Web query:
Transformers 4.56 chat_template dict effective template apply_chat_template tokenizer.chat_template dictionary named templates
💡 Result:
In the Hugging Face Transformers library, the tokenizer.chat_template attribute typically stores a string containing a Jinja2 template used to format conversation history [1][2]. However, for models that support multiple use cases (such as standard chat, tool-use, or retrieval-augmented generation), tokenizer.chat_template can be defined as a dictionary [3][4]. When tokenizer.chat_template is a dictionary, each key represents the name of a specific template [3][5]. The apply_chat_template method handles these named templates with the following logic: 1. Default Behavior: In most scenarios, apply_chat_template automatically looks for a template named default [3][4]. If this key is missing, the method will raise an error [6][7]. 2. Tool Use: If the user passes a tools argument to apply_chat_template and a template named tool_use exists within the dictionary, the method will automatically prioritize and use that template [3][2]. 3. Explicit Selection: To use a template other than the default or the tool-use-specific template, you must explicitly provide the name of the desired template by passing it to the chat_template argument of the apply_chat_template method (e.g., apply_chat_template(messages, chat_template="rag")) [6][4]. Hugging Face generally recommends using a single, unified template with Jinja2 conditional logic (such as if tools is defined) to handle different use cases, rather than relying on a dictionary of multiple templates, to reduce complexity [3][1].
Citations:
- 1: https://huggingface.co/docs/transformers/v4.48.0/en/chat_templating
- 2: https://huggingface.co/docs/transformers/v4.45.1/chat_templating
- 3: https://huggingface.co/docs/transformers/main/en/chat_template_advanced
- 4: https://huggingface.co/docs/transformers/v4.49.0/en/chat_template_advanced
- 5: https://huggingface.co/docs/transformers/v4.49.0/chat_template_advanced
- 6: https://huggingface.co/docs/transformers/v4.53.3/chat_templating
- 7: https://huggingface.co/docs/transformers/v4.53.1/chat_templating
🏁 Script executed:
python3 - <<'PY'
try:
import transformers
import inspect
from transformers import tokenization_utils_base
print("transformers", transformers.__version__)
src = inspect.getsource(tokenization_utils_base.PreTrainedTokenizerBase.apply_chat_template)
for needle in ["chat_template", "dict", "default", "Template"]:
if needle in src:
print(f"FOUND:{needle}")
# Print the relevant chunk around chat_template handling if available
lines = src.splitlines()
for i, line in enumerate(lines):
if "chat_template" in line or "dict" in line or "default" in line:
start = max(0, i - 6)
end = min(len(lines), i + 12)
print("\n--- chunk ---")
print("\n".join(f"{j+1}: {lines[j]}" for j in range(start, end)))
break
except Exception as e:
print(type(e).__name__, e)
PYRepository: NVIDIA/Model-Optimizer
Length of output: 211
Handle dict-valued tokenizer.chat_template before regex matching. tokenizer.chat_template can be a dict of named templates, and passing it straight to re.search raises TypeError before the intended RuntimeError. Resolve the effective template string first (for example, the "default" template) and add a regression test.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/speculative/plugins/hf_streaming_dataset.py` around lines 114
- 115, Update the template handling before the regex check in the dataset
formatting logic to support dict-valued tokenizer.chat_template by selecting the
effective named template, preferably the "default" entry, before calling
re.search. Preserve the existing empty-string fallback and intended RuntimeError
behavior when no generation marker exists, and add a regression test covering a
dictionary-valued chat template.
| environment: | ||
| - HF_MODEL_CKPT: <<global_vars.hf_model>> |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Set the required QUANT_CFG environment variable.
New launcher model configurations must declare the intended quantization config; this recipe currently provides none. Set the appropriate value for this MiniMax-M3 job.
As per path instructions, “Set QUANT_CFG environment variable … when adding a new model config.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml`
around lines 99 - 100, Update the environment configuration in the MiniMax-M3
launcher recipe to include the required QUANT_CFG variable, setting it to the
appropriate quantization configuration for this job while preserving the
existing HF_MODEL_CKPT entry.
Source: Path instructions
2 participants
What does this PR do?
Type of change: new example + bug fixes
Adds a MiniMax-M3 DSpark streaming training recipe under
tools/launcher/examples/MiniMaxAI/MiniMax-M3/, plus the four fixes it needs to actually run. Each fix addresses a failure mode that is silent or misleading without it:modeling_final_norm.py,modeling_fakebase.py): M3 uses a gemma-style final RMSNorm ((1 + weight)scale, fp32 multiply-then-cast). Selecting the norm bymodel_typealone picks plainrmsnorm— MiniMax's VL remote code coerces itsmodel_type-lesstext_configto mixtral — which silently drops the+1and corrupts the distillation target. New_FinalGemmaRMSNorm+ selection by the explicituse_gemma_normconfig flag (only MiniMax sets it).answer_only_loss(hf_streaming_dataset.py): with a fast tokenizer whose chat template has no{% generation %}tags,apply_chat_template(return_assistant_tokens_mask=True)only warns and returns an all-zero mask — training runs at zero loss on every sample with no other symptom. Now raises at tokenization with an actionable message.SERVE_BLOCK_SIZEknob (train_eagle_streaming.sh): nemo_run exports env values unquoted, so a multi-tokenSERVE_EXTRA_ARGS="--trust-remote-code --block-size 128"loses everything after the first token. M3's MSA sparse attention requires KV block 128 (ValueError: No common block size for 16at engine init otherwise), so--block-sizegets a dedicated single-token knob.transformerspin to<5.13(matchpyproject.toml): the old<5.4pin downgrades recent vLLM containers (e.g. transformers 5.12.1, which also provides in-treeminimax_m3_vl) and breaksvllm serve(ALLOWED_LAYER_TYPESneeds >=5.5.3).The example itself encodes the validated M3 specifics: generation-tagged chat template copy (required for
answer_only_loss— see fix 2), draft dims + baserope_theta=5e6set explicitly (not inherited), mask token 200063 (reserved slot; added tokens end at 200060),EAGLE_CAPTURE_IDS= draft defaulttarget_layer_ids+1+ final layer,trust_remote_codeat serve/export, and AWS-EFA NIXL notes (UCX segfaults at agent init on EFA nodes; LIBFABRIC required there).Usage
Testing
tests/unit/torch/speculative/plugins/: 129 passed inside a current vLLM x86_64 nightly container (transformers 5.12.1), including 3 new tests for the assistant-mask guard.input_idsidentical to the original template, mask covers exactly the assistant turns (think prefix + content + eos), contiguous, no user-prompt leak.ruff check/ruff format(0.15.20) clean.Before your PR is "Ready for review"
use_gemma_norm=True, previously mis-normed; the guard turns a silent zero-loss run into an error;SERVE_BLOCK_SIZEis opt-in; the pin relax widens the allowed range)🤖 Generated with Claude Code
Summary by CodeRabbit
SERVE_BLOCK_SIZEsupport for vLLM serve launches.answer_only_lossmasking validation with clear, fail-fast errors when required{% generation %}markers are missing or when using a non-fast tokenizer.answer_only_lossbehavior.