[TRTLLM-13458][feat] Support Minimax M3 NVFP4 checkpoint - #15857
Conversation
|
/bot run |
📝 WalkthroughWalkthroughThis PR adds Minimax M3 NVFP4 mixed-precision MoE support by injecting per-layer routed-expert quant configs during model loading and applying them during MoE construction, fixes a SwigluBias activation mapping in the TRTLLM-Gen fused MoE kernel selector, and adds new accuracy test coverage. ChangesMinimax M3 NVFP4 MoE quant support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant FromPretrained as ModelConfig.from_pretrained
participant QuantHelper as _set_minimax_m3_moe_quant_config
participant MoEInit as MiniMaxM3MoE.__init__
participant ExpertsHelper as _get_experts_quant_config
participant CreateMoE as create_moe
FromPretrained->>QuantHelper: pretrained_config, layer_quant_config
QuantHelper->>QuantHelper: scan for NVFP4 expert entries, derive moe_layer_freq
QuantHelper-->>FromPretrained: updated layer_quant_config
FromPretrained-->>MoEInit: ModelConfig with quant_config_dict
MoEInit->>ExpertsHelper: model_config, layer_idx
ExpertsHelper->>ExpertsHelper: lookup per-layer entry or fallback
ExpertsHelper-->>MoEInit: experts_quant_config
MoEInit->>CreateMoE: override_quant_config=experts_quant_config
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_minimaxm3.py (1)
355-372: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a return type annotation.
_get_experts_quant_configreturns aQuantConfigbut is unannotated. As per coding guidelines: "Always annotate functions with return types."♻️ Suggested annotation
- `@staticmethod` - def _get_experts_quant_config(model_config: "ModelConfig", layer_idx: int): + `@staticmethod` + def _get_experts_quant_config(model_config: "ModelConfig", + layer_idx: int) -> QuantConfig:🤖 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 `@tensorrt_llm/_torch/models/modeling_minimaxm3.py` around lines 355 - 372, The static helper _get_experts_quant_config in modeling_minimaxm3.py is missing a return type annotation. Update its signature to explicitly declare the returned QuantConfig type, matching the values it already returns from model_config.quant_config or model_config.quant_config_dict, so the method follows the project’s function-annotation guideline.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 355-372: The static helper _get_experts_quant_config in
modeling_minimaxm3.py is missing a return type annotation. Update its signature
to explicitly declare the returned QuantConfig type, matching the values it
already returns from model_config.quant_config or
model_config.quant_config_dict, so the method follows the project’s
function-annotation guideline.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7563a0d1-96e6-4b1b-b28b-f92c45b7760d
📒 Files selected for processing (7)
tensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/references/mmlu.yamltests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txt
|
/bot run |
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #57190 [ run ] triggered by Bot. Commit: |
|
LGTM |
|
PR_Github #57190 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57271 [ run ] triggered by Bot. Commit: |
Add NVFP4 quantization support for the routed MoE experts in Minimax M3, enabling mixed-precision checkpoints (MXFP8 base layers + NVFP4 experts). The NVFP4 checkpoint stores per-linear expert entries in hf_quant_config.json under MIXED_PRECISION. ModelConfig cannot resolve these fine-grained keys to a backend, so _set_minimax_m3_moe_quant_config() detects them and injects coarse "model.layers.N.block_sparse_moe.experts" entries that MiniMaxM3MoE._get_experts_quant_config() can look up per layer. MiniMaxM3MoE.__init__ then passes the result as override_quant_config to create_moe(), selecting the NVFP4 backend for experts while leaving the base-layer quant config unchanged. Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com>
…mapping SwigluBias reuses the SwiGlu kernel path (ActType::SwiGlu == 0); the per-expert alpha/beta/clamp_limit are passed as separate tensors. Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com>
|
/bot run |
|
PR_Github #57317 [ run ] triggered by Bot. Commit: |
|
PR_Github #57271 [ run ] completed with state |
Add TestMiniMaxM3::test_nvfp4 (MMLU + GSM8K) for the nvidia/MiniMax-M3-NVFP4 checkpoint (MXFP8 base layers with NVFP4 routed experts), with matching gsm8k/mmlu accuracy references and QA test-list registration. Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com>
|
/bot run |
|
PR_Github #57343 [ run ] triggered by Bot. Commit: |
|
PR_Github #57317 [ run ] completed with state |
|
PR_Github #57343 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #57369 [ run ] triggered by Bot. Commit: |
|
PR_Github #57369 [ run ] completed with state
|
|
/bot run --disabke-fail-fast |
|
PR_Github #57470 Bot args parsing error: usage: /bot [-h] |
|
/bot run --disable-fail-fast |
|
PR_Github #57500 [ run ] triggered by Bot. Commit: |
|
PR_Github #57500 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57506 [ run ] triggered by Bot. Commit: |
|
PR_Github #57506 [ run ] completed with state |
Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com>
8 participants
Description
Adds NVFP4 quantization support for the routed MoE experts in MiniMax-M3, enabling the
mixed-precision checkpoint layout (MXFP8 base layers + NVFP4 routed experts) published as
nvidia/MiniMax-M3-NVFP4.The NVFP4 checkpoint stores per-linear expert quant entries in
hf_quant_config.jsonunderMIXED_PRECISION(e.g.language_model.model.layers.N.block_sparse_moe.experts.E.w1 -> NVFP4).ModelConfigcannot resolve these fine-grained per-linear keys to a MoE backend directly, soModelConfig._set_minimax_m3_moe_quant_config()detects the NVFP4 expert entries and injectscoarse
model.layers.N.block_sparse_moe.expertskeys thatMiniMaxM3MoE._get_experts_quant_config()looks up per layer.
MiniMaxM3MoE.__init__passes the result asoverride_quant_configtocreate_moe(), selecting the NVFP4-capableTRTLLMGenFusedMoEbackend for the routed expertswhile leaving the base-layer (MXFP8) quant config unchanged.
This also includes a required fix in
fused_moe_trtllm_gen.py: MiniMax-M3's MoE experts useActivationType.SwigluBias(SwiGLU with per-expert alpha/beta/clamp-limit bias tensors), whichwas not mapped to a TRTLLM-Gen activation type.
SwigluBiasreuses the same underlying SwiGLUkernel path (
ActType::SwiGlu == 0); the alpha/beta/clamp-limit values are passed separately astensors, so the fix simply extends the activation-type mapping. Without this fix, constructing the
NVFP4 experts through
TRTLLMGenFusedMoEfails sinceSwigluBiasfalls through the mapping.This fix is a hard dependency of the NVFP4 checkpoint support and is included in this PR rather
than split out, since the NVFP4 path cannot function without it.
Test coverage adds
TestMiniMaxM3::test_nvfp4(MMLU + GSM8K accuracy) for the new checkpoint,following the existing
test_mxfp8/test_auto_dtypepattern in the same test class, withmatching
gsm8k.yaml/mmlu.yamlaccuracy references and QA test-list registration.Test Coverage
tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4(MMLU + GSM8K, TP=8/EP=8, Blackwell-only) — new
tests/integration/defs/accuracy/references/{gsm8k,mmlu}.yaml— new reference thresholds fornvidia/MiniMax-M3-NVFP4tests/integration/test_lists/qa/llm_function_core.txt— registers the new testRelated PRs
Summary by CodeRabbit
New Features
Tests