[None][fix] Fix nemotron-h quant and loading config - #16833
Conversation
WalkthroughUpdates Nemotron-H weight preprocessing to detect structured scale tensors by row count and duplicate compatible KV projection scales. Decoder model initialization now centralizes lm_head quantization selection using checkpoint metadata and model configuration. ChangesQuantization updates
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run --disable-fail-fast |
|
PR_Github #61545 [ run ] triggered by Bot. Commit: |
|
PR_Github #61545 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61572 [ run ] triggered by Bot. Commit: |
|
PR_Github #61572 [ run ] completed with state
|
…ckpoints DecoderModelForCausalLM only resolved the lm_head quant entry from the per-layer quant_config_dict (MIXED_PRECISION checkpoints). A homogeneous checkpoint with a single global quant_algo that quantizes lm_head (and does not exclude it) built lm_head unquantized, so loading its packed NVFP4 weight failed with a shape mismatch (e.g. Nemotron-Nano-3.5 NVFP4: 'size of tensor a (2688) must match tensor b (1344)'). Resolve lm_head_quant_config from config.quant_config when quant_config_dict is None and the global config quantizes, and share the exclude_modules / tie_word_embeddings / lm-head-TP-in-ADP guards across both sources. No-op for existing models: the MIXED_PRECISION path is unchanged, and homogeneous checkpoints that keep lm_head in FP16 list it in exclude_modules, so the exclude guard keeps it unquantized as before. Verified on Nemotron-Nano-3.5 NVFP4 (W4A16_NVFP4 + FP8 KV) on H100: single-GPU and 4-GPU attention-DP + EP now load and generate coherently (with and without MTP). Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ole quant family under GQA+TP NemotronH weight loading gated per-output-row block-scale handling on the quant algo (mamba in_proj on quant_algo == "NVFP4"; attention kv-head scale on has_nvfp4()). Both miss W4A16_NVFP4 and the FP8/MXFP8-scaled tensors in MIXED_PRECISION (W4A4) checkpoints, so under GQA + TP the mamba in_proj block scale was not structurally [z,x,B,C,dt]-split while the weight was (garbage output), and the attention fused-QKV kv weight_scale was not duplicated alongside the replicated kv weight (load size mismatch 1152 vs 1280). Detect a per-output-row block scale by shape (shares dim 0 with the weight), scoped to the NemotronH mapper (override of _duplicate_kv_weights + shape-based in_proj split), leaving the shared HfWeightMapper base untouched. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
d8b4bb6 to
40e733a
Compare
|
/bot run --disable-fail-fast |
![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)
tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py (1)
196-210: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winComplete the helper type annotations.
_num_rowsleavestensoruntyped, and_duplicate_kv_weightsuses an unparameterizeddictand has no return annotation. As per coding guidelines, “Annotate every function” and “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/checkpoints/hf/nemotron_h_weight_mapper.py` around lines 196 - 210, Complete the annotations for _num_rows and _duplicate_kv_weights: type the tensor parameter using the appropriate supported tensor or lazy-slice type, parameterize weights with built-in generic syntax, and add the correct return annotation to _duplicate_kv_weights. Use | for unions where applicable and preserve the existing behavior.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 `@tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py`:
- Around line 196-232: Update _duplicate_kv_weights to call _num_kv_heads() and
pass its returned integer to _duplicate_kv instead of passing the bound method.
Keep the existing duplication conditions and tensor_parallel_size handling
unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py`:
- Around line 196-210: Complete the annotations for _num_rows and
_duplicate_kv_weights: type the tensor parameter using the appropriate supported
tensor or lazy-slice type, parameterize weights with built-in generic syntax,
and add the correct return annotation to _duplicate_kv_weights. Use | for unions
where applicable and preserve the existing 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: eca1dc7e-1ff9-451e-bc3c-5929aa8b9ec9
📒 Files selected for processing (2)
tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.pytensorrt_llm/_torch/models/modeling_utils.py
|
PR_Github #61903 [ run ] triggered by Bot. Commit: |
|
PR_Github #61903 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62034 [ run ] triggered by Bot. Commit: |
|
PR_Github #62034 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62116 [ run ] triggered by Bot. Commit: |
|
PR_Github #62116 [ run ] completed with state |
3 participants
Dev Engineer Review
NemotronHHfWeightMapper.preprocess_weightsto splitmixer.in_proj.*_scalestructurally using an expectedd_in_projdim-0 row-count heuristic (via new_num_rows) rather than checkingquant_algo == "NVFP4"/ relying on string-based NVFP4 detection.in_projweights get split; per-tensor scalar scales remain unchanged._duplicate_kv_weightsso that for replicatedk_proj/v_projweights, companionweight_scaletensors are duplicated when they share the same dim-0 row count as the weight (generalizing beyond NVFP4-only assumptions)._checkpoint_has_lm_head_scale(header check forlm_head.weight_scale) and_resolve_lm_head_quant_configinDecoderModelForCausalLMto centralize/guard whenlm_headshould be quantized.lm_headquantization is now derived fromquant_config_dict["lm_head"]when present; otherwise for homogeneous checkpoints it usesconfig.quant_configonly if the checkpoint actually containslm_head.weight_scale.lm_headback to unquantized whenexclude_modulesexcludeslm_head,tie_word_embeddingsis enabled, orenable_attention_dp+enable_lm_head_tp_in_adpmakes quantized/packed weights incompatible.Mistral3InputProcessor→MistralHFInputProcessor).QA Engineer Review
Test-list changes (CI / QA routing)
tests/integration/test_lists/qa/llm_perf_core.ymlqwen3_30b_a3b-bench-pytorch-bfloat16-maxntfromreqs:256→reqs:64.llama_v3.3_nemotron_super_49bperformance test entries to reflect timeout/interconnect constraints (notably moving a previously presentinput_output_len:8000,1000set and adding it back under a different block with accompanying comments).tests/integration/test_lists/waives.txtaccuracy/test_llm_api_pytorch.py::TestMistralLarge3_675Btest_fp8[latency_moe_deepgemm]andtest_nvfp4_4gpus[...]cases.full:DGX_H100/.../test_linear.py::test_column_linear[2-unbalanced].unittest/_torch/modeling -k "modeling_qwen"and.../test_modeling_qwen3_5_vl.py::test_qwen35_dense_vl_resolves_mamba_ssm_cache_dtype.unittest/_torch/multi_gpu/test_linear.py::test_row_linear[2-unbalanced].Verdict: needs follow-up (test-list coverage data / CBTS linkage for the associated changes is not provided here).
Test code changes
tests/unittest/_torch/modeling/test_modeling_mistral.pytest_processor_get_num_tokens_per_imageto useMistralHFInputProcessor._make_dummy_processorusage to constructMistralHFInputProcessordirectly (affecting the dummy-input contract for downstream assertions intest_mistral_attention_swa_layer_types).tests/unittest/_torch/multi_gpu/test_linear.pytest_column_linearandtest_row_linearto remove conditional failure expectations for oddhidden_size; now asserts all results areTrue.Coverage in
tests/integration/test_lists/: not determinable from provided context (no mapping from these unit tests to specific CI lists included here).Verdict: needs follow-up
Description
Test Coverage
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.