[TRTLLM-9179][feat] add pp_partition to customize each rank's layer number #9003
Conversation
📝 WalkthroughWalkthroughThis pull request introduces support for a new Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant BaseLlmArgs
participant _ParallelConfig
participant Mapping
participant MappingBase
User->>BaseLlmArgs: Initialize with pp_partition
activate BaseLlmArgs
BaseLlmArgs->>_ParallelConfig: Create with pp_partition
activate _ParallelConfig
_ParallelConfig->>_ParallelConfig: Store pp_partition
_ParallelConfig->>BaseLlmArgs: Return config
deactivate _ParallelConfig
BaseLlmArgs->>Mapping: to_mapping() includes pp_partition
activate Mapping
Mapping->>MappingBase: Initialize with pp_partition
activate MappingBase
MappingBase->>MappingBase: Store pp_partition
MappingBase->>Mapping: Ready for pp_layers()
deactivate MappingBase
deactivate Mapping
deactivate BaseLlmArgs
rect rgba(100, 150, 200, 0.1)
note right of MappingBase: pp_partition execution
User->>MappingBase: Call pp_layers(num_layers)
activate MappingBase
alt pp_partition provided
MappingBase->>MappingBase: Validate partition length and sum
MappingBase->>MappingBase: Use partition sizes to split layers
else pp_partition not provided
MappingBase->>MappingBase: Fallback to even distribution (tensor_split)
end
MappingBase->>User: Return layer distribution
deactivate MappingBase
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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. 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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/mapping.py (2)
454-485: Critical: pp_partition not passed to parent class.The
pp_partitionparameter is added to theMapping.__init__signature at line 464, but it is not passed tosuper().__init__()in the call at lines 472-485. This meansMappingBasewill not receive thepp_partitionvalue whenMappingis instantiated.While
Mapping.__new__returnsMpiTopologyorDeviceMeshTopologyinstances (which use*args, **kwargsand would work correctly), the explicit signature here should still passpp_partitionfor correctness and consistency.Apply this diff to add the missing parameter:
super().__init__(world_size=world_size, rank=rank, gpus_per_node=gpus_per_node, cp_size=cp_size, cp_config=cp_config, tp_size=tp_size, pp_size=pp_size, + pp_partition=pp_partition, moe_cluster_size=moe_cluster_size, moe_tp_size=moe_tp_size, moe_ep_size=moe_ep_size, attn_tp_size=attn_tp_size, attn_cp_size=attn_cp_size, enable_attention_dp=enable_attention_dp, enable_lm_head_tp_in_adp=enable_lm_head_tp_in_adp)
332-348: Add pp_partition to the to_dict() method.The
to_dict()method should includepp_partitionin the returned dictionary to support proper serialization and deserialization ofMappingobjects. Without this, any custom partition configuration would be lost when converting to/from dictionary format.Apply this diff:
def to_dict(self): return { 'world_size': self.world_size, 'rank': self.rank, 'gpus_per_node': self.gpus_per_node, 'cp_size': self.cp_size, 'tp_size': self.tp_size, 'pp_size': self.pp_size, + 'pp_partition': self.pp_partition, 'moe_tp_size': self.moe_tp_size, 'moe_cluster_size': self.moe_cluster_size, 'moe_ep_size': self.moe_ep_size, 'attn_tp_size': self.attn_tp_size, 'attn_cp_size': self.attn_cp_size, 'cp_config': self.cp_config, 'enable_attention_dp': self.enable_attention_dp, 'enable_lm_head_tp_in_adp': self.enable_lm_head_tp_in_adp, }
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
docs/source/developer-guide/api-change.md(6 hunks)tensorrt_llm/llmapi/llm_args.py(4 hunks)tensorrt_llm/mapping.py(6 hunks)tests/unittest/api_stability/references/llm.yaml(1 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Use only spaces, no tabs; indent with 4 spaces.
Files:
tensorrt_llm/llmapi/llm_args.pytensorrt_llm/mapping.py
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.
Files:
tensorrt_llm/llmapi/llm_args.pytensorrt_llm/mapping.py
**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).
Files:
tensorrt_llm/llmapi/llm_args.pytensorrt_llm/mapping.py
🧠 Learnings (2)
📓 Common learnings
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:42-49
Timestamp: 2025-09-23T14:58:05.372Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/), the token partitioning intentionally uses ceil-like distribution (same token_per_rank for all ranks) to ensure all ranks launch the same number of blocks. This is required for optimal NCCL device API barrier performance, even though it may launch extra blocks for non-existent tokens on later ranks. Runtime bounds checking in the kernel (blockID validation) handles the overshoot cases.
📚 Learning: 2025-08-26T09:37:10.463Z
Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which can contain default `cuda_graph_config` values, so `llm_args` may already have this config before the extra options processing.
Applied to files:
tensorrt_llm/llmapi/llm_args.py
🧬 Code graph analysis (2)
tensorrt_llm/llmapi/llm_args.py (1)
tensorrt_llm/builder.py (1)
default(45-50)
tensorrt_llm/mapping.py (2)
tensorrt_llm/_torch/distributed/communicator.py (2)
pp_size(59-60)pp_rank(75-76)tensorrt_llm/_torch/device_mesh.py (1)
pp_rank(80-81)
🪛 markdownlint-cli2 (0.18.1)
docs/source/developer-guide/api-change.md
50-50: Emphasis used instead of a heading
(MD036, no-emphasis-as-heading)
58-58: Emphasis used instead of a heading
(MD036, no-emphasis-as-heading)
🪛 Ruff (0.14.3)
tensorrt_llm/mapping.py
308-310: Avoid specifying long messages outside the exception class
(TRY003)
312-313: Avoid specifying long messages outside the exception class
(TRY003)
464-464: Unused method argument: pp_partition
(ARG002)
⏰ 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). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (8)
tests/unittest/api_stability/references/llm.yaml (1)
21-24: LGTM! API reference correctly documents the new parameter.The
pp_partitionparameter is properly documented with the correct type annotation, default value, and prototype status, matching the implementation inllm_args.py.tensorrt_llm/llmapi/llm_args.py (3)
370-383: LGTM! Correct propagation of pp_partition.The
pp_partitionfield is properly passed to theMappingconstructor in theto_mapping()method.
1592-1596: LGTM! Well-documented public API field.The
pp_partitionparameter is properly added toBaseLlmArgswith clear description, appropriate prototype status, and correct type annotation.
1844-1856: LGTM! Proper propagation in model validator.The
pp_partitionfield is correctly passed fromBaseLlmArgsto_ParallelConfigduring validation.tensorrt_llm/mapping.py (4)
42-59: LGTM! Parameter properly added and stored.The
pp_partitionparameter is correctly added toMappingBase.__init__signature and stored as an instance attribute. Validation is appropriately deferred to thepp_layers()method where it's actually used.Also applies to: 130-130
151-166: LGTM! Equality comparison properly includes pp_partition.The
pp_partitionfield is correctly included in the equality check.
168-184: LGTM! Hash function correctly handles pp_partition.The hash implementation properly converts the
pp_partitionlist to a tuple for hashing and handles the None case correctly.
305-319: LGTM! Partition-aware layer distribution logic is correct.The
pp_layers()method properly implements custom partitioning:
- Validates partition length against
pp_size- Validates partition sum against
num_layers- Uses
torch.Tensor.split()with partition sizes to distribute layers- Falls back to even distribution when
pp_partitionis not providedThe implementation correctly handles both custom and automatic partitioning scenarios.
| moe_tp_size: int = -1 | ||
| moe_ep_size: int = -1 | ||
| cp_config: dict = Field(default_factory=dict) | ||
| pp_partition: Optional[List[int]] = Field(default=None) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Add description and status to the Field definition.
The pp_partition field in _ParallelConfig should include a description parameter for documentation purposes and a status parameter to indicate its maturity level (likely "prototype" to match the public API in BaseLlmArgs).
Apply this diff:
- pp_partition: Optional[List[int]] = Field(default=None)
+ pp_partition: Optional[List[int]] = Field(
+ default=None,
+ description="Pipeline parallel partition, a list of each rank's layer number.",
+ status="prototype")🤖 Prompt for AI Agents
In tensorrt_llm/llmapi/llm_args.py around line 329, the pp_partition Field lacks
a description and a status indicating maturity; update its Field(...) call to
include a descriptive description string explaining it holds pipeline-parallel
partition indices (e.g., "List of pipeline-parallel partition indices for model
sharding") and add status="prototype" to match the public API in BaseLlmArgs.
| moe_tp_size: int = -1 | ||
| moe_ep_size: int = -1 | ||
| cp_config: dict = Field(default_factory=dict) | ||
| pp_partition: Optional[List[int]] = Field(default=None) |
There was a problem hiding this comment.
Please give a "description" to this Field to introduce this concept. It will be published to the api reference docs as well.
There was a problem hiding this comment.
The attribution in LLMArgs contains description already. Do we need to add description for _ParallelConfig`?
Superjomn
left a comment
Superjomn
left a comment
There was a problem hiding this comment.
Overall LGTM on the llmapi changes. Please add a "description" to the new pp_partition Field.
dc3671
force-pushed
the
add-pp-partition
branch
from
9fdfb23 to
7e39d29
Compare
|
/bot run --disable-fail-fast |
dc3671
force-pushed
the
add-pp-partition
branch
from
7e39d29 to
5777444
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #23961 [ run ] triggered by Bot. Commit: |
|
PR_Github #23962 [ run ] triggered by Bot. Commit: |
|
PR_Github #23961 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #23996 [ run ] triggered by Bot. Commit: |
|
PR_Github #23962 [ run ] completed with state |
|
PR_Github #23996 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #24089 [ run ] triggered by Bot. Commit: |
|
PR_Github #24089 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #24122 [ run ] triggered by Bot. Commit: |
|
PR_Github #24122 [ run ] completed with state |
dc3671
force-pushed
the
add-pp-partition
branch
from
88ae336 to
e79d5f1
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #24150 [ run ] triggered by Bot. Commit: |
|
PR_Github #24150 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #24189 [ run ] triggered by Bot. Commit: |
|
PR_Github #24189 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #24224 [ run ] triggered by Bot. Commit: |
|
PR_Github #24224 [ run ] completed with state |
Signed-off-by: Zhenhuan Chen <[email protected]>
Signed-off-by: Zhenhuan Chen <[email protected]>
Signed-off-by: Zhenhuan Chen <[email protected]>
dc3671
force-pushed
the
add-pp-partition
branch
from
e79d5f1 to
11821a7
Compare
|
/bot run --disable-fail-fast |
|
/bot run --stage-list "DGX_H100-2_GPUs-PyTorch-Others-1" |
|
PR_Github #24292 [ run ] triggered by Bot. Commit: |
|
PR_Github #24294 [ run ] triggered by Bot. Commit: |
|
PR_Github #24292 [ run ] completed with state |
7 participants
Summary by CodeRabbit
New Features
pp_partitionparameter to the LLM API enabling custom pipeline parallelism configuration. Users can now specify precise layer distribution strategies across pipeline stages for optimized performance in distributed LLM inference deployments.Documentation
Description
bs1, [16, 16, 15, 14+MTP3]:
bs1, [16, 15, 15, 15+MTP3]:
bs1, [16, 16, 16, 13+MTP3]:
Can improve from 122ms -> ~106ms
Test Coverage
add test to existed PP tests.
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
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.