[Bug] AutoEP FP32 gradient clipping averages norms across ranks with different expert ownership

[gss10282023]

Summary

DeepSpeed's FP32 gradient-clipping path computes the wrong global norm when AutoEP places different experts on different ranks. It averages rank-local norms as if every rank held the same parameters, but under AutoEP the ranks own different expert parameters.

Redistributing the same unique expert gradients across ranks can therefore change the clipping coefficient and the resulting update when clipping is active.

System info

Environment: DeepSpeed commit cf443004, two CUDA ranks, FP32 gradients, ZeRO-0, L2 clipping, mpu=None.

The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in 0166cc87; the full experiment has not been rerun on that commit.

Separate current-branch CPU validation: Ubuntu 26.04 LTS (x86_64), Python 3.11.15, PyTorch 2.13.0+cpu, Intel Xeon Cascadelake, two local CPU/Gloo processes, source baseline 0166cc87 and PR #8476. No GPU or inter-node connection is used in that CPU run.

To Reproduce

The test enters through DeepSpeedEngine.clip_fp32_gradients() and represents the same four unique expert gradients in two layouts:

• EP size 1: all four experts replicated on both ranks.
• EP size 2: each rank owns two of the four experts.

The unique gradient vector and the clipping threshold are fixed. The expected coefficient comes directly from the L2 norm of the unique vector, counting each expert once.

Comparison	Absolute difference in clipping coefficient
Replicated vs. owner-sharded, current code	0.2000005
Replicated vs. owner-sharded, with repair	0
Fully replicated control	0

The replicated layout matches the direct reference; the owner-sharded layout does not. The discrepancy is identical across three repeated blocks, and the repair was confirmed on a second host.

The regression code is in PR #8476. TestAutoEPFP32Clipping exercises the native engine against an independent unique-gradient reference. PR #8476 includes the exact two-process CPU/Gloo command; all eight cases passed with the repair, and the owner-sharded case fails on unmodified 0166cc87.

Expected behavior

Replicated and owner-sharded layouts of the same unique gradients should give the same global L2 norm, clipping coefficient, and optimizer update. Dense parameters should count once despite replication, and each expert should count once according to its ownership.

Root cause

clip_grad_norm_ takes a local square root before averaging across the data-parallel group. When ranks own different experts, the average of local norms does not reconstruct the norm of the concatenated unique gradients.

Replication counts differ as well: dense parameters live on every DP rank, while AutoEP expert parameters live on only dp_world_size / ep_size ranks. Their contributions need different weights for each logical parameter to count exactly once.

Proposed fix

Record the EP size on each AutoEP expert parameter. For the FP32 L2 path:

1. Sum squared gradients locally, dividing each parameter's contribution by its replica count.
2. Sum those contributions across the DP group.
3. Take a single square root after the reduction and compute the clipping coefficient.

The tested patch adds the ownership metadata in auto_ep_layer.py and implements this reduction in runtime/utils.py. Its validated scope is the ZeRO-0, mpu=None, L2 case above.

ds_report output

Collected with DS_ACCELERATOR=cpu python -m deepspeed.env_report from the current PR checkout. This describes the separate CPU validation environment; it is not the pinned two-GPU environment. Local installation paths have been replaced with placeholders.

CPU environment report

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
deepspeed_not_implemented  [NO] ....... [OKAY]
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
deepspeed_ccl_comm ..... [NO] ....... [OKAY]
deepspeed_shm_comm ..... [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
pin_memory ............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['<local dependency path>']
torch version .................... 2.13.0+cpu
deepspeed install path ........... ['<DeepSpeed checkout>/deepspeed']
deepspeed info ................... 0.19.6, [none], [none]
deepspeed wheel compiled w. ...... torch 0.0 
shared memory (/dev/shm) size .... 30.27 GB

Launcher context

The current regression uses python -m torch.distributed.run --standalone --nproc-per-node=2 with Gloo. The exact driver is included in PR #8476.

Docker context

The current CPU regression ran in an existing Ubuntu environment; no separately specified Docker image is required by the reproducer. The pinned GPU report does not specify an image tag.

[alanhuangyoo]

@gss10282023 Moving the thread here — I had been replying on #8469 before noticing you withdrew it for this one. #8478 now points at this issue rather than the closed one.

Recap of where it stands, so nothing is buried in a closed issue:

The two defects you found in #8478 are fixed (338ed51), both reproduced on 2×H20 first. Your comment about inspecting ownership before the gradient filter was the answer, and there was a harder case underneath it — a rank owning no expert parameter at all, which happens under pipeline parallelism. That is now settled by one all-reduced flag rather than by local inspection, with group names from _EXPERT_PARALLEL_GROUP.

uneven  (rank1 owns no expert at all)     both ranks return 4.4721
nograd  (rank1's expert got no tokens)    both ranks return 4.4721
inf     (expert-only, norm_type=inf)      both ranks return 2.0000

On the split, the open question is still yours to answer: can AutoEP experts carry a group_name and an entry in _EXPERT_PARALLEL_GROUP? If they can, get_norm_with_moe_layers covers both families and the ownership weighting falls out of the group size, so there is one implementation instead of two special cases in clip_grad_norm_. If they cannot, #8476 is the right shape and I will fold legacy MoE into it rather than keep a parallel branch.

You reported this and had a fix up first. If a maintainer wants a single PR, it should be yours.