[https://nvbugs/6412108][fix] AutoDeploy:Shard shared Qwen3.5 experts by greg-kwasniewski1 · Pull Request #16286 · NVIDIA/TensorRT-LLM

greg-kwasniewski1 · 2026-07-12T10:23:35Z

Summary

Fix the Qwen3.5-397B MoE NVFP4 accuracy regression (nvbugs/6412108) by sharding the shared expert instead of keeping it replicated.

Root cause

The shared expert (Qwen3_5MoeMLP) had its torch_linear_simple ops left untagged, and qwen3.5_moe_400b.yaml set apply_sharding_hints.shard_layers: ["moe","delta","mha"], which excludes untagged ("unknown") nodes. So the shared expert stayed replicated (full output on every rank) while the routed experts were sharded. The single merge-point expert_output + shared_expert_output followed by all_reduce then scaled the replicated shared output by world_size (8x), dropping MMLU from ~85 to ~0.05.

Fix

Shard the shared expert like the routed experts rather than special-casing the reduction order:

Tag the shared expert's gate/up (colwise) and down (rowwise) linears with layer_type="mlp" in modeling_qwen3_5_moe.py.
Remove the shard_layers whitelist from qwen3.5_moe_400b.yaml so all tagged nodes shard.
Keep the existing add-then-all_reduce merge point: both routed and shared branches now emit per-rank partials, so one reduction on the sum lifts both to full.

This decouples modeling-code numerical correctness from an optional YAML knob (the whitelist).

Relation to #15922

This supersedes the alternative fix in #15922, which instead reordered the collective (all_reduce the routed partial first, then add the replicated shared output). That approach keeps the shared expert replicated and leaves correctness dependent on the yaml whitelist. Local benchmarking (Qwen3.5-397B-A17B-NVFP4, 8xB200, ISL/OSL 1024, concurrency 1/16/32/64) shows replicating vs sharding the shared expert is within ~1-2% throughput (run-to-run noise, no consistent winner), so there is no perf reason to keep it replicated; sharding both paths is the simpler, self-consistent fix.

Test plan

TestQwen3_5_397B_MoE::test_nvfp4[8] accuracy: MMLU ~86.8 (was ~0.05), passes threshold; unwaived here.
test_sharding_num_correctness.py for modeling_qwen3_5_moe across tp-only / ep-only / tep / attn-dp (normal pass + SHARDING_IR_SABOTAGE=1 true-negative).
tests/unittest/auto_deploy/multigpu/ suite.
CI green (triggered after open).

Risk

library-visible (AutoDeploy sharding / modeling).

Summary by CodeRabbit

Performance
- Improved deployment configuration for Qwen3.5 MoE models with enhanced sharding and multi-stream execution.
- Optimized expert-layer parallelism to support more efficient inference.
Compatibility
- Improved coordination of model outputs during distributed execution.
Testing
- Expanded integration test coverage by removing an outdated waiver for Qwen3.5 MoE accuracy testing.

greg-kwasniewski1 · 2026-07-12T10:24:01Z

/bot run

coderabbitai · 2026-07-12T10:26:09Z

📝 Walkthrough

Walkthrough

Changes

Qwen3.5 MoE sharding

Layer / File(s)	Summary
Sharding and streaming configuration `examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml`	Replaces `shard_layers` with an `lm_head` filter and enables compile-stage MoE/GEMM multi-streaming plus logits gathering.
MLP sharding annotations and waiver `tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py`, `tests/integration/test_lists/waives.txt`	Tags shared-expert MLP projections for sharding and removes the related NVFP4 test waiver.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: yechank-nvidia, Hudayday, litaotju, aswinvisva, lfr-0531

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the required ticket/type format and clearly summarizes the shared-expert sharding fix.
Description check	✅ Passed	The description is complete and well-structured, covering summary, root cause, fix, related work, tests, and risk.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

greg-kwasniewski1 · 2026-07-12T13:48:29Z

/bot run

greg-kwasniewski1 · 2026-07-12T14:47:12Z

/bot run

tensorrt-cicd · 2026-07-12T14:53:55Z

PR_Github #58811 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

tensorrt-cicd · 2026-07-12T14:53:56Z

PR_Github #58811 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

tensorrt-cicd · 2026-07-12T14:55:01Z

PR_Github #58813 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

tensorrt-cicd · 2026-07-12T14:55:02Z

PR_Github #58813 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

tensorrt-cicd · 2026-07-12T14:56:09Z

PR_Github #58815 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

tensorrt-cicd · 2026-07-12T14:56:10Z

PR_Github #58815 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

chzblych · 2026-07-12T17:28:24Z

/bot run

tensorrt-cicd · 2026-07-12T17:38:39Z

PR_Github #58830 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

tensorrt-cicd · 2026-07-12T22:13:46Z

PR_Github #58830 [ run ] completed with state SUCCESS. Commit: 9f8d009
/LLM/main/L0_MergeRequest_PR pipeline #47381 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

greg-kwasniewski1 · 2026-07-13T10:46:45Z

/bot run

The Qwen3.5 MoE shared expert was force-replicated via the qwen3.5_moe_400b.yaml `shard_layers` whitelist (its linears were left untagged so the whitelist excluded them). Because the routed experts are sharded, the single merge-point `expert_output + shared_expert_output` then `all_reduce` scaled the replicated shared output by world_size (8x), collapsing MMLU from ~85 to ~0.05. Instead of special-casing the reduction order to accommodate a replicated shared expert (which couples modeling-code correctness to an optional yaml), shard the shared expert too: tag its gate/up (colwise) and down (rowwise) linears with `layer_type="mlp"` and remove the `shard_layers` whitelist so both routed and shared experts emit per-rank partials. The existing add-then-all_reduce merge point is then correct. Benchmarks show replicating the shared expert yields no throughput benefit over sharding it (within ~1-2% across concurrency 1/16/32/64 on 8xB200), so sharding both paths is the simpler, self-consistent fix. Unwaives accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8]. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

greg-kwasniewski1 · 2026-07-14T08:40:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-14T08:46:12Z

PR_Github #59159 [ run ] triggered by Bot. Commit: de481d3 Link to invocation

tensorrt-cicd · 2026-07-14T13:38:53Z

PR_Github #59159 [ run ] completed with state SUCCESS. Commit: de481d3
/LLM/main/L0_MergeRequest_PR pipeline #47667 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

taylor-yb-lee · 2026-07-15T04:29:32Z

Perf verified

chzblych

Approved for the waive list change.

greg-kwasniewski1 requested a review from a team as a code owner July 12, 2026 10:23

greg-kwasniewski1 requested a review from taylor-yb-lee July 12, 2026 10:23

github-actions Bot assigned greg-kwasniewski1 Jul 12, 2026

greg-kwasniewski1 added this to AutoDeploy Board Jul 12, 2026

github-project-automation Bot moved this to Backlog in AutoDeploy Board Jul 12, 2026

greg-kwasniewski1 added the bug Something isn't working label Jul 12, 2026

nv-guomingz changed the title ~~[https://nvbugs/6412108][fix] Shard shared Qwen3.5 experts~~ [https://nvbugs/6412108][fix] AutoDeploy:Shard shared Qwen3.5 experts Jul 13, 2026

greg-kwasniewski1 requested review from a team as code owners July 14, 2026 07:40

greg-kwasniewski1 requested review from BowenFu, EmmaQiaoCh, schetlur-nv and yuanjingx87 July 14, 2026 07:40

greg-kwasniewski1 force-pushed the gk/shard-qwen35-shared-experts branch from 7578296 to de481d3 Compare July 14, 2026 08:38

greg-kwasniewski1 removed request for BowenFu, EmmaQiaoCh, schetlur-nv and yuanjingx87 July 14, 2026 14:04

taylor-yb-lee approved these changes Jul 15, 2026

View reviewed changes

chzblych approved these changes Jul 15, 2026

View reviewed changes

github-project-automation Bot moved this from Backlog to In review in AutoDeploy Board Jul 15, 2026

greg-kwasniewski1 merged commit adc8c4d into NVIDIA:main Jul 16, 2026
7 checks passed

github-project-automation Bot moved this from In review to Done in AutoDeploy Board Jul 16, 2026

Uh oh!

Conversation

greg-kwasniewski1 commented Jul 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Root cause

Fix

Relation to #15922

Test plan

Risk

Summary by CodeRabbit

Uh oh!

greg-kwasniewski1 commented Jul 12, 2026

Uh oh!

coderabbitai Bot commented Jul 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

greg-kwasniewski1 commented Jul 12, 2026

Uh oh!

greg-kwasniewski1 commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

chzblych commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

greg-kwasniewski1 commented Jul 13, 2026

Uh oh!

greg-kwasniewski1 commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

taylor-yb-lee commented Jul 15, 2026

Uh oh!

chzblych left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

greg-kwasniewski1 commented Jul 12, 2026 •

edited

Loading

coderabbitai Bot commented Jul 12, 2026 •

edited

Loading