[https://nvbugs/6445375][fix] Kept the gated MiniMaxVLLayerNorm fix and added 4 CPU-only meta-init… - #17321
[https://nvbugs/6445375][fix] Kept the gated MiniMaxVLLayerNorm fix and added 4 CPU-only meta-init…#17321trtllm-agent wants to merge 2 commits into
MiniMaxVLLayerNorm fix and added 4 CPU-only meta-init…#17321Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughThe MiniMax vision model now uses ChangesMiniMax vision meta initialization
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
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.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/models/test_minimax_m3_vl.py (1)
833-875: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the new test functions.
Add
-> Noneto all four tests. Typemonkeypatchaspytest.MonkeyPatch. Test coverage is insufficient because these four tests are not listed intests/integration/test_lists/.🤖 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 `@tests/unittest/_torch/models/test_minimax_m3_vl.py` around lines 833 - 875, Add -> None return annotations to test_vision_tower_builds_under_meta_init, test_meta_init_still_rejects_plain_layer_norm_init, test_layer_norm_off_meta_init_matches_upstream, and test_layer_norm_without_affine_params_builds_under_meta_init; annotate the monkeypatch parameter in the second test as pytest.MonkeyPatch. Add all four test names to the appropriate tests/integration/test_lists/ entry so they are included in integration coverage.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_vl.py`:
- Around line 823-875: Register the four MiniMax M3 VL
tests—test_vision_tower_builds_under_meta_init,
test_meta_init_still_rejects_plain_layer_norm_init,
test_layer_norm_off_meta_init_matches_upstream, and
test_layer_norm_without_affine_params_builds_under_meta_init—in an appropriate
CI test-list entry such as l0_cpu.yml, then run pytest tests/unittest/ to verify
discovery and execution.
---
Nitpick comments:
In `@tests/unittest/_torch/models/test_minimax_m3_vl.py`:
- Around line 833-875: Add -> None return annotations to
test_vision_tower_builds_under_meta_init,
test_meta_init_still_rejects_plain_layer_norm_init,
test_layer_norm_off_meta_init_matches_upstream, and
test_layer_norm_without_affine_params_builds_under_meta_init; annotate the
monkeypatch parameter in the second test as pytest.MonkeyPatch. Add all four
test names to the appropriate tests/integration/test_lists/ entry so they are
included in integration coverage.
🪄 Autofix
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: 0d04e3d5-50b0-48e0-a4af-e4f918b31a91
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_minimaxm3_vl.pytests/integration/test_lists/waives.txttests/unittest/_torch/models/test_minimax_m3_vl.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
| def _build_tiny_vision_model() -> MiniMaxVLVisionModel: | ||
| cfg = CLIPVisionConfig.from_dict_or_obj(_tiny_vision_config()) | ||
| return MiniMaxVLVisionModel( | ||
| config=cfg, | ||
| text_hidden_size=16, | ||
| projector_hidden_size=16, | ||
| dtype=torch.float32, | ||
| ) | ||
|
|
||
|
|
||
| def test_vision_tower_builds_under_meta_init(): | ||
| """The vision tower must construct inside ``MetaInitMode``. | ||
|
|
||
| See :class:`MiniMaxVLLayerNorm` for why a plain ``nn.LayerNorm`` here aborts | ||
| meta-init for the whole M3 model. | ||
| """ | ||
| with MetaInitMode(): | ||
| model = _build_tiny_vision_model() | ||
|
|
||
| layer_norms = [m for m in model.modules() if isinstance(m, nn.LayerNorm)] | ||
| assert layer_norms, "expected layer norms in the vision tower" | ||
| # pre_layrnorm + layer_norm1/2 per encoder layer. | ||
| assert len(layer_norms) == 1 + 2 * len(model.vision_model.encoder.layers) | ||
| for ln in layer_norms: | ||
| assert ln.weight.is_meta | ||
| assert ln.bias.is_meta | ||
|
|
||
|
|
||
| def test_meta_init_still_rejects_plain_layer_norm_init(monkeypatch): | ||
| """Control for :func:`test_vision_tower_builds_under_meta_init`. | ||
|
|
||
| Restoring the upstream ``reset_parameters`` must bring the exception back, | ||
| otherwise that test could pass without the skip doing any work. | ||
| """ | ||
| monkeypatch.setattr(MiniMaxVLLayerNorm, "reset_parameters", nn.LayerNorm.reset_parameters) | ||
| with pytest.raises(MetaInitException, match="fill_"): | ||
| with MetaInitMode(): | ||
| _build_tiny_vision_model() | ||
|
|
||
|
|
||
| def test_layer_norm_off_meta_init_matches_upstream(): | ||
| """Off meta the skip must not engage, else a checkpoint-free build keeps | ||
| uninitialized ``torch.empty`` storage instead of ones/zeros.""" | ||
| ln = MiniMaxVLLayerNorm(8, dtype=torch.float32) | ||
| assert torch.equal(ln.weight, torch.ones(8)) | ||
| assert torch.equal(ln.bias, torch.zeros(8)) | ||
|
|
||
|
|
||
| def test_layer_norm_without_affine_params_builds_under_meta_init(): | ||
| """``elementwise_affine=False`` registers ``weight`` as ``None``.""" | ||
| with MetaInitMode(): | ||
| ln = MiniMaxVLLayerNorm(8, elementwise_affine=False, dtype=torch.float32) | ||
| assert ln.weight is None |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
for test_name in \
test_vision_tower_builds_under_meta_init \
test_meta_init_still_rejects_plain_layer_norm_init \
test_layer_norm_off_meta_init_matches_upstream \
test_layer_norm_without_affine_params_builds_under_meta_init
do
echo "== $test_name =="
rg -n -F "$test_name" tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== test file status =="
git status --short -- tests/unittest/_torch/models/test_minimax_m3_vl.py
echo "== matching test-list entries by file or suite =="
rg -n -i -F "test_minimax_m3_vl" tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
rg -n -i -F "minimax" tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
echo "== available list files =="
find tests/integration/test_lists -maxdepth 2 -type f -print | sort | sed -n '1,160p'
echo "== nearby test definitions =="
rg -n "^(def test_|class Test)" tests/unittest/_torch/models/test_minimax_m3_vl.py | sed -n '1,120p'Repository: NVIDIA/TensorRT-LLM
Length of output: 16415
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== test-list scope =="
sed -n '1,180p' tests/integration/test_lists/test-db/README.md
sed -n '1,180p' tests/integration/test_lists/qa/README.md
echo "== unit-test entries in integration lists =="
rg -n -i "unittest/|tests/unittest|pytest.*unittest" tests/integration/test_lists/test-db tests/integration/test_lists/qa | sed -n '1,160p'
echo "== file history summary =="
git log -5 --oneline -- tests/unittest/_torch/models/test_minimax_m3_vl.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 29614
Register the new unit tests in CI.
- Added tests:
test_vision_tower_builds_under_meta_init,test_meta_init_still_rejects_plain_layer_norm_init,test_layer_norm_off_meta_init_matches_upstream, andtest_layer_norm_without_affine_params_builds_under_meta_init. - The test file and these test functions are absent from the CI test lists. Add the four tests to a suitable entry, such as
tests/integration/test_lists/test-db/l0_cpu.yml. - Run
pytest tests/unittest/for this change. - Coverage verdict: insufficient.
🤖 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 `@tests/unittest/_torch/models/test_minimax_m3_vl.py` around lines 823 - 875,
Register the four MiniMax M3 VL tests—test_vision_tower_builds_under_meta_init,
test_meta_init_still_rejects_plain_layer_norm_init,
test_layer_norm_off_meta_init_matches_upstream, and
test_layer_norm_without_affine_params_builds_under_meta_init—in an appropriate
CI test-list entry such as l0_cpu.yml, then run pytest tests/unittest/ to verify
discovery and execution.
Sources: Coding guidelines, Path instructions
|
/bot run --only-qa-verify test accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8],accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The code fix looks right and is scoped correctly: LayerNorm's constant fill is the only init in this tower that MetaInitMode rejects (nn.Linear/nn.Conv3d bottom out in uniform_, which is allowlisted), and gating the skip on weight.is_meta avoids the torch.empty-garbage trap that an unconditional skip would create. The ablation test is worth having.
My concern is the waiver removal. MetaInitException is caught at tensorrt_llm/_torch/pyexecutor/model_loader.py:549 and the loader falls back to regular init — so the exception in the bug's log is a recoverable path, not the abort itself. Fixing it restores the meta-init fast path (real win: no per-rank host materialization), but it doesn't by itself establish that the two unwaived B300 tests now pass. Could you post the actual passing run for both test_auto_dtype[tp_size=8-ep_size=8] and test_nvfp4[use_msa=False] on B300? If only test_auto_dtype was run, I'd keep the test_nvfp4 waiver and drop it in a follow-up.
| full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4gpus_static_eplb[moe_backend=WIDEEP] SKIP (https://nvbugs/6546609) | ||
| full:B300/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True] SKIP (https://nvbugs/6475346) | ||
| full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6445375) | ||
| full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) |
There was a problem hiding this comment.
test_nvfp4[use_msa=False] is being unwaived here for B300, but the same parametrization is still waived on B200 (line 160), GB200 (line 210) and GB300 (line 243) under different bugs, and B300 keeps test_mxfp8[use_msa=False] waived under 6424188 (this line). If B300 is the only platform where this case is expected to pass, that's fine — but please confirm it actually ran green rather than inferring it from the test_auto_dtype fix. Otherwise this re-enables a case that fails everywhere else and the next failure gets triaged as a new regression.
| skipped values are always overwritten at load time. | ||
| """ | ||
|
|
||
| def reset_parameters(self) -> None: |
There was a problem hiding this comment.
The meta gate is the right call — worth noting modeling_nemotron.py:59 has the unconditional variant of this (reset_parameters → pass), which leaves torch.empty storage on any checkpoint-free build. Since this is now the second copy of the same workaround, consider hoisting it to a shared module (e.g. next to the other norms in _torch/modules/) so the next VLM doesn't rediscover it; happy for that to be a follow-up rather than this PR.
| onto its regular-init fallback (hundreds of GB of host allocation per rank | ||
| for M3). Every layer-norm slot here is covered by the checkpoint, so the | ||
| skipped values are always overwritten at load time. | ||
| """ |
There was a problem hiding this comment.
The docstring asserts "every layer-norm slot here is covered by the checkpoint." That's the safety argument for the whole change, and it's checked by nothing — if a future config adds a norm the weight mapper doesn't populate, the parameter stays uninitialized on the meta path and silently produces garbage rather than failing loudly. Is there an existing post-load check that no parameter is left on meta? If so, worth naming it here; if not, that's the guard this comment is standing in for.
|
PR_Github #64168 [ run ] triggered by Bot. Commit: |
|
PR_Github #64168 [ run ] completed with state |
|
/bot run --only-qa-verify test accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8],accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] |
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 (1)
tests/integration/test_lists/waives.txt (1)
1-1: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winAdd the NVIDIA copyright header before merge.
tests/integration/test_lists/waives.txtis modified, but no NVIDIA copyright header is present. Add the repository-standard header and use2026, the year of this modification. Ensure the header format is accepted by the waiver-list parser.As per coding guidelines, modified files must contain the NVIDIA copyright header with the year of the latest meaningful modification.
🤖 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 `@tests/integration/test_lists/waives.txt` at line 1, Add the repository-standard NVIDIA copyright header at the beginning of tests/integration/test_lists/waives.txt, using 2026 as the modification year. Preserve the waiver entry and ensure the header format remains valid for the waiver-list parser.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.
Outside diff comments:
In `@tests/integration/test_lists/waives.txt`:
- Line 1: Add the repository-standard NVIDIA copyright header at the beginning
of tests/integration/test_lists/waives.txt, using 2026 as the modification year.
Preserve the waiver entry and ensure the header format remains valid for the
waiver-list parser.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5b44af86-df24-4904-a1d9-840e5cb7fd18
📒 Files selected for processing (1)
tests/integration/test_lists/waives.txt
|
PR_Github #64189 [ run ] triggered by Bot. Commit: |
Plain nn.LayerNorm.reset_parameters() fills weight/bias with aten.fill_.Scalar, which MetaInitMode does not allow, so the vision tower's LayerNorm construction raised MetaInitException and aborted meta-init for the whole model. The loader fell back to regular init, really allocating every weight on the host. Subclass nn.LayerNorm and short-circuit reset_parameters() when the weight is on the meta device, following NemotronLayerNormPlus1. The gate matters: an unconditional skip leaves uninitialized storage for modules built without a checkpoint and NaNs the vision-tower forward tests. All 65 layer-norm slots are present in the checkpoint, so the skipped values are overwritten at load time. Unwaive the two B300 MiniMax-M3 tests filed under this bug. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…VL LayerNorm The MiniMaxVLLayerNorm meta-init skip had no in-tree coverage, so a regression would only surface as a silent fallback to regular model init (hundreds of GB of host allocation per rank) rather than a test failure. Covers the skip working under MetaInitMode, a control that restores the upstream reset_parameters and asserts the exception returns, and the two edge cases the gate depends on: off-meta init keeping ones/zeros, and elementwise_affine=False registering weight as None. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
27209fa to
c27da3b
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
PR_Github #64189 [ run ] completed with state |
5 participants
Summary
nn.LayerNorm.reset_parameters()emitsaten.fill_.Scalarviainit.ones_/zeros_, whichMetaInitMode's allowlist (aten.empty*+normal_/uniform_/log) rejects, soMetaInitExceptionaborts meta-init for the whole M3 model.MiniMaxVLLayerNormfix and added 4 CPU-only meta-init regression tests — including an ablation control that restores upstreamreset_parametersand asserts the exception returns — so the fix is guarded by CI instead of throwaway scratch.pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] -vTest plan
Links
Dev Engineer Review
MiniMaxVLLayerNormwith meta-device initialization gating.LayerNorminitialization for non-meta parameters.auto_dtypeandnvfp4waiver entries for bug6445375.mxfp8waiver.QA Engineer Review
MetaInitMode.nn.LayerNorm.reset_parameters()failure behavior.elementwise_affine=False.auto_dtypeandnvfp4tests are unwaived. Themxfp8waiver remains.