[None][feat] AutoDeploy: Qwen3.5: Apply whielist based sharding and apply lm_head sharding - #15185
Conversation
18285ce to
897fa1a
Compare
taylor-yb-lee
changed the title
taylor-yb-lee
changed the title
taylor-yb-lee
changed the title
📝 WalkthroughWalkthroughThis PR extends TensorRT-LLM's tensor-parallel sharding system for Qwen 3.5 MoE by introducing a keyword-based simple-shard filter, adding NVFP4 expert weight-scale sharding, simplifying SwiGLU pattern matching, and updating shared-expert layer-type handling in the model. ChangesQwen 3.5 MoE IR Sharding & Config
🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 (1)
tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py (1)
734-745: ⚡ Quick winAdd
strict=Truetozip()for defensive programming.The
zip()calls pair EP-partitioned scale nodes with their original weight shapes. While both lists should have identical lengths (both derived from the same EP partition viaget_partition), addingstrict=Trueprovides fail-fast behavior if a future change introduces a mismatch—rather than silently truncating and producing incorrect sharding.♻️ Proposed fix
- for sn, osh in zip(w1_ws_local, w1_orig): + for sn, osh in zip(w1_ws_local, w1_orig, strict=True): _tp_shard_moe_scale( gm, sn, "weight_scale", SplitDimension.COLUMN, tp_rank, tp_size, osh ) - for sn, osh in zip(w3_ws_local, w3_orig): + for sn, osh in zip(w3_ws_local, w3_orig, strict=True): _tp_shard_moe_scale( gm, sn, "weight_scale", SplitDimension.COLUMN, tp_rank, tp_size, osh ) - for sn, osh in zip(w2_ws_local, w2_orig): + for sn, osh in zip(w2_ws_local, w2_orig, strict=True): _tp_shard_moe_scale( gm, sn, "weight_scale", SplitDimension.ROW, tp_rank, tp_size, osh )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py` around lines 734 - 745, The three zip iterations pairing EP-partitioned scale nodes with original weight shapes (the loops using w1_ws_local/w1_orig, w3_ws_local/w3_orig, and w2_ws_local/w2_orig that call _tp_shard_moe_scale) should use zip(..., strict=True) to fail fast on length mismatches; update each zip invocation to pass strict=True so any future mismatch between the partitioned lists and original shapes raises an immediate error rather than silently truncating.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py`:
- Around line 734-745: The three zip iterations pairing EP-partitioned scale
nodes with original weight shapes (the loops using w1_ws_local/w1_orig,
w3_ws_local/w3_orig, and w2_ws_local/w2_orig that call _tp_shard_moe_scale)
should use zip(..., strict=True) to fail fast on length mismatches; update each
zip invocation to pass strict=True so any future mismatch between the
partitioned lists and original shapes raises an immediate error rather than
silently truncating.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d642dbe1-4a65-406a-a93e-534b028d8c34
📒 Files selected for processing (4)
examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yamltensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.pytensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py
897fa1a to
d93311e
Compare
|
/bot run |
|
PR_Github #53228 [ run ] triggered by Bot. Commit: |
|
PR_Github #53228 [ run ] completed with state |
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" |
|
PR_Github #53401 [ run ] triggered by Bot. Commit: |
|
PR_Github #53401 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" |
|
PR_Github #53476 [ run ] triggered by Bot. Commit: |
|
PR_Github #53476 [ run ] completed with state |
|
/bot run |
|
PR_Github #53621 [ run ] triggered by Bot. Commit: |
|
PR_Github #53621 [ run ] completed with state
|
|
/bot run |
|
PR_Github #53661 [ run ] triggered by Bot. Commit: |
d93311e to
bc2698a
Compare
|
PR_Github #53661 [ run ] completed with state |
|
/bot run |
|
PR_Github #53681 [ run ] triggered by Bot. Commit: |
|
PR_Github #53681 [ run ] completed with state
|
- Upstream pattern-matcher now auto-stamps layer_type, so the explicit hooks are dead.
- Drop modeling's layer_type="shared_expert" tag + fuse_swiglu's preserve_layer_types.
- No functional change: shared expert stays replicated (untagged -> excluded by shard_layers).
shared-expert linear: layer_type="shared_expert" --> untagged ("unknown")
excluded by shard_layers=["moe","delta","mha"] either way -> replicated
preserve_layer_types ctx-manager (3 sites) --> removed (hint auto-stamped)
Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
- simple_shard_filter + _simple_shard_node gather-shard lm_head (column split + all_gather) instead of replicating the 248320-vocab projection (~3% tput loss). - yaml: enable the hint-driven IR sharder + simple_shard_filter: "lm_head". lm_head: replicated (V x H per rank) --> V/tp x H + all_gather Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
- MoEShardableNode TP-splits NVFP4 per-expert weight_scale (via _tp_shard_moe_scale) to follow the weight split (w1/w3 COLUMN, w2 ROW), so moe_tp>1 works for NVFP4 (was crashing on gemm2_weights_scale dim). moe (TP): weight /tp, scale full (crash) --> weight AND scale both /tp Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
MoE TP is not used for this deployment (EP-only); revert to keep the PR to a single concern. The reverted commit stays in history for later cherry-pick. Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
- New AttentionSinksShardableNode splits torch_attention's per-head `sinks` param (dim 0, /tp) to follow the q-head column split; the IR sharder previously left it full, so rank>0 used the wrong per-head sinks. - Gating reuses the apply loop: attention-DP skips it (attention replicated), shard_layers gates via the node's layer_type="mha". No-op when no sinks. - Result: gpt-oss-120b TP2 GSM8K 88.48 -> 90.14 (full 1319, ref 90.30). Reference torch_attention crashed before; the trtllm kernel tolerated full sinks but silently mis-applied them on rank>0. q/k/v heads: /tp (LinearShardableNode) + head-shrink (ViewShardableNode) -- already sharded sinks[num_heads]: full on every rank --> sinks[num_heads/tp] (head-aligned, this commit) Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
bc2698a to
05cffdd
Compare
- A row-parallel linear's bias was left full on every rank, so the trailing all_reduce summed it world_size times. Keep it on rank 0 (zero elsewhere) so the all_reduce contributes it exactly once. - Needs a dedicated always-apply load hook: the shape-gated _load_hook only transforms when the shape changes, so it skips a same-shape zeroing. - Surfaced by gpt-oss attention bias; fp32 sharding-IR equivalence 0.29 -> 0.000. row-parallel out + full bias (every rank) --all_reduce--> out + bias*world_size (was wrong) row-parallel out + bias on rank 0 only --all_reduce--> out + bias (this fix) Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
- _has_ir_markers matched the OpOverloadPacket, but torch.export emits the .default overload -> sharding was skipped for every IR family, so the test was a silent no-op (compared unsharded vs unsharded). Match both forms. - Compare in fp32: sharding is dtype-independent math, and bf16 noise (~0.05 on attention-heavy models) overlapped the sabotage band. - Per-family shard_layers (qwen3_5_moe, gpt_oss): default None shards replicate-only weights (shared expert / lm_head) and corrupts the output. - Result: 26 passed / 12 skipped / 0 failed (previously validated nothing). Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
05cffdd to
fe6df82
Compare
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" |
|
PR_Github #53946 [ run ] triggered by Bot. Commit: |
|
PR_Github #53946 [ run ] completed with state |
|
/bot run |
|
PR_Github #53963 [ run ] triggered by Bot. Commit: |
|
PR_Github #53963 [ run ] completed with state |
3 participants
Summary by CodeRabbit
Release Notes
Chores
Refactor
Description
History
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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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.