From 17322b1c1494951b6ed09a26dbcbcd04b4377bf9 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 3 Jul 2026 04:19:50 -0700 Subject: [PATCH] [nvbugs/6412108][fix] Reduce Qwen3.5 MoE routed output before adding 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> --- .../models/custom/modeling_qwen3_5_moe.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py index b44c5f040982..b68f7e718bd2 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py @@ -780,13 +780,14 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: layer_type="moe", ) - # Single merge-point all_reduce for routed + shared partial sums. - # Both branches produce per-rank partial outputs under TP/EP sharding - # (routed: MoEShardableNode; shared: rowwise down_proj inside the MLP). - # One reduction on the sum lifts both to full; reducing before the add - # would mix a full routed contribution with an unreduced shared one. - expert_output = expert_output + shared_expert_output + # The shared expert is replicated (Qwen3_5MoeMLP intentionally omits + # ``layer_type`` and the yaml ``shard_layers`` whitelist excludes it), + # so its output is already the full value on every rank. All-reduce + # the sharded routed-expert partial first, then add the replicated + # shared output; adding before would scale the shared output by the + # TP world size. expert_output = torch.ops.auto_deploy.all_reduce(expert_output, layer_type="moe") + expert_output = expert_output + shared_expert_output expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim) return expert_output