[https://nvbugs/5448464][fix] Partially fix LoRA overallocation for Nemotron NAS#12817
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThis PR improves LoRA PEFT cache sizing calculation in TensorRT-LLM by introducing a new Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 2
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/_util.py (2)
1574-1578: Replace Unicode multiplication sign with ASCII 'x'.The Unicode
×(MULTIPLICATION SIGN, U+00D7) may cause issues with log parsers or terminals. Use ASCIIxinstead.Suggested fix
logger.info(f"LoRA module-layer count: {total} " - f"(attn: {layers_with_attn}×{len(attn_modules)}, " - f"mlp: {layers_with_mlp}×{len(mlp_modules)}, " - f"other: {num_layers}×{len(other_modules)}, " + f"(attn: {layers_with_attn}x{len(attn_modules)}, " + f"mlp: {layers_with_mlp}x{len(mlp_modules)}, " + f"other: {num_layers}x{len(other_modules)}, " f"uniform would be {num_layers * len(all_target_modules)})")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 1574 - 1578, Log message uses Unicode multiplication sign '×'; update the formatted string in the logger.info call inside _util.py to replace every '×' with ASCII 'x' so log parsers/terminals won't break. Edit the logger.info(...) invocation (referenced by the variables total, layers_with_attn, attn_modules, layers_with_mlp, mlp_modules, num_layers, other_modules, all_target_modules) and change the three occurrences of '×' to 'x' in the f-string.
1538-1550: Add return type annotation.The function should include a return type hint for consistency with coding guidelines.
Suggested fix
-def _compute_num_lora_modules(pretrained_config, all_target_modules): +def _compute_num_lora_modules(pretrained_config, all_target_modules: list[str]) -> int:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 1538 - 1550, The function _compute_num_lora_modules lacks a return type annotation; update its signature to include an explicit return type (e.g., change def _compute_num_lora_modules(pretrained_config, all_target_modules): to def _compute_num_lora_modules(pretrained_config, all_target_modules) -> int:) and ensure all return paths (like the existing return num_layers * len(all_target_modules)) are consistent with that type.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 1559-1568: The code currently inspects nonexistent attributes on
AttentionConfig/FFNConfig (causing AttributeError); update the helper functions
_has_lora_capable_attn and _has_lora_capable_ffn to use the DeciLayerConfig
boolean properties (use bc.is_attention_layer in place of the old
attention.no_op/replace_with_linear checks and bc.is_mlp_layer in place of
ffn.no_op/replace_with_linear), and then compute layers_with_attn and
layers_with_mlp using those updated helpers so the logic works with Nemotron-NAS
block configs.
In `@tests/unittest/others/test_lora_module_count.py`:
- Around line 29-37: The function _compute_num_lora_modules is accessing
bc.attention.no_op and bc.attention.replace_with_linear (and similar ffn fields)
but DeciLayerConfig exposes attention.impl and ffn.impl as enum values; change
_compute_num_lora_modules to inspect bc.attention.impl and bc.ffn.impl and
compare them to the appropriate AttentionImplementation enum members (e.g.,
AttentionImplementation.NO_OP, AttentionImplementation.LINEAR or the equivalent
enum names used in the code) when deciding whether to count LoRA modules, and
remove/replace any references to no_op/replace_with_linear so real
DeciLayerConfig instances do not raise AttributeError. Ensure any logic
branching that used the old booleans maps to the same behavior using the enum
comparisons.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 1574-1578: Log message uses Unicode multiplication sign '×';
update the formatted string in the logger.info call inside _util.py to replace
every '×' with ASCII 'x' so log parsers/terminals won't break. Edit the
logger.info(...) invocation (referenced by the variables total,
layers_with_attn, attn_modules, layers_with_mlp, mlp_modules, num_layers,
other_modules, all_target_modules) and change the three occurrences of '×' to
'x' in the f-string.
- Around line 1538-1550: The function _compute_num_lora_modules lacks a return
type annotation; update its signature to include an explicit return type (e.g.,
change def _compute_num_lora_modules(pretrained_config, all_target_modules): to
def _compute_num_lora_modules(pretrained_config, all_target_modules) -> int:)
and ensure all return paths (like the existing return num_layers *
len(all_target_modules)) are consistent with that type.
🪄 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: Pro
Run ID: 6bf88908-12c6-4cbf-aa72-2c9050e6bc69
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/_util.pytests/integration/test_lists/test-db/l0_a10.ymltests/unittest/llmapi/test_llm_pytorch.pytests/unittest/others/test_lora_module_count.py
💤 Files with no reviewable changes (1)
- tests/unittest/llmapi/test_llm_pytorch.py
|
PR_Github #42191 [ run ] triggered by Bot. Commit: |
15341ad to
d0be243
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #42193 [ run ] triggered by Bot. Commit: |
brb-nv
changed the title
d0be243 to
d43004f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #42204 [ run ] triggered by Bot. Commit: |
|
PR_Github #42204 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #42337 [ run ] triggered by Bot. Commit: |
|
PR_Github #42337 [ run ] completed with state |
…Nemotron NAS Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
d43004f to
13c56d1
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #42380 [ run ] triggered by Bot. Commit: |
|
PR_Github #42380 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #42487 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #42508 [ run ] triggered by Bot. Commit: |
|
PR_Github #42487 [ run ] completed with state |
|
PR_Github #42508 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #42573 [ run ] triggered by Bot. Commit: |
|
PR_Github #42573 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #42607 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #42696 [ run ] triggered by Bot. Commit: |
|
PR_Github #42696 [ run ] completed with state |
3 participants
Description
This is to fix https://nvbugspro.nvidia.com/bug/5448464.
Background:
num_device_module_layer, which is computed in Python asmax_lora_rank×num_lora_modules×max_loras. On Nemotron-Super-49B with LoRA rank 64, this cache requests ~2.46 GiB but only ~2.44 GiB remains after model + KV cache allocation, causingcudaMallocAsyncOOM during executor creation.Root cause
There are two forms of LoRA cache over-allocation:
Inflated page size:
_infer_nemotron_ffn_multusesmax(ffn_mult)across all layers as a single uniformmlp_hidden_size. This inflates LoraModule dimensions and therefore cache page size, since the C++ API requires uniform module dimensions across layers. Fixing this requires C++ changes to support per-layer LoraModule dimensions (tracked by TRTLLM-5045).Inflated module count:
num_lora_modulesis computed asnum_hidden_layers×len(target_modules), assuming every layer has every module type. For Nemotron-NAS, this allocates cache slots for MLP LoRA on layers with no MLP (no_oporreplace_with_linear), and attention LoRA on layers with no attention. These layers don't create those submodules in the PyTorch model, and LoRA adapters can't have weights for them — the slots are pure waste.Fix:
_compute_num_lora_modules()which, for models with block_configs, inspects each layer to determine whether it has LoRA-capable attention and/or MLP. Target modules are classified as attention, MLP, or other, and only counted for layers that actually have that module type.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)
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.Summary by CodeRabbit
Bug Fixes
Tests