AutoEP + AutoTP folding: allow CPU optimizer offload under ZeRO-3

[0z5a]

Summary

#8670 left optimizer and parameter offload rejected for AutoEP + AutoTP folding ("not validated yet ...
run a follow-up proof for per-family replica groups"). This PR provides that proof for ZeRO-3 with the
optimizer offloaded to the CPU, and accepts that configuration:

{
  "zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}},
  "tensor_parallel": {"autotp_size": 2},
  "expert_parallel": {"enabled": true, "autoep_size": 2}
}

Everything else stays rejected, now with a message that names the unsupported part: optimizer offload
with ZeRO stages 1 and 2, NVMe optimizer offload, parameter offload, and Muon with optimizer offload.
An explicit "device": "none" is treated as no offload (it was rejected before).

Why ZeRO-3 + CPU optimizer offload is already correct

• ZeRO-3 applies the folded per-family gradient correction (_apply_autoep_folding_tp_correction:
  AVERAGE for TP-replicated parameters, /tp_size for routed experts, SKIP for TP-sharded ones) in
  reduce_independent_p_g_buckets_and_remove_grads, before every reduce-scatter. That path is the same
  with and without offload; offload only changes where the reduced partitions are stored.
• The offloaded norm (complete_grad_norm_calculation_for_cpu_offload) uses the same folding-aware
  counting as the device norm (_grad_norm_counts_every_rank plus the expert-parallel all-reduce), from
  per-partition norms recorded at the accumulation boundary.

The test below checks the claim end to end instead of relying on this reading.

Changes

• validate_folding_global / validate_autoep_config: take the offload devices instead of
  "offload configured" booleans (offload_optimizer_device, offload_param_device); accept stage 3 +
  "cpu"; reject the rest as above. Internal helpers; the two in-repo callers are updated.
• DeepSpeedEngine: passes the configured devices.
• ZeRO-3: rejects Muon with optimizer offload under folding. The offloaded Muon update gathers full
  gradients on its own path, which this PR does not validate (and which overlaps ongoing Muon offload work).
• Docs: the AutoEP folding constraint in autoep.rst.

Tests

• test_autoep_autotp_folding_config.py: accepted and rejected offload combinations (replaces the
  blanket-rejection test; the duplicate in test_autoep_autotp_grad_parity.py is removed).
• test_autoep_autotp_sharded_attention.py, new TestFoldedZero3ClippedUpdate: reuses AutoEP + AutoTP folding under ZeRO-3 #8670's
  per-gradient comparison against an unfolded ZeRO-0 reference and adds one clipped SGD step. For each
  of 4 folding shapes (world 4; TP x EP = 2x2, 2x4, 4x2, 4x4), zero.Init on/off, and offload on/off,
  it compares losses, every gradient, the global norm, and every parameter's update. Clipping (clip
  1e-2, lr 1/clip) makes the update depend on the norm the optimizer computed, so a wrong offloaded norm
  fails the test; the comparison allows one fp32 ulp of the parameter for recovering an update as
  after - before, and asserts that a 0.1% error in the update is still detected.

Results (on-device)

Hardware: 1 node, 4x NVIDIA H100 (64 GB, NVLink NV6 between all pairs), driver 595.91.07 (CUDA 13.2),
torch 2.12.1+cu130, NCCL 2.29.7, transformers 5.15.1, Python 3.12.1 (MareNostrum 5 ACC).

Command (from tests/)	This PR	Unpatched base 1c862f7df
pytest unit/v1/moe/test_autoep_autotp_sharded_attention.py -k TestFoldedZero3ClippedUpdate	16 passed	offload off: 8 passed; offload on: 8 fail with the old "not validated yet" ValueError (the restriction this PR lifts)
pytest unit/v1/moe/test_autoep_autotp_sharded_attention.py -k "not TestFoldedZero3ClippedUpdate"	18 passed, 15 skipped, 2 failed	18 passed, 15 skipped, 2 failed
pytest unit/v1/moe/test_autoep_autotp_{bf16_folding_parity,checkpoint,deepcompile_reject,dispatch,folding_config,folding_groups,grad_parity,runtime,zero1_overlap_folding}.py	81 passed, 9 skipped	75 passed, 9 skipped
pytest unit/v1/moe/test_autoep_autotp_folding_config.py (CPU)	33 passed	—

The 2 failures are TestTiedOutputHead::test_row_parallel_tied_head_is_refused[False/True] ("AutoTP sharded an
output head whose weight is tied to the input embedding"). They fail identically on the unpatched base in
the same job, so they are pre-existing and unrelated to this change (likely AutoTP tied-head detection
with transformers 5.15.1); not addressed here. The global norms match the reference exactly, for example
folded 2.380656 vs reference 2.380656 (TP 2) and 3.552228 vs 3.552228 (TP 4).

8-rank shapes skip on the 4-GPU node. The GPU runs used base 1c862f7df; the branch is rebased onto
5f266e26d without conflicts (the 4 commits in between touch DeepEP slicing, XPU op builders, ZeRO-2
partition storage and a test footprint, none of the paths changed here). pre-commit run --files on the
changed files passes (yapf, flake8, codespell, check-torchdist, check-license, check-torchcuda).

Compatibility matrix

Validated: ZeRO-3 x optimizer offload {none, cpu} x zero.Init {off, on} x fp32 x folding TP {2, 4}
x EP {2, 4} (world 4), torch SGD on the offloaded CPU fp32 partitions.

Not validated (and either rejected or left unchanged): DeepSpeedCPUAdam's offloaded step (the tests run
torch SGD on the CPU fp32 partitions with zero_force_ds_cpu_optimizer: false; DeepSpeedCPUAdam steps
the same partitions but through its own kernel), bf16/fp16 with offload, 8-rank shapes, NVMe offload
(rejected), parameter offload (rejected), ZeRO-1/2 optimizer offload (rejected; its offloaded norm path is
being changed by #8823), Muon with offload (rejected).

Not a duplicate

No open or merged PR enables offload with AutoEP + AutoTP folding (searched "folding offload",
"AutoEP offload", "AutoTP offload"). #8823 fixes ZeRO-1/2 routed-expert clipping under folding and does not
touch ZeRO-3 or offload; ZeRO-1/2 offload is left rejected so the two do not interact. Tracked in the
Q4 roadmap #8672 (AutoEP + AutoTP folding milestone).

[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-15T14:02:02.973040Z	5c6be23	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.