[https://nvbugs/6412108][fix] Restore original order — all_reduce the routed partial first, then add the… - #15922
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthroughThe reduction order in Qwen3_5MoeSparseMoeBlock.forward is changed so the routed expert output is all-reduced before adding the shared expert output, rather than summing first then reducing. A test waiver entry for a Qwen3.5 MoE NVFP4 accuracy test is removed. ChangesMoE reduction fix and waiver update
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
da126a4 to
285807f
Compare
|
/bot run |
|
PR_Github #58496 [ run ] triggered by Bot. Commit: |
|
PR_Github #58496 [ run ] completed with state
|
285807f to
b656654
Compare
1db06b6 to
93d7a64
Compare
…replicated shared expert The shared expert (Qwen3_5MoeMLP) intentionally omits the layer_type hint on its torch_linear_simple ops and the qwen3.5_moe_400b.yaml shard_layers whitelist excludes it, so its output is already the full value on every rank. The previous single-merge-point ordering (add then all_reduce) scaled the replicated shared output by world_size and dropped MMLU from ~85% to ~0.07%. Restore the original order: all_reduce the routed partial first, then add the replicated shared output. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
93d7a64 to
17322b1
Compare
Summary
expert_output + shared_expert_outputandall_reduce, scaling the replicated (unsharded) shared_expert output by world_size (8×) and corrupting MMLU/GSM8K outputs.all_reducethe routed partial first, then add the replicated shared output — and update the comment; also removed the nvbugs/6412108 waiver line.Test plan
Links
Summary by CodeRabbit
Bug Fixes
Tests