[moe] remove DPGradScaler that over-scales expert grads under moe_dp

[JazyJiang]

Related issue

Fixes #6464

Related to #6465 (ZeRO grad norm with EP) and #6466 (expert grad sync with SP + EP). Both are separate bugs and are not changed here.

Background

With MoeHybridParallelPlugin (EP > 1, moe_dp > 1), DPGradScalerIn/Out multiply an expert's gradient by moe_dp / k, where k is the number of moe_dp replicas that routed at least one token to that expert. The correct gradient is the average over the whole data-parallel group, in which a replica that skips the expert contributes zero. ZeRO already divides expert gradients by moe_dp and EPGradScaler supplies the 1/ep, so the extra factor inflates the gradient of any expert that is cold on some replica, by up to moe_dp×. Megatron-Core scales expert gradients by 1/dp_size with no routing-dependent factor.

Reproduction (from #6464, latest main, 4 GPUs, ep=2, moe_dp=2, routing forced so that each expert is used by k=1 replica):

	ZeRO grad / DDP grad
layers.0.self_attn.q_proj.weight (dense)	1.000
layers.0.block_sparse_moe.experts.0.w1.weight	2.000 (expected 1.000)

With 8 GPUs (moe_dp=4) the ratios are 4.000 (k=1) and 1.337 (k=3). The existing test_moe_ep_zero.py did not catch this: it routes randomly, so k == moe_dp and the factor is 1, and it is skipped.

Changes

• colossalai/moe/_operation.py: remove DPGradScalerIn and DPGradScalerOut.
• colossalai/shardformer/modeling/{mixtral,deepseek,deepseek_v3}.py: remove their use around each expert, and the per-layer activate_experts computation and all-reduce over the moe_dp group (including its all_reduce_fp8 variant) that only fed them. EPGradScalerIn/Out are unchanged.
• tests/test_moe/test_moe_ep_zero_grad_scale.py (new): ep=2, ZeRO-1, bf16, routing forced so that each expert is used by only part of the moe_dp group. Every parameter's gradient is compared, with relative L2 error < 2e-2, against two references: DDP over the dp group, and a single-GPU fp32 copy of the model run on the concatenation of all ranks' batches. The single-GPU reference has no parallel logic, so it also checks that DDP is the right target. Parametrized for 4 GPUs (moe_dp=2) and 8 GPUs (moe_dp=4); the 8-GPU case skips when fewer GPUs are visible.

Not in this PR:

• Qwen3-MoE EP ([shardformer]: support Qwen3-MoE with expert parallelism #6450) copies DPGradScaler and should drop it there.
• The gradient-norm bug ([zero] Gradient norm (and clipping) is wrong and differs across ranks with expert parallelism #6465) and the SP + EP sync bug ([moe] With all_to_all sequence parallelism, expert gradients are not reduced over the sp group #6466).
• EPGradScaler stays in the model forward. Moving the 1/ep factor into the gradient sync belongs to the EP layer refactor.

User and compatibility impact

• Model accuracy / training numerics: changes, and this is the fix. Expert gradients are now correct when some moe_dp replica routes no token to an expert. Runs with EP > 1, moe_dp > 1 and ZeRO will produce different (correct) expert updates than before. Nothing changes when every expert is used on every replica, or when moe_dp = 1.
• Performance: faster. Each MoE layer drops one all-reduce over the moe_dp group, and the assert activated_experts != 0 on a GPU tensor no longer forces a host sync per expert chunk. In a single-layer benchmark of the old EP forward (Qwen3.5-like shape, 256 experts top-8, EP=4, forward + backward), time goes from 94.6 ms to 60.2 ms, with bit-identical outputs.
• Python API: DPGradScalerIn/Out are removed from colossalai.moe._operation, a private module. No other code in the repository uses them.
• Checkpoint format, default behavior of other plugins, Python/PyTorch/CUDA versions, C++/CUDA kernels, Docker images: none.
• Existing training scripts: no changes needed.

Test environment

• OS: Ubuntu 22.04.5 LTS
• Python: 3.10.12
• GPU: 8 × NVIDIA H20, 1 node
• NVIDIA driver: 595.71.05 (nvidia-smi CUDA Version: 13.2)
• PyTorch: 2.11.0+cu128 (torch.version.cuda: 12.8, NCCL 2.28.9)
• CUDA toolkit (nvcc --version): 12.8 (V12.8.61)
• Transformers: 4.51.3
• flash-attn: 2.7.4.post1
• ColossalAI: this branch, 4c745136 on top of main 250c903c

Note: this is the team's torch 2.11 dev environment. CI uses torch 2.13 / CUDA 13.0, and the CI results will be added below.

Test commands and results

All runs used FAST_TEST unset, on GPUs reserved via gpu-run 0,1,2,3,4,5,6,7.

python -m pytest -p no:cacheprovider -rA tests/test_moe
python -m pytest -p no:cacheprovider -rA tests/test_shardformer/test_model/test_shard_mixtral.py tests/test_shardformer/test_model/test_shard_deepseek.py tests/test_shardformer/test_model/test_shard_deepseek_v3.py

Run	Result
tests/test_moe (this branch)	3 passed, 6 skipped, 0 failed, exit 0
Mixtral / DeepSeek / DeepSeek-V3 ShardFormer (this branch)	5 passed, 0 failed (after one rerun of DeepSeek-V3, see Known issues)
tests/test_moe/test_moe_ep_zero_grad_scale.py on main 250c903c (control)	2 failed: [4] and [8], as expected

Passed: test_moe_checkpoint::test_mixtral_moe_layer[4], test_moe_ep_zero_grad_scale[4], test_moe_ep_zero_grad_scale[8], test_mixtral[4], test_mixtral_3d[8], test_deepseek[4], test_deepseek_3d[8], test_deepseek_v3[4].

The 6 skips are all pre-existing markers in the repository and are not touched by this PR: 4 × "tested in corresponding sharderformer" (test_deepseek_layer, test_mixtral_layer, test_moe_ep_tp, test_moe_ep_zero) and 2 × "requires MOE_TENSOR_PATH reference tensors" (test_kernel).

Logs and JUnit XML are kept on the team cluster in test-results/pr_a_regression_20261015/.

CI (torch 2.13.0+cu130 image sha256:a7c63f5a…, head 799f367c):

Layer	Result
Layer 1 (formatting, lint, CPU tests)	pass
Layer 2 (batches 2, 4a, 5)	pass
Layer 3, batches 8 9, fast_mode=false (run 38729167313)	pass

Layer 3 details from the JUnit reports:

• Batch 8: 47 passed, 0 failed, 21 skipped. test_moe_ep_zero_grad_scale[4] and test_moe_checkpoint::test_mixtral_moe_layer[4] passed; test_moe_ep_zero_grad_scale[8] skipped (needs 8 GPUs). The other MoE skips are the pre-existing markers listed above.
• Batch 9: 51 passed, 0 failed, 8 skipped. test_shard_mixtral::test_mixtral[4] passed; the DeepSeek ShardFormer tests are not collected offline (see Known issues).

Known issues

• DeepSeek-V3 ShardFormer test was rerun once. The first run failed before reaching the test body: hf-mirror.com returned 429 Too Many Requests when downloading the remote-code config.json for deepseek-ai/DeepSeek-V3, and it was not in the local cache. It passed when rerun alone on GPUs 0–3. Not introduced by this PR.
• DeepSeek and DeepSeek-V3 are not covered by CI. Batch 9 runs offline and ignores test_shard_deepseek.py and test_shard_deepseek_v3.py when no local Hub snapshot is present. Coverage for these two models comes from the local run above.
• The 8-GPU cases are not covered by CI. The current Layer 3 pool has 4 GPUs, so test_moe_ep_zero_grad_scale[8] skips there (confirmed in the batch 8 report). It passed locally, as did test_mixtral_3d[8] and test_deepseek_3d[8] (above).
• No new skip, xfail, --ignore or loosened tolerance is added. The new test's tolerance (relative L2 < 2e-2) is explained in the file: bf16 noise is below 1e-2, while the removed factor was off by 0.33×–3× on experts.

Risk and rollback

• Risk: low. The change only removes code. The most likely surprise is that a user's existing EP run with moe_dp > 1 trains slightly differently after upgrading, which is the intended correction. The removed all-reduce only fed the removed scaler, so no other code depends on it.
• Rollback: revert this single commit.

Documentation and release impact

None: no README, installation, API docs, Docker, version.txt, PyPI or DockerHub changes. A changelog line for the next release: "Fix expert gradients being over-scaled under expert parallelism when an expert is not used by every moe_dp replica (#6464)."