DeepSpeed

7d77213d - AutoEP + AutoTP folding under ZeRO-3 (#8670)

AutoEP + AutoTP folding under ZeRO-3 (#8670) # AutoEP + AutoTP folding under ZeRO-3, and correct gradients for tensor-parallel layers under folding ## Summary * **ZeRO-3 support.** AutoTP and AutoEP can now share the same ranks under ZeRO stage 3. AutoTP shards the layers its partition specification names (the tests shard attention only), AutoEP distributes the routed experts, and `expert_tensor_parallel_size` stays 1. This also works for models built under `deepspeed.zero.Init`. * **A gradient bug in folding at every ZeRO stage, fixed.** When a layer below an MoE layer is really tensor-parallel, that layer and every parameter below it get wrong gradients. On unmodified `master` at ZeRO-2, 36 of the 45 gradients of a small GLM-4.5-Air-shaped model differ from plain data parallelism by 1.0 % to 5.5 %. The existing folding tests shard nothing ("skip" for every weight), so they could not see it. * **Verified.** A new test compares every gradient and the gradient norm with plain data parallelism, at ZeRO-2 and ZeRO-3, on 4 and 8 ranks, with and without `zero.Init`. On GLM-4.5-Air with 23 of its 46 layers (53.0B parameters) on 8 H200s, the folded ZeRO-3 run gives the same losses and gradient norms as the unfolded run on the same samples (largest norm difference 0.1 %). ## Problem 1: ZeRO-3 was rejected Two checks rejected AutoEP + AutoTP under ZeRO-3: * `validate_folding_global` (`deepspeed/module_inject/auto_ep_folding.py`): "AutoEP+AutoTP with ZeRO stage 3 is reserved for the separate ZeRO-3 composition lane". * `DeepSpeedEngine._validate_zero3_moe_compatibility`: "AutoEP with ZeRO Stage 3 does not support AutoTP yet". Removing them is not enough. ZeRO-3 did not apply the per-family gradient corrections folding needs, counted the gradient norm wrongly, and AutoTP could not shard a model built under `zero.Init` (see "The change"). ## Problem 2: wrong gradients below every folded MoE layer (all ZeRO stages) In the folded MoE layer each TP peer dispatches only its share of the (token, expert) assignments, and `restore_combined` all-gathers the results back into the full output. The backward of that all-gather (`_AllGatherVariableRows`) all-reduces the replicated output gradient, so each peer gets `tp_size` times the gradient of the assignments it kept. The gradient the layer returns for its input is therefore, on peer `p`: ``` dx_p = tp_size * E_p + S E_p : the routed-expert and router part from peer p's own assignments S : the part every peer computes alike (the shared expert) true gradient = E + S = average over the peers of dx_p ``` No peer holds the true gradient; only their average is right. * **Replicated parameters recover it.** Their gradients are linear in `dx`, and the folding reduction averages them over the TP group, which gives the gradient of the averaged `dx`. This is why the existing folding tests pass: every weight there is replicated. * **Tensor-parallel layers cannot.** A row-parallel layer (`o_proj`) computes its shard's weight gradient from its own `dx_p`, and its gradient is not reduced over the TP group. A column-parallel layer (`q_proj`, `k_proj`, `v_proj`) sums the peers' partial input gradients, each computed from a different `dx_p`. Both are wrong, and every parameter below them inherits the error. Measured with the new test on unmodified `master` at ZeRO-2 (4 ranks, TP 2 x EP 2): 36 of 45 gradients differ, relative errors 1.0 % to 5.5 %, scale 0.995 to 1.021 (wrong direction, not a scale factor). The 9 correct gradients are exactly those that do not pass back through a tensor-parallel layer after leaving an MoE layer: `lm_head`, the final norm, the last layer's MoE block, and the norm that feeds it. **Fix:** `average_gradient_over_tp` (`deepspeed/moe/ep_tp_dispatch.py`) is an identity in forward and averages the gradient over the TP group in backward. The folded MoE layer applies it to its input, so every peer returns the true input gradient. The cost is one all-reduce of `[tokens, hidden]` over the TP group per MoE layer in backward, the same size as the one a column-parallel layer already does. The router still needs the AVERAGE reduction and the routed experts the division by `tp_size`, because those corrections are about gradients computed inside the layer. ## The change | file | change | |---|---| | `deepspeed/moe/ep_tp_dispatch.py` | `average_gradient_over_tp`: identity forward, TP average in backward | | `deepspeed/module_inject/auto_ep_layer.py` | the folded layer averages its input gradient over the TP group | | `deepspeed/module_inject/auto_ep_folding.py` | ZeRO-3 no longer rejected; offload, DeepCompile and `expert_tensor_parallel_size > 1` still are | | `deepspeed/runtime/zero/stage3.py` | per-parameter folding correction before a gradient enters the reduce bucket, as ZeRO-2 does; gradient norm counted once per distinct partition (below) | | `deepspeed/runtime/engine.py` | AutoTP check removed from `_validate_zero3_moe_compatibility`; ZeRO-3 partition groups resolved whenever AutoTP runs; caller's optimizer remapped after AutoTP; the engine's fallback correction skips ZeRO-3 | | `deepspeed/module_inject/auto_tp.py` | AutoTP gathers weights partitioned by `zero.Init` before sharding them | | docs (`autoep.rst`, `config-json.md`) | folding works at ZeRO stages 0 to 3; the old "cannot be combined with AutoTP" text removed | Details: * **Gradient correction in ZeRO-3.** `DeepSpeedZeroOptimizer_Stage3` gets `configure_autoep_folding_tp_gradient_reduction`, which the engine already calls on any optimizer that has it. Each micro-batch gradient is corrected with the existing `apply_folding_correction_to_grad_buffer` before it is reduce-scattered: AVERAGE over TP for the router and other replicated parameters, a division by `tp_size` for routed experts, nothing for TP-sharded parameters. The engine's own after-backward correction now skips ZeRO-3: a gradient still waiting in the reduce bucket has been corrected already, and a second pass would divide expert gradients by `tp_size` twice. * **Gradient norm in ZeRO-3.** ZeRO-3 counts a parameter that is not tensor-parallel only on model-parallel rank 0, then sums over the partition group and the model-parallel group. That assumes each TP rank holds its own copy partitioned over the data-parallel group. Two cases break it, and `_grad_norm_counts_every_rank` detects them: * routed experts under folding: the expert-parallel group spans the TP ranks, so expert-data-parallel x expert-parallel already covers every rank. Counting only on TP rank 0 and then summing over TP misses half the experts and doubles the rest (8 ranks, TP 2 x EP 4: the result is `2 * (norms on ranks 0, 2, 4, 6)`); * parameters partitioned over a group that contains the TP peers, as `zero.Init` partitions every parameter over all ranks before AutoTP creates the data-parallel group. In both cases every rank counts its partition and the model-parallel all-reduce is skipped. * **AutoTP with `zero.Init`.** A weight partitioned by `zero.Init` holds an empty 1-D placeholder, so AutoTP both skipped the layer (its "is this a 2-D linear weight" check saw 1 dimension) and would have sliced the placeholder. AutoTP now reads `ds_shape` for that check, gathers the weight once, and gives the layer an ordinary parameter holding the full value to slice. ZeRO-3 partitions the new shard over the data-parallel group (`_resolve_zero3_param_placement` now also runs when AutoTP is on). The replaced parameters are returned as a `ReplacementSourceMap`, and the engine re-points a caller-built optimizer or parameter list with the existing AutoEP remap functions. The `use_muon` and `muon_num_heads` tags are carried over. * **A hang avoided.** `deepspeed.comm.get_all_ranks_from_group` stops only when `get_global_rank` raises, and torch never raises for the world group, so it never returns for that group. The norm code enumerates ranks bounded by the group size instead. `get_all_ranks_from_group` itself is unchanged here. ## Tests `tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py` (new). The model is a small transformers `glm4_moe` (the GLM-4.5-Air architecture): attention with biases, one dense MLP layer, two MoE layers with a sigmoid router and a shared expert. AutoTP shards only `q_proj`, `k_proj`, `v_proj` (column) and `o_proj` (row); AutoEP replaces the MoE blocks. The reference is the same model, float32, plain data parallelism at ZeRO-0, reading the same micro-batches (2 per rank, gradient accumulation 2). The test checks: * the structure: exactly the attention projections are tensor-parallel, every MoE block is an AutoEP layer; * the losses of both micro-batches; * every one of the 45 gradients, rebuilt at the reference's names and shapes (TP shards gathered, experts gathered over the expert-parallel group), to 1e-4 relative; * the gradient norm against the norm of the reference gradients, to 1e-4 relative, at ZeRO-3. ZeRO-2 takes an approximate norm for MoE parameter groups with or without folding (`_average_expert_grad_norms` averages the expert groups' norms over the expert-data-parallel group), so its norm is printed but not asserted. Folding only needs TP and EP to divide the world size (dense layers use TP x world/TP, experts EP x world/EP), so the test covers TP x EP smaller than, equal to and larger than the world size: | ranks | TP | EP | dense data parallel | expert data parallel | shape | ZeRO-2 | ZeRO-3 | ZeRO-3 under `zero.Init` | |---|---|---|---|---|---|---|---|---| | 4 | 2 | 2 | 2 | 2 | TP x EP = world | pass | pass | pass | | 4 | 2 | 4 | 2 | 1 | each EP group spans both TP groups | pass | pass | pass | | 4 | 4 | 2 | 1 | 2 | all ranks are one TP group | pass | pass | pass | | 4 | 4 | 4 | 1 | 1 | the EP group is the TP group | pass | pass | pass | | 8 | 2 | 2 | 4 | 4 | TP x EP < world | pass | pass | pass | | 8 | 2 | 4 | 4 | 2 | TP x EP = world | pass | pass | pass | | 8 | 4 | 2 | 2 | 4 | TP x EP = world, larger TP | pass | pass | pass | | 8 | 4 | 4 | 2 | 2 | TP x EP = 2 x world | pass | pass | pass | | 8 | 2 | 8 | 4 | 1 | one expert per rank | pass | pass | pass | 27 passed in 741 s on one node of 8 H200. The earlier two-shape version of this file (4 ranks TP 2 x EP 2, 8 ranks TP 2 x EP 4) on unmodified `master`, ZeRO-2 cases only (`master` rejects ZeRO-3): both fail, 36 of 45 gradients differ. `test_autoep_autotp_folding_config.py`: the rule test that asserted the ZeRO-3 rejection now asserts that ZeRO-3 is accepted. The duplicate rejection test in `test_autoep_autotp_grad_parity.py` is removed. Existing tests, rerun on this branch: `test_autoep_autotp_grad_parity.py`, `test_autoep_autotp_folding_config.py`, `test_autoep_autotp_runtime.py`, `test_autoep_autotp_dispatch.py`, `test_autoep_autotp_checkpoint.py`, `test_autoep_grad_parity.py`, `unit/v1/autotp/test_autotp_training.py`: 123 passed, 0 failed (17 skipped and 29 deselected by their own conditions). Hardware: one node of 8 H200, torch 2.8.0+cu126, transformers 5.15.1. `pre-commit run --files` passes on every changed file. ## GLM-4.5-Air, 1 node of 8 H200 Config-only GLM-4.5-Air (random weights) built under `zero.Init`. Sequence 4,096, micro-batch 1, bf16 with bf16 optimizer states, AdamW at learning rate 1e-5, activation checkpointing, ZeRO-3 with `leaf_module` on the AutoEP layer, AutoEP preset `deepseek_v3`, AutoTP on the attention projections only. Both runs of a pair read the same 4 sequences per step: the folded run's 4 dense data-parallel ranks and the unfolded run's 4 ranks take the same samples. The loss is the running mean of the per-step losses since step 1, as the benchmark prints it. | layers (parameters) | run | GPUs | step 1 loss | grad norm, steps 1 / 2 / 3 / 4 | s / step | peak allocated / reserved | |---|---|---|---|---|---|---| | 8 (17.9B) | folded, ZeRO-3, TP 2 x EP 4 (this PR) | 8 | 12.738950 | 2.9206 / 2.9031 / 2.9248 / 2.9810 | 0.727 | 32.5 / 36.4 GiB | | 8 (17.9B) | unfolded, ZeRO-3, TP 1 x EP 4 | 4 | 12.739278 | 2.9212 / 2.9036 / 2.9242 / 2.9814 | 0.541 | 52.3 / 55.6 GiB | | 23 (53.0B) | folded, ZeRO-3, TP 2 x EP 4 (this PR) | 8 | 12.743540 | 2.6618 / 2.6882 / 2.7329 / 2.7304 | 1.774 | 73.9 / 76.9 GiB | | 23 (53.0B) | unfolded, ZeRO-3, TP 1 x EP 4 | 4 | 12.743791 | 2.6635 / 2.6875 / 2.7355 / 2.7302 | 1.849 | 134.5 / 134.9 GiB | * The running-mean losses agree to 1.2e-3 or better and the gradient norms to 0.1 % or better at every step, the size of bf16 rounding. * For comparison, an earlier run of the unfixed ZeRO-2 folding path on the same 8-layer model and samples (TP 4 x EP 16 on 16 GPUs, attention, dense MLP and shared expert sharded) had a step-1 gradient norm of 3.2311 against 2.9245 for the unfolded run, 10.5 % higher. * These runs check correctness, not speed. The folded run on 8 GPUs is slower than the unfolded run on 4 GPUs at 8 layers (0.727 against 0.541 s per step) and 4 % faster at 23 layers. The unfolded 23-layer run reserves 134.9 of 139.8 GiB per GPU, the folded one 76.9 GiB. --------- Signed-off-by: pengdurice <pengduhit@gmail.com>