[zero] compute the grad norm over all gradient stores at once

[JazyJiang]

Related issue

Fixes #6465

Related to #6464 / #6467 (expert grad scaling) and #6466 (expert grad sync with SP + EP). Those are separate bugs and are not touched here.

Background

With MoeHybridParallelPlugin + ZeRO and EP > 1, dense and expert parameters live in two ZeRO gradient stores: dense on the dp group, experts on the moe_dp group. The grad norm used for clipping (and returned by get_grad_norm()) was wrong in two ways:

1. LowLevelZeroOptimizer.step() added the norms of the two stores after each was square-rooted, giving D + E instead of sqrt(D² + E²).
2. The expert part was only reduced over the moe_dp group, so the experts held by other EP ranks were never counted. Each EP rank ended up with D + E_local, and the value differed across ranks.

In a real MoE model the experts carry most of the gradient, so the norm is under-estimated by roughly 1/sqrt(ep) and clipping fails to trigger. Because each rank clips its own ZeRO shard with its own norm, shards of the same parameter are also scaled differently. Reproduction and numbers are in #6465.

Changes

The norm is now computed the way Megatron-Core does it (get_grad_norm_fp32): each rank sums the p-th powers of the gradient shards it owns across all stores, the sum is all-reduced once, and the root is taken once.

• colossalai/zero/low_level/low_level_optim.py
  • _compute_grad_norm(gradients, norm_type) takes the grads of all stores. The local sum (or max for inf) is all-reduced over _grad_norm_process_groups(), and the root is taken once. Every rank participates even without gradients: the old early return 0.0 skipped the collective.
  • Two hooks:
    • _grad_norm_process_groups(): the groups to reduce over. Defaults to the new grad_norm_process_group argument.
    • _grad_norm_weight(grad): how much this rank's copy of a gradient counts. Defaults to 1.
  • New constructor argument grad_norm_process_group. It defaults to the store's group when there is a single store, so plugins with one store behave as before. It is required when pg_to_param_list has several groups.
  • step() calls _compute_grad_norm once per param group with the grads of all stores, instead of adding per-store norms.
• colossalai/booster/plugin/hybrid_parallel_plugin.py (HybridParallelZeroOptimizer)
  • The old _compute_grad_norm override is replaced by the two hooks: tp and pp groups are added to the reduction groups.
  • The existing duplicate handling moves unchanged into _grad_norm_weight: a param not split by TP counts 1/tp_size per TP rank, and a param shared by pipeline stages counts 1/num_stages per stage.
  • This also removes the inf branch that called super()._compute_grad_norm(gradients, norm_type) with the arguments shifted.
• colossalai/booster/plugin/moe_hybrid_parallel_plugin.py: MoeHybridParallelZeroOptimizer passes its dp group as grad_norm_process_group. That group spans the moe_dp and ep axes, so it covers every dense shard and every expert exactly once; no separate ep reduction is needed.
• Tests
  • New tests/test_moe/test_moe_ep_zero_grad_norm.py (4 GPUs, ZeRO-1, bf16, random routing). For ep=2 and ep=4, and for training only the MoE blocks or all parameters, it checks two things:
    • every rank reports the same norm;
    • the norm matches an fp32 single-GPU reference run on the concatenated global batch, within relative 5e-3.
  • tests/test_shardformer/test_hybrid_parallel_grad_clip_norm/test_zero_optimizer.py called the private _compute_grad_norm(dp_group, gradients=...) and is updated to the new signature.

Not in this PR:

• Switching the TP duplicate handling from 1/tp_size weights to Megatron's "count on TP rank 0 only". The copies are not guaranteed to be bit-identical, so that would change TP training numerics.
• With all_to_all SP, expert grads are currently replicated, and not synchronized, across the sp axis ([moe] With all_to_all sequence parallelism, expert gradients are not reduced over the sp group #6466). The dp group then counts them sp times. That case is wrong before and after this PR and becomes correct once [moe] With all_to_all sequence parallelism, expert gradients are not reduced over the sp group #6466 is fixed, with no change needed here.
• The zero_stage=0 path of HybridParallelPlugin / MoeHybridParallelPlugin, which has its own norm code.

User and compatibility impact

• Model accuracy / training numerics: with EP > 1, ZeRO and max_norm > 0, clipping now uses the correct norm, identical on all ranks, so clipped runs change; this is the fix. get_grad_norm() changes the same way. Without EP (one gradient store) the result is unchanged up to floating-point summation order.
• Performance: negligible. One all-reduce of a single float per param group per extra group, instead of one per store; with EP that is fewer collectives than before.
• Python API: the private LowLevelZeroOptimizer._compute_grad_norm drops its dp_pg argument. LowLevelZeroOptimizer and HybridParallelZeroOptimizer get an optional grad_norm_process_group argument. A custom subclass that passes several groups in pg_to_param_list must now also pass grad_norm_process_group; the only one in the repository, MoeHybridParallelZeroOptimizer, is updated.
• Checkpoint format, default behavior, 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_hybrid_parallel_grad_clip_norm -k "not 3d"
python -m pytest -p no:cacheprovider -rA tests/test_zero/test_low_level
python -m pytest -p no:cacheprovider -rA tests/test_booster/test_plugin/test_low_level_zero_plugin.py tests/test_booster/test_plugin/test_3d_plugin.py
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 (this branch)	Result
tests/test_moe	2 passed, 6 skipped, 0 failed
test_hybrid_parallel_grad_clip_norm, 4-GPU cases (AMP, naive, ZeRO optimizers with TP/PP)	3 passed, 0 failed
tests/test_zero/test_low_level	4 passed, 1 failed (test_grad_acc, fails identically on main, see Known issues)
test_low_level_zero_plugin.py, test_3d_plugin.py	2 passed, 0 failed
Mixtral / DeepSeek ShardFormer, 4-GPU cases	2 passed, 0 failed

New test, per config (4 GPUs), ZeRO norm / reference norm:

Config	This branch	main 250c903c
ep=2, MoE blocks only	0.9989, equal on all ranks	0.58–0.86, differs across ranks
ep=2, all params	1.0001, equal on all ranks	1.012–1.018, differs across ranks
ep=4, MoE blocks only	0.9988, equal on all ranks	0.33–0.73, differs across ranks
ep=4, all params	1.0001, equal on all ranks	1.006–1.015, differs across ranks

test_moe_ep_zero_grad_norm.py: 1 passed on this branch, 1 failed on main. Its tolerance (5e-3) is 4× the largest error with the fix and smaller than the smallest error on main, so the value check and the all-ranks-equal check each catch the bug on their own.

The first run of the grad-clip suite also failed test_zero_optimizer.py::test_grad_clip_norm, because the test called the private _compute_grad_norm(dp_group, gradients=...). After updating the call it passes (counted above). The 6 tests/test_moe skips are pre-existing markers.

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

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

Layer	Result
Layer 1 (formatting, lint, CPU tests)	pass
Layer 2 (batches 2, 4a, 5; includes test_low_level_zero_plugin.py)	pass
Layer 3, batches 8 9, fast_mode=false (run 38734809131)	pass

Layer 3 details from the JUnit reports:

• Batch 8: 47 passed, 0 failed, 20 skipped. Passed include test_moe_ep_zero_grad_norm[4], test_moe_checkpoint::test_mixtral_moe_layer[4], test_low_level_zero_checkpointIO and test_grad_acc::test_grad_accumulation; the MoE skips are the pre-existing markers.
• Batch 9: 51 passed, 0 failed, 8 skipped. Passed include the 4-GPU test_grad_clip_norm of the AMP, naive and ZeRO optimizers and test_shard_mixtral::test_mixtral[4].

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

• tests/test_zero/test_low_level/test_grad_acc.py fails locally on torch 2.11, on this branch and on main alike. The reference model is built with DDP(..., bucket_cap_mb=0), which raises AssertionError: bucket_cap_mb should be set by now before any ZeRO code runs. Not introduced by this PR. It passes in CI on torch 2.13 (batch 8 above).
• 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_grad_clip_norm_3d of the AMP, naive and ZeRO optimizers (3 passed), test_mixtral_3d[8] and test_deepseek_3d[8] (2 passed).
• DeepSeek is not covered by CI (batch 9 runs offline and ignores the DeepSeek ShardFormer tests); coverage comes from the local run above.
• SP + EP: the norm is still wrong when all_to_all SP is combined with EP, until [moe] With all_to_all sequence parallelism, expert gradients are not reduced over the sp group #6466 is fixed. See "Not in this PR".
• No new skip, xfail, --ignore or loosened tolerance is added. One existing test is changed only to follow the new private signature.

Risk and rollback

• Risk: the norm code is shared by every LowLevelZero-based plugin. With a single store the reduction groups and weights are the same as before, and the existing grad-clip tests pass. With TP/PP the tp and pp reductions now happen once on the combined sum instead of inside each per-store call. That is mathematically the same, but any subclass overriding the old _compute_grad_norm(dp_pg, ...) signature would break; none exists in the repository.
• 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 the ZeRO gradient norm and clipping with expert parallelism; the norm was wrong and differed across ranks (#6465)."