[#11548][feat] AutoDeploy: Optimize Qwen3.5 perf#12265
Conversation
taylor-yb-lee
changed the title
📝 WalkthroughWalkthroughThis PR adds tensor parallelism-aware MoE scale sharding support to the auto-deploy framework. Changes include updating the Qwen 3.5 MoE 400B configuration with new MoE transforms blocks and parameter adjustments, introducing new NVFP4 and FineGrained FP8 scale sharding helpers, and enhancing lm_head node handling for column-sharded architectures. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 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.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py (2)
1671-1691: Add parameter documentation to the docstring.The function docstring explains the purpose well but lacks parameter documentation, which would help future maintainers understand the expected inputs (especially
orig_weight_shapewhich represents the uint8-packed weight shape before any sharding).📝 Suggested docstring enhancement
def _shard_nvfp4_moe_scale( scale: torch.Tensor, orig_weight_shape: torch.Size, dim: int, rank: int, world_size: int, ) -> torch.Tensor: """Shard NVFP4 weight_scale for MoE TP, preserving 2D cutlass format. Unlike _shard_fp4_weight_scale (which returns 1D), this returns a 2D tensor with the correct padded shape, matching the format expected by MoE stacking. + + Args: + scale: The weight_scale tensor in cutlass format (1D or 2D). + orig_weight_shape: Shape of the corresponding uint8-packed weight tensor + before TP sharding (used for format conversion). + dim: Dimension along which to split (0 for column, 1 for row). + rank: Current TP rank. + world_size: Total TP world size. + + Returns: + 2D tensor with padded shape (m + pad_m, n + pad_n) in cutlass format. """🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py` around lines 1671 - 1691, Update the _shard_nvfp4_moe_scale docstring to include parameter and return documentation: describe each parameter (scale: torch.Tensor of per-element scales; orig_weight_shape: torch.Size representing the uint8-packed weight shape prior to sharding—note last dim was doubled to convert from bytes to elements; dim: int tensor split dimension; rank: int current TP rank; world_size: int total TP ranks) and the return value (torch.Tensor shaped to the padded 2D cutlass format matching MoE stacking). Also mention padding behavior (pads m to multiple of 128 and n to multiple of 4) so maintainers know why reshape uses those pads.
1846-1865: Index arithmetic relies on implicit arg layout.The scale index calculation
6 + s_idx * 3depends on the implicit structure ofnode.args. While the comments explain the layout, this coupling between the function and op schema is fragile. Consider extracting the base index as a constant or usingextract_op_args/set_op_argsfor consistency with the pattern used elsewhere (e.g., line 1922).♻️ Optional: extract base index as constant
+ # Scales start at arg index 6 (after: hidden_states, selected_experts, + # final_scales, w_up, w_down, w_gate) + _SCALE_ARGS_BASE_IDX = 6 + if tp_size > 1 and scale_names: _BLOCKED_SCALE_NAMES = {"weight_scale", "weight_scale_inv"} for s_idx, s_name in enumerate(scale_names): if s_name not in _BLOCKED_SCALE_NAMES: continue - # For each scale_name, the 3 lists correspond to w_up, w_down, w_gate - # w_up/w_gate use COLUMN split (dim=0), w_down uses ROW split (dim=1) scale_dim_groups = [ - (6 + s_idx * 3 + 0, SplitDimension.COLUMN, w_up_orig_shapes), - (6 + s_idx * 3 + 1, SplitDimension.ROW, w_down_orig_shapes), - (6 + s_idx * 3 + 2, SplitDimension.COLUMN, w_gate_orig_shapes), + (_SCALE_ARGS_BASE_IDX + s_idx * 3 + 0, SplitDimension.COLUMN, w_up_orig_shapes), # w_up + (_SCALE_ARGS_BASE_IDX + s_idx * 3 + 1, SplitDimension.ROW, w_down_orig_shapes), # w_down + (_SCALE_ARGS_BASE_IDX + s_idx * 3 + 2, SplitDimension.COLUMN, w_gate_orig_shapes), # w_gate ]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py` around lines 1846 - 1865, The index math for scale args is fragile: replace the repeated magic expression (6 + s_idx * 3) with a clearly named base index and/or use the existing arg helper functions; specifically, inside the loop over scale_names compute a single base_index = 6 + s_idx * 3 (or obtain the named args via extract_op_args/set_op_args) and then build scale_dim_groups using base_index + 0/1/2, keeping references to _BLOCKED_SCALE_NAMES, args, scale_dim_groups and the call to _tp_shard_moe_scale unchanged so the rest of the logic is identical but the indexing is explicit and maintainable.tensorrt_llm/_torch/auto_deploy/utils/_graph.py (1)
574-582: Consider handling both all_gather backends for consistency.The unwrapping logic only handles
trtllm_dist_all_gather, but the codebase supports bothtrtllm_dist_all_gatherandtorch_dist_all_gatherbackends (see_get_dist_opsin sharding.py). If the torch distributed backend is used for lm_head sharding, this unwrapping won't trigger.♻️ Proposed fix to handle both backends
# Unwrap all_gather for sharded lm_head: when lm_head weight is column- # sharded the graph contains lm_head_linear -> all_gather -> output. # We look through the all_gather so that callers (e.g. # gather_logits_before_lm_head) see the underlying linear and can insert # gather_tokens *before* the sharded GEMM + all_gather, keeping both out # of the main CUDA graph and avoiding NVLink contention with layer # AllReduces. - if is_op(lm_head_node, torch.ops.auto_deploy.trtllm_dist_all_gather): + if is_op(lm_head_node, [ + torch.ops.auto_deploy.trtllm_dist_all_gather, + torch.ops.auto_deploy.torch_dist_all_gather, + ]): lm_head_node = lm_head_node.all_input_nodes[0]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/utils/_graph.py` around lines 574 - 582, The unwrap currently only checks for torch.ops.auto_deploy.trtllm_dist_all_gather; update the condition that tests lm_head_node to also detect the torch distributed all_gather backend used elsewhere in the repo (the op registered as torch_dist_all_gather) so sharded lm_head -> all_gather is handled regardless of backend. Concretely, modify the check around lm_head_node (the block that calls is_op(lm_head_node, torch.ops.auto_deploy.trtllm_dist_all_gather) and then reassigns lm_head_node = lm_head_node.all_input_nodes[0]) to test for both torch.ops.auto_deploy.trtllm_dist_all_gather and the torch ops variant (or obtain the set from _get_dist_ops in sharding.py) and unwrap when either matches.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py`:
- Around line 1671-1691: Update the _shard_nvfp4_moe_scale docstring to include
parameter and return documentation: describe each parameter (scale: torch.Tensor
of per-element scales; orig_weight_shape: torch.Size representing the
uint8-packed weight shape prior to sharding—note last dim was doubled to convert
from bytes to elements; dim: int tensor split dimension; rank: int current TP
rank; world_size: int total TP ranks) and the return value (torch.Tensor shaped
to the padded 2D cutlass format matching MoE stacking). Also mention padding
behavior (pads m to multiple of 128 and n to multiple of 4) so maintainers know
why reshape uses those pads.
- Around line 1846-1865: The index math for scale args is fragile: replace the
repeated magic expression (6 + s_idx * 3) with a clearly named base index and/or
use the existing arg helper functions; specifically, inside the loop over
scale_names compute a single base_index = 6 + s_idx * 3 (or obtain the named
args via extract_op_args/set_op_args) and then build scale_dim_groups using
base_index + 0/1/2, keeping references to _BLOCKED_SCALE_NAMES, args,
scale_dim_groups and the call to _tp_shard_moe_scale unchanged so the rest of
the logic is identical but the indexing is explicit and maintainable.
In `@tensorrt_llm/_torch/auto_deploy/utils/_graph.py`:
- Around line 574-582: The unwrap currently only checks for
torch.ops.auto_deploy.trtllm_dist_all_gather; update the condition that tests
lm_head_node to also detect the torch distributed all_gather backend used
elsewhere in the repo (the op registered as torch_dist_all_gather) so sharded
lm_head -> all_gather is handled regardless of backend. Concretely, modify the
check around lm_head_node (the block that calls is_op(lm_head_node,
torch.ops.auto_deploy.trtllm_dist_all_gather) and then reassigns lm_head_node =
lm_head_node.all_input_nodes[0]) to test for both
torch.ops.auto_deploy.trtllm_dist_all_gather and the torch ops variant (or
obtain the set from _get_dist_ops in sharding.py) and unwrap when either
matches.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 119f372a-2765-40bb-8ff6-cc5145e00132
📒 Files selected for processing (3)
examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yamltensorrt_llm/_torch/auto_deploy/transform/library/sharding.pytensorrt_llm/_torch/auto_deploy/utils/_graph.py
nvchenghaoz
left a comment
nvchenghaoz
left a comment
There was a problem hiding this comment.
Shall we add the test coverage for new functions?
be0f385 to
e553c79
Compare
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse 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. |
7a9606f to
3b16a43
Compare
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
17aacaa to
3c1e8b0
Compare
Added new unittest for NVFP4 MoE TP sharding |
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #40085 [ run ] triggered by Bot. Commit: |
|
PR_Github #40087 [ run ] triggered by Bot. Commit: |
|
PR_Github #40085 [ run ] completed with state |
|
PR_Github #40087 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #42060 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #42085 [ run ] triggered by Bot. Commit: |
|
PR_Github #42085 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #42103 [ run ] triggered by Bot. Commit: |
|
PR_Github #42103 [ run ] completed with state
|
|
/bot run --disable-multi-gpu-test |
|
PR_Github #42177 [ run ] triggered by Bot. Commit: |
|
PR_Github #42177 [ run ] completed with state
|
|
/bot run --disable-multi-gpu-test |
|
PR_Github #42185 [ run ] triggered by Bot. Commit: |
|
PR_Github #42185 [ run ] completed with state
|
|
/bot run --disable-multi-gpu-test --reuse-test |
|
PR_Github #42267 [ run ] triggered by Bot. Commit: |
|
PR_Github #42267 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #42371 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast --disable-multi-gpu-test --reuse-test |
|
PR_Github #42371 [ run ] completed with state
|
|
PR_Github #42401 [ run ] triggered by Bot. Commit: |
|
PR_Github #42401 [ run ] completed with state |
|
/bot run --disable-fail-fast --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #42466 [ run ] triggered by Bot. Commit: |
|
PR_Github #42466 [ run ] completed with state |
|
/bot skip --comment "my change is only about AutoDeploy. I ran the relevant tests separately: Single GPU only : https://prod.blsm.nvidia.com/sw-tensorrt-top-1/job/LLM/job/main/job/L0_MergeRequest_PR/33174/ |
|
PR_Github #42560 [ skip ] triggered by Bot. Commit: |
|
PR_Github #42560 [ skip ] completed with state |
5 participants
Description
This PR enables lm_head sharding for Qwen3_5MoeModel and adds TP8 MoE support for Qwen3.5 with NVFP4 quantization.
Key changes:
Known Limitations
As of #12114, the default model factory for Qwen3.5 has been changed to Qwen3_5MoeForConditionalGeneration to support both VLM and text-only modes. However, this factory currently has several outstanding issues:
Recommendation: Until these issues are resolved, use AutoModelForCausalLM for text-only Qwen3.5 workloads.
(See the comments in examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml)
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.