[https://nvbugs/6417488][fix] Rename M3 override to _prepare_page_table_tensor(index_mapper_capacity) and… - #15984
[https://nvbugs/6417488][fix] Rename M3 override to _prepare_page_table_tensor(index_mapper_capacity) and…#15984trtllm-agent wants to merge 3 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughMiniMax-M3 vision normalization now uses ChangesMiniMax-M3 vision normalization
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: 🔵 Low · up to The change removes a GB300 NVFP4 test waiver, but the exact passing result is not documented while a related waiver remains. The PR is otherwise mergeable with explicit owner confirmation of that test outcome. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description explains the root cause, the implementation fix, the test plan, and the related bug. It does not use the template headings or include the PR checklist, but the required technical information is mostly complete. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Running upstream PR NVIDIA#15984's variant (which keeps the base V-K stride derivation for kv_offset) against MiniMax-M3 at TP4 showed the derivation does NOT assert on coalesced pools — our comment's claimed justification was wrong. The real reason for pinning zero: no M3 consumer reads the value, and zero keeps init independent of the base's stride assumptions. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
2f0c920 to
b9d6414
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58453 [ run ] triggered by Bot. Commit: |
|
My local try shows this change does not solve all the problem. MiniMaxM3KVCacheManagerV2._get_batch_cache_indices_by_pool_id() |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py`:
- Around line 435-443: The MinimaxM3 cache manager override is bypassing the
base SWA scratch-copy setup, so the scratch-reuse tensors are never initialized
when enable_swa_scratch_reuse is enabled. In CacheManager’s override where
host_kv_cache_block_offsets is allocated, make sure to preserve the base hook
behavior by calling _prepare_swa_scratch_copy_tensors(index_mapper_capacity) on
the enabled path, or otherwise delegate to the base initialization before
returning, so later scratch-reuse code in CacheManager and related SWA paths can
find the expected tensors.
🪄 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: 9dce3103-99a0-4840-98a3-6ed2447e3d1a
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
|
/bot run --disable-fail-fast |
|
PR_Github #58461 [ run ] triggered by Bot. Commit: |
|
PR_Github #58453 [ run ] completed with state |
|
PR_Github #58461 [ run ] completed with state
|
|
Waiting #16218 for a more proper fix. |
59133cb to
339132d
Compare
|
/bot run --only-qa-verify |
|
PR_Github #63371 [ run ] triggered by Bot. Commit: |
|
PR_Github #63371 [ run ] completed with state |
41a28d1 to
1562764
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. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py`:
- Around line 1184-1197: Update MiniMaxVLLayerNorm.reset_parameters to
initialize affine parameters during ordinary construction while skipping
initialization under MetaInitMode; use the existing mode-detection mechanism and
standard LayerNorm initialization behavior. Ensure incomplete strict=False VL
checkpoint loads cannot reach forward with undefined normalization weights, and
add coverage for both meta and non-meta construction paths.
🪄 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: e254b143-6671-4098-9133-1f2b7cbf890e
📒 Files selected for processing (2)
tensorrt_llm/_torch/models/modeling_minimaxm3_vl.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
[by Codex] @yizhang-nv Could you please review PR #15984 for the KV-cache manager changes? Thanks! |
brnguyen2
left a comment
There was a problem hiding this comment.
Two things to sort out before this merges.
Title/description don't match the diff. The PR text is entirely about renaming a MiniMaxM3KVCacheManagerV2 override to _prepare_page_table_tensor and populating the base pool-mapping state. None of that is in this branch — the diff is a MiniMaxVLLayerNorm subclass plus one waiver removal. Please rewrite the description to describe what actually changed (the title commit message [nvbugs/6417488][fix] Skip MiniMax M3 VL LayerNorm init under MetaInitMode is accurate).
The LayerNorm change doesn't obviously justify un-waiving the test. model_loader.py:545-553 wraps construction in try: with MetaInitMode() ... except Exception: logger.info("Fallback to regular model init"); model = AutoModelForCausalLM.from_config(config). So the MetaInitException this fixes was already non-fatal — it cost a host-side init pass, it didn't fail the run. Whatever made TestMiniMaxM3::test_nvfp4[use_msa=False] fail on GB300 is downstream of that, so removing the waiver needs a passing run on GB300 as evidence, not a mechanism argument. Please link the CI job.
The change itself is fine and matches NemotronLayerNormPlus1 in modeling_nemotron.py:45-61.
| full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=True] SKIP (https://nvbugs/6546909) | ||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771) | ||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6473161) | ||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6422502) |
There was a problem hiding this comment.
test_mxfp8[use_msa=False] on the line above is waived under the same NVBug (6422502) as the test_nvfp4 line being removed here. If the LayerNorm fix resolves that bug, both should be un-waived; if it doesn't, neither should be. Please make the two consistent and say which GB300 run demonstrates the pass.
There was a problem hiding this comment.
The number 6422502 should not be waived, and this bug no longer exists.
|
[by Codex] @yizhang-nv Could you review this PR? Thanks! |
1562764 to
3e4140e
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. |
|
[by Codex] @yizhang-nv Could you review this PR? Thanks! |
3e4140e to
7f4d7a6
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. |
…tMode MiniMaxVLVisionTransformer built plain nn.LayerNorm modules, whose reset_parameters() calls init.ones_/init.zeros_ -> aten.fill_.Scalar. That op is not on MetaInitMode's allowlist, so building the model raised MetaInitException at pre_layrnorm and model_loader fell back to regular init, materializing the whole 233 GB checkpoint on the host per rank. Add MiniMaxVLLayerNorm, an nn.LayerNorm subclass whose reset_parameters() is a no-op, and use it at all three vision-tower LN sites. This follows the existing NemotronLayerNormPlus1 convention. All 130 vision LN tensors are present in the checkpoint, so the loaded weights supply the values the skipped init would have written; state_dict keys and forward numerics are unchanged. Meta-init now succeeds for the full model (params meta: 1290, cpu: 697 -- 24.0 GB materialized of 233.1 GB total). Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
7f4d7a6 to
f7941ba
Compare
crazydemo
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The LayerNorm change is a sound, established pattern (matches NemotronLayerNormPlus1), but the waiver removal lacks CI evidence and the PR description does not match the diff — resolve both before merging.
Issues
- [MAJOR]
tests/integration/test_lists/waives.txt:212- nvfp4 waiver removed on a mechanism argument; sibling mxfp8 waiver for the same bug 6422502 stays - [MINOR]
tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py:1196- reset_parameters no-op leaves affine params uninitialized on the non-meta path
QA view
- Test coverage: missing - only a waiver line was removed; no unit test exercises MiniMaxVLLayerNorm's meta vs non-meta init paths. The re-enabled integration test is not backed by a linked passing run.
- SM coverage: the LayerNorm code is architecture-independent, but the removed waiver re-enables a GB300/nvfp4 (Blackwell) test with no linked GB300 CI job — unverified on that arch.
- Test code: nvfp4 un-waived while mxfp8 for the identical bug 6422502 remains waived (inconsistent); no positive test added.
- Test time: small - re-activates one previously-skipped GB300 nvfp4 accuracy case.
- Needs
/qa-verify: yes - waiver removal justified by mechanism not evidence, plus an untested source change; confirm test_nvfp4[use_msa=False] passes on GB300.
Does this actually fix 6417488?
Unclear/no. The description attributes the bug to a KV-cache-manager hook rename (_build_pool_mapping_tensors → _prepare_page_table_tensor) tripping exact_div — none of that code is in this diff. The diff only adds MiniMaxVLLayerNorm and removes a waiver, which is the separate MetaInit/LayerNorm issue tracked by 6422502. The KV-cache incompatibility (num_blocks_per_seq, follow-up #16218) noted in the release-notes block is still open. Please rewrite the description to match the actual change.
Possible new issues
- The re-enabled GB300 nvfp4 test may still fail on main: model_loader already caught MetaInitException and fell back to regular init, so that exception was non-fatal and the real failure is downstream of it.
- reset_parameters returning without initializing means a strict=False checkpoint load that misses an LN key would reach forward() with uninitialized normalization weights and silently wrong output.
What I could not verify
Whether TestMiniMaxM3::test_nvfp4[use_msa=False] now passes on GB300 (no CI job in context), whether all LN keys are guaranteed present in every supported MiniMax VL checkpoint, and the state of the KV-cache path described in 6417488 (not shown in this diff).
Automated review by NVCortex Lite, run by @crazydemo.
| full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771) | ||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6473161) | ||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True] SKIP (https://nvbugs/6525010) | ||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6422502) |
There was a problem hiding this comment.
[MAJOR] nvfp4 waiver removed without GB300 evidence; sibling mxfp8 waiver for same bug stays
The removed line un-waived full:GB300/.../TestMiniMaxM3::test_nvfp4[use_msa=False] (nvbugs/6422502), yet the line just above (212, kept) still waives TestMiniMaxM3::test_mxfp8[use_msa=False] for the SAME bug 6422502. If one code change resolved 6422502, both should clear; keeping mxfp8 waived while dropping nvfp4 is internally inconsistent and suggests the underlying GB300 failure is not fully understood. More importantly, the fix is a mechanism argument: model_loader wraps MetaInitMode construction in try/except and already falls back to regular init on MetaInitException, so that exception was non-fatal — whatever made test_nvfp4 fail on GB300 is downstream of it. Re-enabling this test on main without a linked passing GB300 CI run risks re-introducing a red test into the pipeline. Fix: link the passing GB300 nvfp4 job in the PR, and either also clear mxfp8 or explain why it stays waived.
| skipped init would have written. | ||
| """ | ||
|
|
||
| def reset_parameters(self) -> None: |
There was a problem hiding this comment.
[MINOR] reset_parameters no-op leaves affine params uninitialized outside meta path
nn.LayerNorm.init (elementwise_affine=True) allocates weight/bias and calls reset_parameters to fill them with ones/zeros. Overriding it to pass means that under ordinary (non-MetaInitMode) construction the affine tensors are left as uninitialized memory until the checkpoint overwrites them. This is safe only if every LN key is actually present in the load. If a MiniMax VL checkpoint is ever loaded with strict=False and misses a layer_norm1/layer_norm2/pre_layrnorm key, forward() will normalize with garbage weights and produce silently wrong outputs rather than erroring. The pattern mirrors NemotronLayerNormPlus1, so it is acceptable, but consider gating on MetaInitMode detection (init only when not in meta mode) so the non-meta path keeps standard init. At minimum add a unit test asserting the three LN weights are populated after a normal (non-meta) build + load.
crazydemo
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: A waiver is being removed on a mechanism argument rather than a linked passing GB300 run, the sibling mxfp8 waiver for the same bug stays, and there is no unit test for the LayerNorm change. QA should confirm test_nvfp4[use_msa=False] actually passes on GB300 before this is trusted.
Automated review by NVCortex Lite, run by @crazydemo.
|
[by Codex] @yizhang-nv Friendly reminder: could you please review or revisit this PR when you have a chance? Thanks! |
Signed-off-by: WeiHaocheng <fredw@nvidia.com>
Signed-off-by: WeiHaocheng <fredw@nvidia.com>
|
Looks like the issue is resolve by |
|
The issues that this PR aims to fix have already been addressed in the main branch. |
Summary
Test plan
Links
Dev Engineer Review
MiniMaxVLLayerNorm, which skipsreset_parameters.TestMiniMaxM3::test_nvfp4[use_msa=False]waiver for NVBug 6422502.num_blocks_per_seq. This requires follow-up with the proper fix from PR#16218.QA Engineer Review
tests/integration/test_lists/waives.txt.TestMiniMaxM3::test_nvfp4[use_msa=False].Verdict: needs follow-up