fix: Add default configs in LLMAPI. Fixes OOM issues #2198
Conversation
tanmayv25
requested review from
a team,
GuanLuo,
PeaBrane,
alec-flowers,
biswapanda,
grahamking,
ishandhanani,
jthomson04,
kkranen,
nnshah1,
paulhendricks,
piotrm-nvidia,
ptarasiewiczNV,
rmccorm4,
ryanolson,
tedzhouhk and
tmonty12
as code owners
WalkthroughThe changes introduce enhanced configuration options for the TensorRT LLM backend, including dynamic GPU detection, new parallelism parameters, and resource-related settings. Command-line argument parsing and the configuration class are updated to support these options, and the engine initialization logic is refactored to incorporate the new configuration objects and parameters. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15 minutes Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![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 (1)
components/backends/trtllm/src/dynamo/trtllm/utils/trtllm_utils.py (1)
127-140: Inconsistent help text capitalization.The help text for
--expert-parallel-sizeuses lowercase "expert" while other similar parameters use proper capitalization.- help="expert parallelism size.", + help="Expert parallelism size.",
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
components/backends/trtllm/src/dynamo/trtllm/main.py(2 hunks)components/backends/trtllm/src/dynamo/trtllm/utils/trtllm_utils.py(6 hunks)
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: ptarasiewiczNV
PR: ai-dynamo/dynamo#2027
File: container/deps/vllm/install_vllm.sh:0-0
Timestamp: 2025-07-22T10:22:28.972Z
Learning: The `--torch-backend=auto` flag works with vLLM installations via uv pip install, even though it's not a standard pip option. This flag is processed by vLLM's build system during installation to automatically match PyTorch distribution with container CUDA versions.
Learnt from: tanmayv25
PR: ai-dynamo/dynamo#1391
File: examples/tensorrt_llm/common/base_engine.py:171-176
Timestamp: 2025-06-05T01:10:51.865Z
Learning: In examples/tensorrt_llm/common/base_engine.py, the _init_engine method is called only once during initialization, so direct mutation of the _default_sampling_params object during setup is safe and appropriate.
components/backends/trtllm/src/dynamo/trtllm/main.py (2)
Learnt from: tanmayv25
PR: #1391
File: examples/tensorrt_llm/common/base_engine.py:171-176
Timestamp: 2025-06-05T01:10:51.865Z
Learning: In examples/tensorrt_llm/common/base_engine.py, the _init_engine method is called only once during initialization, so direct mutation of the _default_sampling_params object during setup is safe and appropriate.
Learnt from: ptarasiewiczNV
PR: #2027
File: container/deps/vllm/install_vllm.sh:0-0
Timestamp: 2025-07-22T10:22:28.972Z
Learning: The --torch-backend=auto flag works with vLLM installations via uv pip install, even though it's not a standard pip option. This flag is processed by vLLM's build system during installation to automatically match PyTorch distribution with container CUDA versions.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: pre-merge-rust (lib/runtime/examples)
- GitHub Check: Build and Test - vllm
- GitHub Check: pre-merge-rust (.)
- GitHub Check: pre-merge-rust (lib/bindings/python)
🔇 Additional comments (7)
components/backends/trtllm/src/dynamo/trtllm/main.py (4)
11-17: LGTM! Well-organized imports for new configuration classes.The imports are properly structured and include all necessary TensorRT LLM API configuration classes and the torch CUDA device counting function.
Also applies to: 20-20
95-101: Good defensive programming with GPU detection.The dynamic GPU detection logic properly handles the case where
gpus_per_nodeis not explicitly provided, with appropriate error handling when no GPUs are found.
102-122: Confirm DynamicBatchConfig values for OOM preventionPlease verify that these hardcoded values align with your memory-safety goals and workload characteristics in components/backends/trtllm/src/dynamo/trtllm/main.py (lines 102–122):
- enable_batch_size_tuning=True
• Adapts batch size at runtime—ensure it never exceeds available GPU memory under peak load.- enable_max_num_tokens_tuning=False
• Keeps the maximum token count fixed; consider enabling if dynamic token trimming could reduce padding overhead.- dynamic_batch_moving_average_window=128
• The default moving-average window is often 64; confirm that 128 strikes the right balance between responsiveness and stability for your models.- capacity_scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT
• Guarantees no in-flight eviction (conservative memory usage) but may impact throughput; validate this trade-off against MAX_UTILIZATION if higher throughput is needed.Monitor GPU memory (e.g., with
nvidia-smi) under representative workloads and adjust these parameters as needed.
123-139: Double-check overlapping argument mappings forbuild_config/kv_cache_configvs. individual parametersIt looks like you’re passing both individual settings (
max_num_tokens,max_seq_len, etc.) and entire config objects (build_config,kv_cache_config) intoarg_map. Please verify that:
- Neither
build_confignorkv_cache_configalready include the same fields (e.g., max sequence length, batch size, etc.).- There are no silent overrides or conflicts when
arg_mapis later merged withconfig.extra_engine_argsviaupdate_llm_args_with_extra_options.- The final engine invocation receives each parameter exactly once, with the intended precedence.
Key location:
- File: components/backends/trtllm/src/dynamo/trtllm/main.py
- Lines: 123–139 (where
arg_mapis defined)If the config objects already encapsulate these values, consider removing the duplicates or consolidating into one source of truth to avoid confusion.
components/backends/trtllm/src/dynamo/trtllm/utils/trtllm_utils.py (3)
7-7: LGTM! Appropriate import for BuildConfig constants.The import enables using BuildConfig default values for the new configuration parameters.
32-41: Well-structured configuration fields with sensible defaults.The new configuration fields are properly typed and use appropriate defaults from BuildConfig constants, which should help prevent OOM issues by providing tested default values.
153-189: Comprehensive resource configuration options.The new command-line arguments provide good coverage of resource-related parameters that should help prevent OOM issues. The defaults using BuildConfig constants are appropriate.
components/backends/trtllm/src/dynamo/trtllm/utils/trtllm_utils.py
Outdated
Show resolved
Hide resolved
|
In general, I want to ask why the file structure is the way it is. We have |
|
@nv-kmcgill53 See the discussion here. It is to support launch like |
Overview:
These defaults are picked from trtllm-serve stack. Without these optimizations, larger models like DS R1 can run into OOM.
See here for more information: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/commands/serve.py#L94-L149
Summary by CodeRabbit
New Features
Refactor