[moe] sync and shard expert grads over moe_dp * sp with all_to_all sequence parallelism

[JazyJiang]

Related issue

Fixes #6466

Related to #6464 / #6467 (expert grad scaling) and #6465 (ZeRO grad norm with EP). Those are separate bugs and are not touched here.

Background

With MoeHybridParallelPlugin, EP > 1 and sequence_parallelism_mode="all_to_all", every sequence is split along the sequence dimension and the MoE blocks run on each rank's part. The ranks of one sp group hold the same experts but see different tokens, so their expert gradients must be summed over the sp group, exactly as for dense parameters.

Dense parameters are synced over mixed_dp_group, which spans the (moe_dp, ep, sp) axes. Expert parameters, however, were put in a ZeRO store on moe_dp_group, which spans only the moe_dp axis. Each sp rank therefore kept the expert gradient of its own part of the sequence. With ep=2, sp=2 on 4 GPUs, expert grads were 0.8–1.0 off (relative L2) from a single-GPU reference, while the mean of the two sp ranks matched it to 0.007. The expert copies on different sp ranks drift apart over training, and sharded checkpoints keep only the copy from sp rank 0. Reproduction in #6466.

The existing test_shard_mixtral.py / test_shard_deepseek.py include ep=2, sp=2 configs but only compare parameters after training with atol=4e-3, which is larger than the gradients of the tiny test models.

Changes

• colossalai/booster/plugin/moe_hybrid_parallel_plugin.py
  • New mixed_moe_dp_group, the expert counterpart of mixed_dp_group. With all_to_all SP it is a group along the (moe_dp, sp) axes; otherwise it is moe_dp_group, so nothing changes without SP.
  • MoeHybridParallelZeroOptimizer and MoECheckpointIO use it for expert parameters. Expert grads are therefore reduced over moe_dp × sp and divided by its size. The existing grad_scale=sp_size on the output gather makes this the correct average.
  • Expert params, grads and optimizer states are now sharded over moe_dp × sp instead of being replicated on every sp rank.
  • plugin.moe_dp_group itself is unchanged; the model code still uses it.
• colossalai/checkpoint_io/moe_checkpoint.py: when loading a sharded optimizer, expert states were padded by the global dp size before being split by the moe_dp size. LowLevelZero pads each param by the size of the group it is sharded over, so the two disagreed whenever an expert param's numel is not divisible by the global dp size, and the loaded states ended up in the wrong place. They are now padded by the moe_dp size. This mismatch exists on main independently of SP (e.g. ep=2, moe_dp=2, a 30-element expert param); this PR makes it more likely to hit, so it is fixed here.
• New tests/test_moe/test_moe_ep_sp_grad_sync.py (4 GPUs, ep=2, sp=2 all_to_all, bf16):
  • Gradients, ZeRO-1 and ZeRO-2: every parameter's grad is compared with an fp32 single-GPU model run on every data replica's batch (relative L2 < 3e-2). Grads must also be identical on the ranks of each sp group.
  • Checkpoint: one Adam step, then save the sharded model and optimizer, reload, and require equal parameters and bit-identical optimizer states. Sizes are chosen so expert params (30 elements) are not divisible by the dp size (4), which exercises the padding fix.

Not in this PR:

• The zero_stage=0 / DDP path. With EP > 1 it already requires the dp and moe_dp groups to be identical and raises otherwise.
• Tightening the post-training parameter check in test_shard_mixtral.py / test_shard_deepseek.py.

User and compatibility impact

• Model accuracy / training numerics: EP + all_to_all SP runs now get correct expert grads, so they train differently; this is the fix. Nothing changes without all_to_all SP.
• Performance / memory: expert grads gain an all-reduce over the sp axis, the same communication dense params already pay. Expert optimizer states and fp32 master weights are now sharded over sp too, so their memory per rank drops by a factor of sp.
• Checkpoint format: the files on disk are unchanged; the full state is gathered as before. With all_to_all SP, loading a sharded optimizer checkpoint now re-shards expert states over moe_dp × sp. Old checkpoints of EP + SP runs contain the diverged expert copy from sp rank 0, which loads but carries the effect of the bug.
• Python API: new attribute MoeHybridParallelPlugin.mixed_moe_dp_group; nothing removed.
• Default behavior of other plugins, Python/PyTorch/CUDA versions, C++/CUDA kernels, Docker images, existing training scripts: none.

Test environment

• OS: Ubuntu 22.04.5 LTS
• Python: 3.10.12
• GPU: 8 × NVIDIA H20, 1 node (4 used for the runs below)
• 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 on top of main 250c903c

CI uses torch 2.13 / CUDA 13.0; CI results will be added below.

Test commands and results

All runs: FAST_TEST unset, 4 GPUs reserved via gpu-run.

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 -k "not 3d"

Run	Result
tests/test_moe (this branch)	2 passed, 6 skipped, 0 failed
Mixtral / DeepSeek ShardFormer, 4-GPU cases (this branch; include ep=2, sp=2)	2 passed, 0 failed (2 deselected: the 8-GPU _3d cases)
tests/test_moe/test_moe_ep_sp_grad_sync.py on main 250c903c (control)	1 failed: e.g. experts.2.w1 grad 0.884 off the reference, grads differ across the sp group
Same test on this branch with only the padding line reverted	1 failed: grads pass, the optimizer checkpoint round trip reports differing states

Passed: test_moe_checkpoint::test_mixtral_moe_layer[4], test_moe_ep_sp_grad_sync[4], test_mixtral[4], test_deepseek[4].

The 6 skips are pre-existing markers, not touched by this PR: 4 × "tested in corresponding sharderformer" and 2 × "requires MOE_TENSOR_PATH reference tensors".

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

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

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 38770628022, attempt 2)	pass

Layer 3 details from the JUnit reports:

• Batch 8: 47 passed, 0 failed, 20 skipped. test_moe_ep_sp_grad_sync[4] and test_moe_checkpoint::test_mixtral_moe_layer[4] passed; the MoE skips are the pre-existing markers.
• Batch 9: 51 passed, 0 failed, 8 skipped. test_shard_mixtral::test_mixtral[4] passed.

Attempt 1 of the Layer 3 run failed during job setup, before any test ran: the runner timed out downloading actions/upload-artifact from GitHub. It passed on re-run. The Draft-only container smoke (batch 4a) check shows as failed because it was cancelled when the PR was marked ready; batch 4a ran and passed in Layer 2.

Known issues

• The 8-GPU cases are not covered by CI (the Layer 3 pool has 4 GPUs). They were run locally on 8 × H20 and pass: test_mixtral_3d[8] and test_deepseek_3d[8] (2 passed).
• DeepSeek is not covered by CI. Batch 9 runs offline and ignores test_shard_deepseek.py / test_shard_deepseek_v3.py without a local Hub snapshot; coverage comes from the local run above.
• No new skip, xfail, --ignore or loosened tolerance is added. The new test's gradient tolerance (relative L2 < 3e-2) is explained in the file: bf16 noise stays below 2e-2, while the bug is 0.8–1.0 off. The checkpoint check is exact (torch.equal).

Risk and rollback

• Risk: only runs with all_to_all SP and EP are affected; otherwise mixed_moe_dp_group is moe_dp_group and every call site gets the same group as before. The most likely issue is a code path that relies on expert states being complete on every sp rank. The ZeRO store and MoECheckpointIO are the two consumers and both are updated and covered by the new test.
• 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: "Fix expert gradients not being synchronized across sequence-parallel ranks with all_to_all SP + EP, and the re-sharding of expert optimizer states when loading checkpoints (#6466)."