Add AutoEP + AutoTP parallel folding (#8064)
This PR adds **parallel folding** for AutoEP: tensor parallelism
(AutoTP) for the dense/attention path can now coexist with expert
parallelism (AutoEP) for the routed-expert path on the same set of
ranks, **without forcing EP to be a subset of DP** — an EP group may
span TP lanes and dense-DP ranks (cross-lane EP).
(This PR should be adjusted for ZeRO3 support after #8060 is merged)
## Design
Attention/dense and MoE are treated as two independent partitionings of
the same rank set, parameterized per parameter family:
- Dense / attention / shared-expert params: `stage_size = tp * dp`
- Routed-expert params: `stage_size = ep * etp * edp`
`dp` and `edp` are always derived, never user-configured, so the
invariant `tp * dp == ep * etp * edp == stage_size` cannot be broken
from config. The only structural requirement is that the expert width
tiles the stage (`stage_size % (ep * etp) == 0`); EP groups are then
laid across a TP-lane-major rank ordering, so they may span TP lanes and
dense-DP ranks.
### Configuration
No new config section. Folding is expressed by the coexistence of the
existing `tensor_parallel` and `expert_parallel` sections:
```json
{
"tensor_parallel": { "autotp_size": 4 },
"expert_parallel": { "enabled": true, "autoep_size": 4,
"expert_tensor_parallel_size": 1 }
}
```
`expert_tensor_parallel_size` is carried as a config field but currently
must be `1` (expert-internal TP is reserved as follow-up and rejected
fail-fast). Validation enforces stage divisibility, TP/sequence-parallel
exclusivity, and `preset_model` consistency between the two sections.
### Cross-lane expert parallelism
Expert parallelism no longer has to be a subset of data parallelism.
Shapes where the expert width exceeds (or does not divide) the dense
data-parallel size are supported, for example:
- `world=4, TP=4, EP=4` (`dp=1`): the EP group is the whole TP group —
one expert per rank.
- `world=4, TP=2, EP=4` (`dp=2`): the EP group spans both TP lanes and
both DP ranks.
- `world=8, TP=4, EP=4` (`dp=2, edp=2`): EP groups span TP lanes with
expert replication.
The per-family gradient convention is keyed to each parameter's
replication structure, not to the EP layout, so it holds across the
whole `tp*dp` pool:
- **Router/gate and dense/LayerNorm** are AVERAGE over the TP
(token-replication) group. The folded router runs redundantly on every
TP peer; its partitioned work is reconstructed into a replicated full
view by `restore_combined`, whose all-gather backward injects a
`tp_size` factor that AVERAGE divides out.
- **Routed experts** cancel that same `restore_combined` `tp_size`
factor (divide by `tp_size`, no TP all-reduce) and reduce data-parallel
over the expert-data-parallel (EDP) group. Without the cancellation,
folded expert gradients are over-scaled by `tp_size` — invisible to
scale-invariant Adam, but real for non-adaptive optimizers and for
gradient clipping (it inflates the expert contribution to the global
grad norm). This is now fixed for all folded shapes (the MVP TP2×EP4
shape included).
## What's included
- Folded process-group derivation using the generalized
expert/data-parallel group creation (`mp_mode` TP-strided vs
SP-consecutive ordering), including cross-lane EP group tables.
- Route-full / partition-dispatch path for folded MoE
(`deepspeed/moe/ep_tp_dispatch.py`), with AutoTP skipping AutoEP
subtrees.
- **Per-family folded gradient reduction**: AVERAGE for replicated
router/gate and dense/LayerNorm; a dedicated `tp_size` cancellation for
routed experts; SKIP for genuinely TP-sharded params; SUM contracts
reserved for a future true sequence-parallel path.
- Per-parameter-family ZeRO checkpoint metadata (routed-expert vs
dense/router/shared placement) and folded ZeRO-1/2 optimizer-state
handling.
## Correctness & validation
- Router/gate and LayerNorm gradient parity against a non-folded ZeRO
baseline (atol=1e-1, rtol=5e-3, fp32), on TP2×EP4 (8-rank) and the
cross-lane shapes TP2×EP4 (4-rank, `ep>dp`) and TP4×EP4 (4-rank,
`dp=1`); scale 1.0.
- Routed-expert weight parity against a non-folded baseline, verified
with SGD (Adam is scale-invariant and would mask a uniform
gradient-scale error), for the MVP TP2×EP4 shape and cross-lane TP4×EP4
with `edp=1` and `edp=2`.
- New folding unit tests for config, cross-lane group layout, dispatch,
runtime, gradient parity, and checkpoint save/load (multi-rank GPU cases
gated for GPU runners; CPU/Gloo parity runs on CI).
- Real-H100 confirmation (8×H100): router/gate, LayerNorm, and
routed-expert gradient parity to the non-folded baseline hold (scale
1.0) for MVP TP2×EP4 and cross-lane TP4×EP4 (`edp=1` and `edp=2`);
cross-lane folded training runs with finite loss and finite, non-zero
expert/router gradients.
- Passes the full unit test suite (`aws-torch-latest-full`) on H100
GPUs.
## Scope / follow-ups
- This PR covers AutoEP + AutoTP folding, including cross-lane EP with
`etp=1`. The replicated-grad reduction is mode-aware so the
sequence-parallel (Ulysses) folding case fits the same contract; AutoTP
+ AutoEP is the validated path here.
- Expert-internal tensor parallelism (`expert_tensor_parallel_size > 1`)
is reserved for a follow-up.
- ZeRO-3 composition with folding is planned as separate follow-up work
(after #8060 is merged).
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
