Fix routed-expert gradient clipping for AutoEP + AutoTP folding - #8823

…epspeedai#8829)

### Summary

deepspeedai#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:

```json
{
  "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 deepspeedai#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 deepspeedai#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"). deepspeedai#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 deepspeedai#8672 (AutoEP + AutoTP folding milestone).

Signed-off-by: 0z5a <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
Co-authored-by: pengdurice <[email protected]>