[TRTLLM-14255][fix] migrate MiniMax M3 to loader v2 for TP8 support - #16468
Conversation
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
📝 WalkthroughWalkthroughMiniMax-M3 gains a dedicated HF weight mapper for MoE gate renaming and KV tensor duplication across quantization modes. Text and VL loaders now accept structured mapper controls, with tests covering registration, tensor handling, parameter mapping, and v2 loading. ChangesMiniMax-M3 weight loading
Sequence Diagram(s)sequenceDiagram
participant Caller
participant MiniMaxM3ForCausalLM
participant MiniMaxM3HfWeightMapper
participant DecoderModelForCausalLM
Caller->>MiniMaxM3ForCausalLM: call load_weights
MiniMaxM3ForCausalLM->>MiniMaxM3HfWeightMapper: initialize mapper and parameter map
MiniMaxM3ForCausalLM->>DecoderModelForCausalLM: forward weights and loader controls
🚥 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.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/models/modeling_minimaxm3.py (1)
1473-1479: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse Python 3.10 built-in generic annotations consistently.
tensorrt_llm/_torch/models/modeling_minimaxm3.py#L1473-L1479: replace unparameterizedDict/Optionalannotations with precisedict[...]and| Nonetypes.tensorrt_llm/_torch/models/modeling_minimaxm3.py#L1625-L1631: apply the same public loader signature convention.As per coding guidelines, “prefer built-in generic types and
|.”🤖 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 1473 - 1479, Update the load_weights signature in tensorrt_llm/_torch/models/modeling_minimaxm3.py at lines 1473-1479 to use precise built-in dict[...] annotations and | None instead of unparameterized Dict/Optional types. Apply the same public loader signature convention at lines 1625-1631, preserving the existing parameter value types and behavior.Source: Coding guidelines
tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py (1)
35-55: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winParameterize the weight mapping types.
weights: dictand-> dictleave this callback contract unchecked; use precisedict[...]types for the tensor mapping. As per coding guidelines, “Annotate every function” and “prefer built-in generic types.”🤖 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/checkpoints/hf/minimaxm3_weight_mapper.py` around lines 35 - 55, Update _duplicate_kv_weights to use precise built-in generic dict[...] annotations for both the weights parameter and return type, specifying the key and tensor value types used by the mapping. Keep the callback behavior and duplication logic unchanged.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.
Inline comments:
In `@tests/unittest/_torch/models/test_minimax_m3.py`:
- Around line 883-894: Update the environment handling around
_load_weights_impl_v2 to preserve the existing
TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL value before setting it, then restore
that value in finally; remove the variable only when it was previously unset,
ensuring later tests retain their original loader-parallelism setting.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py`:
- Around line 35-55: Update _duplicate_kv_weights to use precise built-in
generic dict[...] annotations for both the weights parameter and return type,
specifying the key and tensor value types used by the mapping. Keep the callback
behavior and duplication logic unchanged.
In `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 1473-1479: Update the load_weights signature in
tensorrt_llm/_torch/models/modeling_minimaxm3.py at lines 1473-1479 to use
precise built-in dict[...] annotations and | None instead of unparameterized
Dict/Optional types. Apply the same public loader signature convention at lines
1625-1631, preserving the existing parameter value types and behavior.
🪄 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: 72b3a08f-1250-46a5-8a75-bce316a39326
📒 Files selected for processing (4)
tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.pytests/unittest/_torch/models/test_minimax_m3.py
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run |
|
/bot run |
|
PR_Github #60434 [ run ] triggered by Bot. Commit: |
|
PR_Github #60434 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60463 [ run ] triggered by Bot. Commit: |
|
PR_Github #60463 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60723 [ run ] triggered by Bot. Commit: |
|
PR_Github #60723 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60765 [ run ] triggered by Bot. Commit: |
|
PR_Github #60765 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60818 [ run ] triggered by Bot. Commit: |
|
PR_Github #60818 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61005 [ run ] triggered by Bot. Commit: |
|
PR_Github #61005 [ run ] completed with state |
…VIDIA#16468) Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
3 participants
Summary by CodeRabbit
New Features
Bug Fixes
Description
MiniMax M3 defines variadic
load_weightsmethods, so the model loader does not detect explicitweight_mappersupport and silently routes M3 through the legacy weight loader. This prevents correct TP8 loading when the checkpoint has fewer K/V heads than the TP degree because the MXFP8 K/V scale tensors are not duplicated with the K/V weights.This PR migrates MiniMax M3 models to loader v2 using a model-specific
MiniMaxM3HfWeightMapper. The mapper keeps the existing routed-gate bias rename local to M3 and duplicates K/V weights plus both supported MXFP8 scale spellings (weight_scaleandweight_scale_inv) when required by GQA tensor parallelism. It preserves the existing NVFP4 behavior and accepts arbitraryBaseWeightMapperimplementations without assuming mapper-specific attributes.Test Coverage
_load_weights_impl_v2; full GPU execution is left to CI.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions).
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.