[https://nvbugs/6417488][fix] Rename M3 override to _prepare_page_table_tensor(index_mapper_capacity) and… - #15984
Conversation
WalkthroughMiniMax-M3 vision normalization now uses ChangesMiniMax-M3 vision normalization
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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() |
![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
🤖 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
4803c16 to
70034ff
Compare
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
70034ff to
89ffcfd
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59844 [ run ] triggered by Bot. Commit: |
|
PR_Github #59844 [ run ] completed with state
|
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM — the M3-specific override rename fixes a silently-dead hook (exact_div assert); the base kv_cache_manager_v2 change swaps try/except-KeyError for an all_buffer_ids membership check that is behavior-preserving — default non-sparse models still return None. Covered by a new unit test.
|
[by Codex] @yizhang-nv Could you please review this PR? Thank you! |
27dfb12 to
ca863cf
Compare
|
[by Codex] @yizhang-nv Friendly review reminder: this PR is awaiting your review. Thanks! |
ca863cf to
d373326
Compare
|
/bot run --only-qa-verify |
|
PR_Github #63371 [ run ] triggered by Bot. Commit: |
|
PR_Github #63371 [ run ] completed with state |
d373326 to
41a28d1
Compare
…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>
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
| class MiniMaxVLLayerNorm(nn.LayerNorm): | ||
| """``nn.LayerNorm`` whose parameter init is skipped under ``MetaInitMode``. | ||
|
|
||
| ``nn.LayerNorm.reset_parameters`` calls ``init.ones_``/``init.zeros_``, | ||
| which lower to ``aten.fill_.Scalar``. That op is not on ``MetaInitMode``'s | ||
| allowlist, so a plain ``nn.LayerNorm`` anywhere in the module tree raises | ||
| ``MetaInitException`` and forces the whole model onto the regular-init | ||
| path (a full host copy of the checkpoint). Every LN tensor here is | ||
| present in the checkpoint, so the loaded weights supply the values the | ||
| skipped init would have written. | ||
| """ | ||
|
|
||
| def reset_parameters(self) -> None: | ||
| pass |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Expected: every construction path reaches a complete checkpoint load
# before the first forward through MiniMaxVLLayerNorm.
rg -n -C 8 \
'MiniMaxVLLayerNorm|MiniMaxVLVisionTransformer|load_state_dict|strict\s*=|MetaInitMode'Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py"
printf '%s\n' '--- candidate symbols ---'
rg -n -C 6 'MiniMaxVLLayerNorm|MiniMaxVLVisionTransformer|MetaInitMode|load_state_dict|load_checkpoint|state_dict' "$file" || true
printf '%s\n' '--- repository references ---'
rg -n -C 4 'MiniMaxVLLayerNorm|MiniMaxVLVisionTransformer|MetaInitMode' tensorrt_llm tests examples || true
printf '%s\n' '--- relevant source ranges ---'
sed -n '1,80p' "$file"
sed -n '1140,1320p' "$file"Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- tracked files related to MiniMaxM3 ---'
git ls-files | rg -i 'minimax|vl' | head -200
printf '%s\n' '--- LayerNorm construction and loading APIs ---'
rg -n -C 5 'LayerNorm|load_state_dict|load_checkpoint|from_pretrained|state_dict' tensorrt_llm/_torch/models tensorrt_llm/_torch | head -500Repository: NVIDIA/TensorRT-LLM
Length of output: 47777
🏁 Script executed:
#!/bin/bash
set -euo pipefail
model="tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py"
test="tests/unittest/_torch/models/test_minimax_m3_vl.py"
printf '%s\n' '--- MiniMax vision model and loader ---'
sed -n '1525,1720p' "$model"
printf '%s\n' '--- MiniMax VL tests ---'
rg -n -C 12 'MiniMaxVLLayerNorm|MiniMaxVLVisionModel|load_minimax_m3_vl_state_dict|load_state_dict|forward|missing|strict' "$test"
printf '%s\n' '--- all call sites of the custom loader ---'
rg -n -C 8 'load_minimax_m3_vl_state_dict|MiniMaxVLVisionModel' tensorrt_llm testsRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- model loader integration for MiniMax M3 VL ---'
rg -n -C 12 'minimax|load_weights|vision_tower|load_state_dict' tensorrt_llm/_torch/pyexecutor/model_loader.py tensorrt_llm/_torch/models/modeling_minimaxm3.py tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py
printf '%s\n' '--- MetaInitMode behavior and materialization ---'
sed -n '45,135p' tensorrt_llm/_torch/models/modeling_utils.py
rg -n -C 10 'is_meta_init|MetaInitMode|materialize|load_weights' tensorrt_llm/_torch/pyexecutor/model_loader.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tensorrt_llm/_torch/models/modeling_minimaxm3.py"
printf '%s\n' '--- VL-specific symbols in modeling_minimaxm3.py ---'
rg -n 'MiniMaxM3VL|vision_tower|MiniMaxVLVisionModel|load_minimax|split_multimodal|_VL_PREFIXES|load_weights' "$file"
printf '%s\n' '--- focused VL sections ---'
python3 - <<'PY'
from pathlib import Path
path = Path("tensorrt_llm/_torch/models/modeling_minimaxm3.py")
lines = path.read_text().splitlines()
needles = ("vision_tower", "MiniMaxVLVisionModel", "load_minimax_m3_vl_state_dict",
"split_multimodal_weights", "def load_weights", "class MiniMaxM3")
for i, line in enumerate(lines, 1):
if any(n in line for n in needles):
start = max(1, i - 18)
end = min(len(lines), i + 35)
print(f"\n--- lines {start}-{end} ---")
for j in range(start, end + 1):
print(f"{j}: {lines[j-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- production VL load path ---'
sed -n '2108,2160p' tensorrt_llm/_torch/models/modeling_minimaxm3.py
sed -n '2190,2245p' tensorrt_llm/_torch/models/modeling_minimaxm3.py
printf '%s\n' '--- construction and forward tests without a complete load ---'
sed -n '140,230p' tests/unittest/_torch/models/test_minimax_m3_vl.py
sed -n '298,370p' tests/unittest/_torch/models/test_minimax_m3_vl.py
sed -n '711,742p' tests/unittest/_torch/models/test_minimax_m3_vl.py
printf '%s\n' '--- source-only verifier for custom LayerNorm usage ---'
python3 - <<'PY'
import ast
from pathlib import Path
paths = [
Path("tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py"),
Path("tests/unittest/_torch/models/test_minimax_m3_vl.py"),
]
for path in paths:
tree = ast.parse(path.read_text(), filename=str(path))
print(f"\n{path}")
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == "MiniMaxVLLayerNorm":
print("custom class:", node.lineno)
if isinstance(node, ast.Call):
func = node.func
if isinstance(func, ast.Name) and func.id == "MiniMaxVLLayerNorm":
print("construction:", node.lineno, "in", type(next(
(p for p in ast.walk(tree)
if isinstance(p, (ast.FunctionDef, ast.AsyncFunctionDef))
and p.lineno <= node.lineno <= getattr(p, "end_lineno", p.lineno)),
ast.Module
)).__name__)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 13610
Initialize MiniMaxVLLayerNorm outside MetaInitMode. reset_parameters() leaves torch.empty affine parameters undefined during normal construction. The VL loader also uses strict=False, so partial loads can reach forward() with missing normalization weights. Initialize non-meta parameters, or reject incomplete loads before forward. Add tests for both paths.
🤖 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_vl.py` around lines 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.
Source: MCP tools
|
[by Codex] @yizhang-nv Could you please review PR #15984 for the KV-cache manager changes? Thanks! |
8 participants
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