Fix routed-expert gradient clipping for AutoEP + AutoTP folding

[0z5a]

AutoEP + AutoTP folding gives TP peers different routed experts. ZeRO-1/2 currently filters non-TP expert parameters to TP rank zero when computing the clipping norm, then applies the ordinary expert-norm averaging path. This can omit expert owners and produce an incorrect clipping coefficient even when reconstructed logical gradients are correct.

Compute folded expert norms over the actual expert-DP partition group and EP owner group, without another TP sum or expert-norm average. Only groups tagged with the AutoEP expert placement family use this path. Dense groups and native MoE groups retain their existing reductions and native expert-norm averaging, including when native ETP and AutoEP groups coexist. ZeRO-3 is unchanged and is checked as a sibling reference. The existing folding contract rejects optimizer/parameter offload.

Refs #8672 and #8670. The live duplicate check found no open PR fixing this ZeRO-1/2 folding path. #8476, #8478 and #8691 address FP32/ZeRO-0 clipping in runtime/utils.py; they do not change the partitioned optimizer reductions fixed here. #8739 implements DeepEP weight-gradient overlap and does not change these norm reductions.

Final change

This PR contains three files:

• deepspeed/runtime/zero/stage_1_and_2.py: correct the folded expert clipping reductions.
• tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py: reconstruct full logical gradients and FP32 master parameters, then compare an actual SGD update against an unsharded ZeRO-0 reference across all 45 parameter families. Thresholds of 0.5 and 2.0 times the reference norm cover active clipping and an inactive control. Snapshots are cloned before in-place updates; distributed diagnostics finish before assertions.
• docs/code-docs/source/autoep.rst: document expert ownership and the reduction contract.

The sharded-attention regression supports 4- and 8-rank TP2/EP4 configurations, ZeRO-1/2/3 and ZeRO-3 zero.Init. An additional two-rank analytic SGD regression covers native ETP alone and mixed native ETP/folded AutoEP groups at ZeRO-1/2.

Incremental review

Incremental review report and validation · diff from the reviewed head.

The correction preserves native ETP reductions alongside folded AutoEP. On the correction commit, 8 positive cases passed on 2 H100 GPUs, and 2 published-head negative controls failed as expected. ZeRO-1/2 optimizer cases validate native norm 5, mixed norm 13 and the actual clipped-SGD updates; TP2/EP2 tiny GLM4-MoE active/inactive clipping cases validate all 45 parameter families against an unsharded reference. The old mixed cases produce partial norms 12.369317 / 12.649111 instead of 13. These runs used FP32, TF32 disabled, Torch 2.12.1+cu130 / CUDA 13.0, Transformers 5.15.1 and driver 595.91.07. Changed-file pre-commit checks and independent archive/report readback passed; owned workers and locks closed. The report records the exact execution method and CPU diagnostic separately. The 8/16-GPU results below belong to the original submitted head.

Original submitted-head validation

The original submitted three-file diff at 352ee55cd14f77f30e5170939c7e60c4d4bd7754 was tested on source based on 1c862f7df93f0345b977c67f5408d967b94bf4ae. The extended finite workload used 16 NVIDIA H100 GPUs across four nodes, random tiny GLM4-MoE, FP32 with TF32 disabled, Torch 2.12.1+cu130, CUDA 13.0 and Transformers 5.15.1. The original receipt did not capture the NVML driver package version.

Coverage	Final result
Clipping and initialization	TP2/EP4 and TP2/EP8, ZeRO-1/2/3, active/inactive clipping and ZeRO-3 zero.Init passed. An unfixed ZeRO-2 control failed the norm/update oracle.
Routing and topology	Rotating hotspots, idle expert owners, GAS4, and TP4/EP16 with 16 experts and ZeRO-3 zero.Init passed.
Communication	Overlap off/on and small-bucket configurations with actual repeated/oversize flushes passed.
Save and resume	Continuous four-update training matched save after two updates, process exit and new-process resume through four updates, including parameters, momentum 0.9, losses and WarmupLR state.
Sustained training	200 actual updates across four 50-update routing stages passed the reference loss, gradient and update checks. Late live allocated-memory spread stayed below 64 MiB for this tiny model.
Independent readback	All 28 cases, 448 rank packets and 247,635 raw arrays from the 16-GPU workloads passed separate NumPy arithmetic and complete archive checksum verification.
Source checks	All applicable changed-files pre-commit hooks passed, including yapf, flake8, codespell, license and accelerator/distributed import checks. git diff --check and a clean three-way merge check against current upstream master passed.

The extended distributed drivers used the same contributed reference helpers. An additional 8-H100 clipping matrix passed 16 cases and 28,800 raw arrays. All owned GPU workers and contexts closed, the original GPU locks were released, and complete outputs were archived.

Idle-owner checks exclude alignment padding using the frozen permutation formula ceil((n_tokens + E_local*8)/8)*8; equality to E_local*8 proves zero actual received rows. The original finalizer's incorrect padding assertion and failed scheduler status are preserved in the private evidence. The separate corrected reader passed against the same completed workload and unchanged numerical tolerances.

These results establish tiny-model correctness, optimizer-state continuity and bounded-run stability. They do not establish pretrained-model convergence, a speedup, BF16/FP16 or Muon behavior, offload, PP/SP, DeepEP overlap, or completion of every Q4 roadmap item.

Reproduce the contributed regression and source checks:

pytest --forked -q tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py \
  -k TestShardedAttentionClippedUpdate
pre-commit run --files deepspeed/runtime/zero/stage_1_and_2.py \
  tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py \
  docs/code-docs/source/autoep.rst

AI assistance was used for implementation, test preparation and incremental review. The human submitter reviewed the original submission, confirmed its test results, and authorized this follow-up correction and DCO sign-off.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 352ee55cd1

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Restrict folded reductions to AutoEP expert groups

When folded AutoEP is enabled in a ZeRO-1/2 model that also contains a native deepspeed.moe.layer.MoE group with enable_expert_tensor_parallelism=True, this condition selects the native group solely because the global folding spec exists. Native ETP experts can have gradients sharded over the model-parallel group, but the arguments below replace _model_parallel_all_reduce with only EDP+EP reductions, whose rank tables cover one TP lane, so TP ranks can compute partial norms and use different clipping coefficients. Restrict this branch to groups tagged with the AutoEP expert placement family, or retain the legacy model-parallel reduction for native ETP groups.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-14T17:19:48.830600Z	352ee55	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[0z5a]

Incremental review and corrections

Reviewed the original PR head 352ee55cd14f77f30e5170939c7e60c4d4bd7754. This review concerns the ZeRO-1/2 change in this PR and agrees with the existing native-ETP review finding.

Finding addressed: A global folding spec previously routed every MoE optimizer group through the folded EDP+EP norm path. A native MoE group using expert tensor parallelism still needs its model-parallel reduction; replacing that reduction can leave each TP lane with a partial norm and a different clipping coefficient.

The correction selects the folded path only for groups whose parameters all carry ds_zero_placement_family="autoep_expert". Native groups keep their existing model-parallel reduction and expert-norm averaging, including when folded and native groups coexist. Dense behavior and ZeRO-3 are unchanged.

The new two-rank regression uses actual ZeRO optimizers, collectives, backward and SGD updates. Native ETP gradients with shard norms 3 and 4 must produce norm 5. Adding a folded expert with norm 12 on TP rank one must produce norm 13 on both ranks, followed by the analytically expected clipped parameter update. Both ZeRO-1 and ZeRO-2 are covered, with native-only controls.

Corrected in fb5d0e5. Review only this increment.

Incremental validation completed on 2026-10-14 against the correction commit:

• 8 positive GPU cases passed; 2 published-head negative controls failed as expected. Four optimizer cases cover native ETP alone and mixed native ETP/folded AutoEP groups at ZeRO-1/2. Both corrected mixed cases produce norm 13 on both ranks and pass the analytic clipped-SGD update check. Both published-head controls reproduce the partial norms 12.369317 / 12.649111 instead of 13.
• Four full tiny GLM4-MoE cases cover TP2/EP2, ZeRO-1/2 and thresholds at 0.5 or 2.0 times the reference norm. The folded norm is 3.5522284508 versus 3.5522283833 from the unsharded reference. Reconstructed gradients and actual updates across all 45 parameter families pass, with active clipping and an inactive control.
• Hardware/runtime: 2 NVIDIA H100 GPUs, FP32 with TF32 disabled, Torch 2.12.1+cu130 / CUDA 13.0, Transformers 5.15.1 and driver 595.91.07. A separate two-rank torchrun process ran each case, directly calling the unmodified contributed mixed-expert test method or _clipped_step_against_reference(2, 2, zero_stage, False, clip_ratio). This execution used the test bodies and reference helpers, not the pytest launcher. Source hashes match the correction commit.
• The earlier two-process CPU/Gloo diagnostic also passed on the fixed mixed ZeRO-2 case and reproduced the published-head failure. It used Intel Xeon Platinum 8460Y+ CPUs, FP32, Torch 2.12.1+cu130 and DeepSpeed's native CPU communication extension. Only the test's accelerator-only skip gate was removed in that private CPU adapter; runtime and collective implementations were not mocked. This remains separate CPU evidence.
• All applicable changed-file pre-commit hooks passed, including yapf, flake8, codespell, license and accelerator/distributed import checks. git diff --check passed.
• Complete outputs were archived and independently read back. The separate reader checked archive/member SHA256 values, bound every rank report to the archived bytes, checked the case matrix and norm results, and verified worker closure and GPU-lock release. The autonomous batch completed normally and released its allocation. Earlier private setup/driver failures were preserved and are excluded from the GPU pass counts.

Run the contributed accelerator regression with:

pytest --forked -q tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py \
  -k TestMixedExpertClipping
pre-commit run --files deepspeed/runtime/zero/stage_1_and_2.py \
  tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py \
  docs/code-docs/source/autoep.rst

The mixed-expert regression validates the optimizer boundary with explicit parameter/group metadata. The new TP2/EP2 tiny GLM4-MoE cases validate the full AutoEP/AutoTP model path on the correction commit; they are not a complete mixed native/AutoEP application model or pretrained convergence test. The earlier 8/16-GPU results in the description belong to the original submitted head; they have not all been rerun for this incremental correction.

AI assistance was used for this review, correction and validation.