Repository navigation
[https://nvbugs/6891030][fix] Default mamba_ssm_cache_dtype to None so "auto" follows the checkpoint dtype… - #20051
Conversation
…he dtype to float32 NemotronHConfig (registered over CONFIG_MAPPING["nemotron_h"] for checkpoints whose hybrid_override_pattern contains MLP layers) defaulted mamba_ssm_cache_dtype to "float32". Dense checkpoints such as NVIDIA-Nemotron-Nano-12B-v2 do not declare the field, so kv_cache_config.mamba_ssm_cache_dtype="auto" resolved to an fp32 recurrent-state cache instead of the checkpoint dtype, doubling the state memory. The V2 Mamba cache manager then rejected the GPU quota at max_batch_size=512 on H100. Leave the field unset unless the checkpoint declares it, restoring the resolution these checkpoints had through their remote config. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (2)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughNemotron-H now leaves ChangesNemotron-H cache dtype resolution
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to The cache dtype fix and regression coverage are ready to merge with no identified current-head risk. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Stamp on behalf of runtime devs, delegating proper review to @NVIDIA/trt-llm-devs; please ping me if you think this is not accurate
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
This fixes the regression, but only inside the TRT-LLM shim. The native transformers NemotronHConfig (5.5.4, configuration_nemotron_h.py:121) also defaults mamba_ssm_cache_dtype to "float32", and NemotronHConfig.from_dict hands every natively representable pattern to that class ([tensorrt_llm/_torch/configs/nemotron_h.py:106](https://github.com/NVIDIA/TensorRT-LLM/pull/20051/files#diff-4e82e7f8e39268e5ed3b250721baa962a54644600e0e28b6bb698f2b6e8ab1eeR106)-111). Two consequences:
- Once
requirements.txtmoves past transformers 5.6.2, which fixes the-pattern upstream, dense checkpoints will load through the native class and"auto"will go back to fp32. Theisinstance(cfg, NemotronHConfig)assertion in the new test will fail at that point, which is good, but the fix itself won't carry over. A more durable fix is to makeresolve_ssm_cache_dtypeignore a value that only comes from the class default (for example, compare againstconfig.to_diff_dict(), or check whether the key was in the raw config dict). That can be a follow-up, but please at least add a TODO next to the default that links it to the transformers pin. - The class docstring (
[nemotron_h.py:79](https://github.com/NVIDIA/TensorRT-LLM/pull/20051/files#diff-4e82e7f8e39268e5ed3b250721baa962a54644600e0e28b6bb698f2b6e8ab1eeR79)-82) says the shim keeps "exact current behaviour" relative to native. That's no longer true for this field, so please update it.
Please also add the post-fix result of the repro (the two nemotron_nano_12b_v2 bench cases on H100) to the description. Right now the only signature shown is the failing ToT run, so the "verified on the same GPU" checkbox has no evidence behind it.
| # do not, and kv_cache_config.mamba_ssm_cache_dtype="auto" must then | ||
| # resolve to the checkpoint dtype rather than an implicit fp32 state | ||
| # cache (twice the recurrent-state memory). | ||
| mamba_ssm_cache_dtype=None, |
There was a problem hiding this comment.
Native transformers NemotronHConfig defaults this to "float32", so this line makes the shim deliberately differ from the native class. Please say so in the comment and add a TODO: when the transformers pin reaches ≥5.6.2, dense checkpoints stop using this class and the fp32 default comes back unless resolve_ssm_cache_dtype learns to ignore a class-default value. The current comment also explains the motivation ("twice the recurrent-state memory"). That belongs in the PR description; the comment should just state the contract: None means not declared by the checkpoint, so "auto" falls back to torch_dtype.
| ids=["undeclared", "declared_fp32"], | ||
| ) | ||
| def test_nemotron_h_dense_ssm_cache_dtype(tmp_path, declared, expected): | ||
| # Dense Nemotron-H checkpoints (e.g. NVIDIA-Nemotron-Nano-12B-v2) do not |
There was a problem hiding this comment.
Nit: the H100 / max_batch_size=512 detail describes the incident, not the behavior under test, and it will go stale. Something like "undeclared → checkpoint dtype; declared value wins" is enough here. The incident context fits better in the PR description.
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
|
NVBug 6891030 is closed as Bug - Fix unknown. The linked bug appears resolved elsewhere or for a reason that does not prove this PR is redundant. This PR should be judged on its own merits; repair-bot is not auto-closing it. |
|
#20284 merged this morning and makes the same change ( |
Summary
bash .repair-bot/reproduce/6891030-0c4d6fa6cb0d-1791533546825211065/tot/repro_driver.shTest plan
Links
Reproduction comparison
Signature: ========== 2 failed, 3860 deselected, 6 warnings in (0:06:27) ==========
Dev Engineer Review
NemotronHConfignow defaultsmamba_ssm_cache_dtypetoNone. For checkpoints without a declared value,"auto"can resolve to the checkpoint dtype instead of implicitly selecting FP32.QA Engineer Review
tests/unittest/_torch/test_custom_config_registration.py. The parameterizedtest_nemotron_h_dense_ssm_cache_dtypechecks behavior when the dtype is absent and when it is explicitly set tofloat32.tests/integration/test_lists/test-db/l0_cpu.yml.Per-File QA Perspective
tensorrt_llm/_torch/configs/nemotron_h.py: Verify that a dense Nemotron-H checkpoint withoutmamba_ssm_cache_dtypeleaves the field unset, and that"auto"resolves from the checkpoint dtype. Also verify that explicit checkpoint values continue to take precedence.tests/unittest/_torch/test_custom_config_registration.py: Covers automatic cache-dtype resolution for absent and explicitly declared values. The test is listed in the CPU CI test list attests/integration/test_lists/test-db/l0_cpu.yml.