[moe] Expert gradients are scaled by moe_dp/k when an expert is used by only k of the moe_dp replicas #6464

Checklist

• I have searched the existing issues and did not find a duplicate.
• I have reproduced this on the latest main (250c903).

Describe the bug

With MoeHybridParallelPlugin (EP > 1, moe_dp > 1, ZeRO), the gradient of an expert is multiplied by moe_dp / k, where k is the number of moe_dp replicas that routed at least one token to that expert in this step. The correct gradient does not depend on k, so whenever an expert is cold on some replica, its gradient is inflated by up to moe_dp×.

The factor comes from DPGradScalerIn / DPGradScalerOut (colossalai/moe/_operation.py), applied around each expert in the EP forward of Mixtral, DeepSeek and DeepSeek-V3 (colossalai/shardformer/modeling/{mixtral,deepseek,deepseek_v3}.py). The k is obtained from an extra all-reduce of activate_experts over the moe_dp group in every MoE layer.

Why the factor is wrong. The target is the gradient of the global mean loss, i.e. the average of the per-rank gradients over the whole data-parallel group. A replica that sends no token to an expert contributes a zero gradient to that average; the zero is part of the correct result, not something to compensate for. ZeRO already reduces expert gradients over the moe_dp group and divides by moe_dp, and EPGradScaler supplies the remaining 1/ep, so the result is correct without DPGradScaler. Megatron-Core does the same: expert gradients are scaled by 1/dp_size regardless of routing (megatron/core/distributed/distributed_data_parallel.py), with no activation-dependent factor.

Why tests did not catch it. tests/test_moe/test_moe_ep_zero.py uses random routing, under which every expert is hit on every replica, so k == moe_dp and the factor is 1. The test is also skipped.

Impact. With fine-grained experts and small per-rank batches, "expert e got no token on some replica" is common, so expert gradients get a different, routing-dependent multiplier at every step. In addition, the per-layer all-reduce and the assert activated_experts != 0 on a GPU tensor in every chunk add host syncs; removing them makes the old EP layer about 1.5× faster (forward + backward of one layer, Qwen3.5-like shape, EP=4: 94.6 → 60.2 ms) with bit-identical outputs.

Minimal reproduction

# torchrun --nproc_per_node=4 repro_dp_grad_scale.py
# ep=2, moe_dp=2, ZeRO-1, bf16. The router is forced so that moe_dp replica 0 only uses expert 0 and replica 1 only
# uses expert 1, i.e. each expert is used by k=1 of the moe_dp=2 replicas. Reference: DDP over the whole dp group.
from copy import deepcopy

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import MixtralConfig, MixtralModel

import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import MoeHybridParallelPlugin

colossalai.launch_from_torch(seed=0)
rank = dist.get_rank()
torch.cuda.set_device(rank)
EP, H = 2, 8

plugin = MoeHybridParallelPlugin(pp_size=1, tp_size=1, ep_size=EP, zero_stage=1, precision="bf16")
booster = Booster(plugin=plugin)
torch.manual_seed(0)
config = MixtralConfig(hidden_size=H, intermediate_size=2 * H, num_hidden_layers=2, num_attention_heads=2,
                       num_key_value_heads=2, num_local_experts=4, num_experts_per_tok=1)
model = MixtralModel(config).to(torch.bfloat16).cuda()
with torch.no_grad():  # tokens with hidden[0] > 0 go to expert 0, the others to expert 1
    for layer in model.layers:
        layer.block_sparse_moe.gate.weight.zero_()
        layer.block_sparse_moe.gate.weight[0, 0], layer.block_sparse_moe.gate.weight[1, 0] = 10.0, -10.0

zero_model = deepcopy(model)
zero_model, zero_opt, *_ = booster.boost(zero_model, torch.optim.SGD(zero_model.parameters(), lr=1))
ddp_model = DDP(model, process_group=plugin.dp_group, find_unused_parameters=True)

torch.manual_seed(1 + rank)
x = torch.rand(4, 7, H).cuda().to(torch.bfloat16)
x[..., 0] = 5.0 if rank // EP == 0 else -5.0  # moe_dp replica 0 -> expert 0, replica 1 -> expert 1
zero_opt.backward(zero_model(inputs_embeds=x).last_hidden_state.mean())
ddp_model(inputs_embeds=x).last_hidden_state.mean().backward()

ref = dict(ddp_model.module.named_parameters())
for n, p in zero_model.named_parameters():
    n = n.removeprefix("module.")
    g = zero_opt.get_param_grad(p)  # collective, call on every rank
    if rank == 0 and n.startswith("layers.0.") and (".experts.0.w1" in n or "self_attn.q_proj" in n):
        print(f"{n:45s} |ZeRO grad| / |DDP grad| = {g.float().norm() / ref[n].grad.float().norm():.3f}")

Output

layers.0.self_attn.q_proj.weight              |ZeRO grad| / |DDP grad| = 1.000
layers.0.block_sparse_moe.experts.0.w1.weight |ZeRO grad| / |DDP grad| = 2.000

Expected: 1.000 for both. The expert gradient is exactly moe_dp / k = 2 / 1 times too large. With 8 GPUs (moe_dp=4) and experts used by k=1 and k=3 replicas, the ratios are 4.000 and 1.337 (≈ 4/3). To rule out the DDP reference itself, we also compared against a single-GPU fp32 model run on the concatenation of all ranks' batches (whose loss is the global mean loss by definition): DDP matches it within bf16 noise, ZeRO is 2× off on the experts.

Environment

• ColossalAI: main @ 250c903
• PyTorch: 2.11.0+cu128 (CUDA runtime 12.8, NCCL 2.28.9)
• Transformers: 4.51.3
• CUDA toolkit: 12.8; driver 595.71.05
• GPU: 8 × NVIDIA H20, 1 node (4 used)
• OS / Python: Ubuntu 22.04.5, Python 3.10.12
• Plugin: MoeHybridParallelPlugin
• Parallel config: ep=2, moe_dp=2, tp=1, pp=1, sp=1, ZeRO stage 1, bf16

Additional context

• Proposed fix: remove DPGradScalerIn/Out and the activate_experts all-reduce from the three models; keep EPGradScaler. PR to follow, with a regression test that forces partial routing and compares against both DDP and the single-GPU reference.
• The open PR [shardformer]: support Qwen3-MoE with expert parallelism #6450 (Qwen3-MoE EP) copies the same DPGradScaler usage and should drop it as well.
• Related: ZeRO grad norm with EP ([zero] Gradient norm (and clipping) is wrong and differs across ranks with expert parallelism #6465) and expert gradient sync with SP + EP ([moe] With all_to_all sequence parallelism, expert gradients are not reduced over the sp group #6466), filed separately.