Overlap AutoEP's expert weight gradients with the DeepEP backward combine by yh0903 · Pull Request #8739 · deepspeedai/DeepSpeed

yh0903 · 2026-10-04T04:23:48Z

The fused gate+up prerequisite, #8733, has merged. This branch is synchronized with current master and retains only the backward-overlap feature delta; no force-rewrite of the earlier submitted commits is required.

In AutoEP's DeepEP backward, the experts compute their input gradient (dX) and their weight gradients (dW). DeepEP's combine then sends dX back to the token ranks. dW and the combine are independent, yet they run serially: a grouped GEMM's backward computes dX and dW together, and only then does the dispatch backward start the combine. In a kernel profile of Qwen3-30B-A3B on 8 H100s, DeepEP ran concurrently with no compute kernel at all.

This adds an opt-in expert_parallel.overlap_weight_grad, which runs the combine and the dW GEMMs side by side.

Change

GroupedExperts.forward takes an optional ExpertWeightGradSlot.
- With a slot, _DeferredWeightGradExperts runs the same forward as the torch._grouped_mm path, for separate or fused gate+up and every registered activation.
- Its backward returns dX at once and leaves the dW GEMMs in the slot.
- The GEMMs are the products autograd forms for torch._grouped_mm (_grouped_mm_mat1_backward, _grouped_mm_mat2_backward).
The DeepEP dispatch takes the expert weights as extra inputs. Its backward:
1. launches the dX combine with async_with_compute_stream=True;
2. runs the slot's dW GEMMs on the compute stream;
3. waits on the combine's event;
4. returns dW through autograd, so ZeRO's gradient handling is unchanged.
SM carveout:
- torch._grouped_mm launches persistent kernels sized to every SM, which leaves the combine nowhere to run.
- The dW GEMMs are therefore launched with comm_num_sm SMs withheld through torch._C._set_sm_carveout_experimental.
- The previous carveout is restored even if the GEMMs raise.
- The pending combine's event is waited on in a finally block as well, so a dW error does not leave the queued communication unaccounted for.
Activation backward:
- The activation registry gains optional standalone backward functions for the fused kernels. swiglu provides them, via new swiglu_backward and swiglu_packed_backward, which the existing autograd functions now also call.
- Forms without them are recomputed and differentiated, so the deferred path always matches the forward that ran.
The default, off, keeps the existing path.
Master's per-layer checkpoint-dispatch replay and exact-size buffer-slice optimization are retained. Overlap slots are separate per invocation and do not replace replay handles.
The deferred custom backward is first-order only: create_graph=True raises rather than returning an incomplete derivative.
overlap_weight_grad is rejected, rather than ignored, unless all of these hold:
- comm_backend="deepep" with autoep_size > 1;
- use_grouped_mm=true, and the experts do not select the Triton grouped GEMM;
- ZeRO stage 1 or 2, explicitly rejecting both stage 0 and stage 3;
- PyTorch provides both the SM carveout getter and setter APIs.
Docs: autoep.rst and the config reference.

Tests

CPU, TestDeferredExpertWeightGrad:
- Output, input gradient and w1/w2/w3 gradients match the plain path's values exactly in the tested CPU cases, for every activation, separate and fused gate+up, FP32 and BF16, with weight gradients immediate or deferred, including an empty expert. This does not assert universal bit-pattern equivalence.
- Also: slot lifecycle errors, the torch._grouped_mm requirement, and config validation for every rejected setting.
CPU, TestWeightGradOverlap, the route against a DeepEP contract buffer:
- Exact numerical agreement with the serial backward in the tested CPU cases, separate and fused gate+up, with no checkpointing, reentrant and non-reentrant checkpointing, now composed with upstream's replay coordinator.
- The order is async combine, then dW with comm_num_sm withheld, then wait.
- The carveout is restored, including on error.
- A dW exception still waits for the async combine before propagating; it does not widen optimizer/backend support.
GPU, H100, 402 passed:
- test_overlapped_weight_gradients_match_the_serial_backward, on 4 GPUs with DeepEP, uniform and skewed routing, separate and fused gate+up: input and every parameter gradient within 1e-2 relative L2, the shared parity checks, and no leaked carveout.
- The CUDA bitwise check with the Triton kernels.
- swiglu_backward and swiglu_packed_backward bitwise against autograd.
- The rest of the DeepEP parity, AutoEP unit and activation suites.
- Not part of this change: the 2 failures are TestMiniMaxM3Block::test_replaced_block_matches_hf_block. That preset needs Transformers >= 5.15.0, and the image has 5.12.0.

End to end, Qwen3-30B-A3B on 8 H100s:

Setup: expert parallel size 8, micro-batch 4, 16 accumulation steps, sequence length 4096, DeepEP, reentrant checkpointing, fused gate+up, 30 steps with 10 warmup.
Measured with an equivalent benchmark patch on our benchmark base: the same deferral, async combine and carveout.
Two pairs in one job:

Arm Mean step

Serial backward 21,569.5 ms

Overlap without the carveout 21,418.5 ms (-0.70%)

Overlap with the carveout 21,097.3 ms (-2.19%)
A kernel profile showed DeepEP running alongside grouped GEMMs for 1,325 ms of the step, against none before.
Step-1 loss identical. Later losses agree within same-arm repeat noise.
Peak allocated memory unchanged; reserved memory about 0.7 GB higher.
A follow-up sweep in one job: withholding more SMs than DeepEP uses (24 with comm_num_sm=12) was 0.9% slower. That supports tying the carveout to comm_num_sm.

The overlap's gain depends on how much backward communication is left to hide. A second job separated the SM budget from the overlap:

Arm	Mean step
Serial backward, `comm_num_sm=12`	21,569.9 ms
Serial backward, `comm_num_sm=20`	20,942.6 ms
Overlap, `comm_num_sm=20`	20,696.9 ms (-1.17% vs serial at 20)
Overlap, `comm_num_sm=28`	20,638.8 ms

On this single node, more DeepEP SMs speed up every DeepEP call regardless of this option. The overlap saves 472 ms at 12 SMs and 246 ms at 20.

Review follow-up

The ZeRO-0 feedback is addressed by failing closed: zero_stage not in (1, 2) is rejected while the disabled default still permits the ordinary stage-0 path.
The engine now passes its actual ZeRO stage into that validator. Previously it omitted the keyword, so tightening the validator alone would incorrectly treat every optimizer as stage 0; both direct-config and engine-wiring cases cover stages 0/1/2/3.
Master synchronization preserves the already merged numerical-reference policy, packed-kernel input-device/stream guard and higher-order rejection. Extracted standalone backward helpers use the input device rather than whichever CUDA device happens to be current.
The real-DeepEP comparison now reads trainable accumulated gradients through safe_get_full_grad before engine.step(), including stage-2 cases where param.grad is released. It does not accept an empty expert-gradient dictionary as proof of parity.
Prior targeted H100 acceptance: 40 passed, 0 skipped, on an 8x NVIDIA H100 80GB HBM3 host (distributed cases use four ranks), driver 580.159.04, torch 2.10.0.2/CUDA 13.0, Triton 3.6.0 and transformers 5.12.0. It covers both supported ZeRO stages, uniform/skewed routing, separate/fused gate+up, replay composition, standalone helper device selection and every trainable gradient before clipping. The 1e-2 relative-L2 criterion is unchanged.
On Apple M4 Pro CPU with torch 2.14.1, the comm suite passes 66/66 in forked mode; targeted config/folding/activation checks pass. Removing only the three guards reproduces all four expected failures (ZeRO 0 accepted, missing event wait, and separate/fused incomplete higher-order derivatives). CPU mocks are not presented as physical DeepEP validation.
Existing 402-passing-case H100 and E2E results above are reused, not described as fresh broad-hardware or full-training validation. The new follow-up retains the old numerical thresholds and performs no throughput sweep.
The activation-weighting child, AutoEP: opt-in routing weights inside the expert activation #8748, is reviewed after this parent and remains a separate opt-in feature; its shifted-rounding semantics are not part of this PR.

AI assistance was used for review and test preparation. The human author reviewed the exact task-specific diff before signoff/publication. This updates an existing PR rather than opening a duplicate feature implementation.

DCO, test collection and changed-files pre-commit passed on the published head. Full GPU CI and infrastructure smoke jobs were skipped; they are not described as a full GPU CI run. On-device evidence for this review follow-up is the targeted H100 execution above.

Latest review follow-up

Applied the requested minimal test-double repair: record_validate accepts and records zero_stage, and the existing sequence-parallel engine test asserts the actual stage is 0. It no longer rejects the engine's keyword argument, and it would still catch missing stage forwarding instead of silently ignoring it. The original test reproduced TypeError: ... unexpected keyword argument 'zero_stage' before the repair.

Synchronized with master to resolve the test-fixture import conflict. The worker-only fixture retains both this PR's standalone SwiGLU backward helpers and master's merged clamped activation helpers and tests. The only manual edits in this follow-up are in those two test files; no manual kernel, autograd, scheduling, numerical-threshold or default changes were made. The activation-weighting child remains separate.

Validation on 1x Apple M4 Pro CPU, macOS 26.7, Python 3.11.15, Torch 2.14.1, pytest 9.1.1; CUDA and accelerator driver are not applicable:

Configuration, deferred expert gradients, ZeRO stage 0/1/2/3 engine wiring, and complete comm/replay coverage: 159 passed, 8 CUDA-only skipped, 49 subtests passed. The repaired sequence-parallel test passes.
The repaired test also passes under --forked. Three selected SwiGLU cases reach the merged worker fixture and skip explicitly because CUDA/Triton is unavailable; this is import/worker coverage, not kernel execution.
All 155 SwiGLU cases collect, including both standalone backward and clamped forms. Changed-files pre-commit passes for all 14 files in the final PR delta.

CPU reproduction command, from the repository root:

DS_ACCELERATOR=cpu PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONDONTWRITEBYTECODE=1 \
  python -m pytest -q -p no:cacheprovider \
  tests/unit/v1/moe/test_autoep_unit.py::TestAutoEPConfig \
  tests/unit/v1/moe/test_autoep_expert_activation.py::TestDeferredExpertWeightGrad \
  tests/unit/v1/moe/test_autoep_expert_activation.py::TestFusedGateUp::test_overlap_weight_grad_settings_are_validated \
  tests/unit/v1/moe/test_autoep_expert_activation.py::TestFusedGateUp::test_overlap_weight_grad_rejects_the_triton_grouped_gemm \
  tests/unit/v1/moe/test_autoep_expert_activation.py::TestFusedGateUp::test_engine_passes_the_actual_zero_stage_to_overlap_validation \
  tests/unit/module_inject/test_auto_ep_comm.py

The prior H100 and E2E results above are reused evidence for the overlap implementation, not a new device run of this merged tree. GPU integration, throughput, memory and convergence experiments were not repeated for this test-only review repair and import-conflict resolution. No additional ZeRO/offload/TP/PP or hardware combinations are claimed.

AutoEP's grouped experts run the gate and up projections of the same input as two grouped GEMMs, and in the backward pass add the two input-gradient GEMM results. expert_parallel.gate_up_impl="fused" concatenates w1 and w3 on each call and runs one torch._grouped_mm with twice the output columns. A single input-gradient GEMM then accumulates both projections in FP32. The activation reads the two halves of the packed output in place: swiglu through a new Triton kernel, swiglu_packed, with the same FP32 math as the existing swiglu kernel; the other registered forms through their eager expression on views. Parameters keep the separate w1/w3 layout, so checkpoints, expert tensor parallelism and optimizer state are unchanged. The default, "separate", keeps the existing path. "fused" is rejected when the experts would not run torch._grouped_mm. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>

…bine In the DeepEP backward, the experts' input gradient is combined back to the token ranks while their weight gradients are independent of it, yet a grouped GEMM's backward computes both before the combine starts. With expert_parallel.overlap_weight_grad, the experts' backward returns the input gradient first and leaves its weight-gradient GEMMs in an ExpertWeightGradSlot. The DeepEP dispatch, which takes the expert weights as extra inputs, launches its backward combine asynchronously, runs those GEMMs on the compute stream with comm_num_sm SMs withheld through PyTorch's SM carveout so the combine has room to run, waits on the combine and returns the weight gradients through autograd. The GEMMs are the products autograd forms for torch._grouped_mm, so gradients are those of the serial backward. The activation registry gains optional standalone backward functions for the fused kernels, provided for swiglu; forms without one are recomputed and differentiated. The option is rejected unless comm_backend is deepep with autoep_size > 1, the experts run torch._grouped_mm, ZeRO stage is 1 or 2, and the SM carveout API exists. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>

…e grouped GEMM in the DeepEP case Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>

0z5a

The documented contract here says overlap_weight_grad is supported with ZeRO stage 1 or 2, but the validation only rejects stage 3, so ZeRO-0 currently passes through.

Could we either explicitly reject zero_stage not in (1, 2), or add ZeRO-0 correctness coverage and document it as supported? Since this path reroutes expert dW through the dispatch autograd node, I'd prefer failing closed for an unvalidated optimizer path.

…dai#8733) AutoEP's grouped experts run the gate (`w1`) and up (`w3`) projections of the same routed rows as two grouped GEMMs. In the backward pass, they compute two input-gradient GEMMs and add the two BF16 results. This adds an opt-in `expert_parallel.gate_up_impl="fused"`, which runs both projections as one `torch._grouped_mm` over the two weights, concatenated on each call. ### Change - `GroupedExperts(gate_up_impl="fused")` concatenates `w1` and `w3` per call and runs one grouped GEMM with `2 * hidden_dim` output columns. - The backward pass therefore has one input-gradient GEMM, which accumulates both projections in FP32, and no separate addition. - The activation reads the two halves of the packed output in place: - `swiglu` uses a new Triton kernel, `swiglu_packed`, with the same FP32 math as the existing `swiglu` kernel. - The other registered forms (`geglu_tanh`, `swiglu_clamped`, `swiglu_oai`) run their eager expression on views of the halves. - New forms can register a `packed_fused_fn`. - The parameters keep their separate `w1`/`w3` layout. Checkpoints, universal checkpoints, expert tensor parallelism and optimizer state, including Muon's per-expert treatment, are unchanged. - The default, `"separate"`, keeps the existing path. - `"fused"` is rejected rather than ignored when the experts would not run `torch._grouped_mm`: - `use_grouped_mm=false`, at config validation and in `GroupedExperts`; - devices that select the Triton grouped GEMM, unless `disable_triton_grouped_mm=true`. - Docs: `autoep.rst` and the config reference. ### Numerical and autograd contract The forward expression and parameter-gradient semantics are unchanged. Exact output/weight-gradient equality is an observation in the tested H100 cases, not a guarantee across GEMM shapes, implementations or hardware. The input gradient can differ because it is accumulated in one FP32 GEMM rather than as two rounded BF16 results added. The FP64 regression now uses the same represented BF16 weights as the tested paths, rather than the original pre-quantization FP32 weights; its existing relative-L2 bound (fused error <= 1.05 * separate error) is unchanged. This is not a universal accuracy-improvement or bitwise-equivalence claim. Cross-GEMM output and parameter-gradient assertions now use the existing grouped-GEMM BF16 policy (`rtol=0.02`, `atol=0.02 * max(abs(reference))`) rather than zero tolerance. The skewed/empty-expert forward and smooth-form parameter gradients use the independent FP64 reference. Clamped-form parameter gradients use an independent dense BF16 reference, because rounding a projection across a clamp boundary changes its derivative discontinuously; the bound is not enlarged to absorb those discontinuities. The input-gradient relative-L2 comparison is unchanged. The multi-step test uses the same policy for comparisons across algorithms, while saved/reloaded parameters, optimizer momentum tensors and parameter-group configuration must still match exactly. Exact parity checks between the two implementations of the same elementwise SwiGLU kernel remain unchanged. The packed SwiGLU Triton path has first-order gradients only: `create_graph=True` raises rather than silently dropping its second derivative. Forward and backward select the input's CUDA device and its current stream, restoring the caller's current device. The eager expressions for the other activation forms and the default separate path are unchanged. ### Tests - CPU, `TestFusedGateUp` in `tests/unit/v1/moe/test_autoep_expert_activation.py`: - packed versus separate halves for every registered form, forward and backward; - only `swiglu` has a packed kernel; - the default is unchanged; - unsupported settings are rejected (bad value, `use_grouped_mm=false`, Triton grouped GEMM path), at config validation and in `GroupedExperts`; - the config reaches the replaced experts, whose parameters remain `w1`, `w2`, `w3`. - Earlier GPU evidence, on one H100 (305 passed): `test_swiglu_triton.py`, `test_autoep_expert_activation.py` and `test_autoep_unit.py`. - `swiglu_packed` matches the existing `swiglu` kernel's values exactly, forward and backward (`rtol=0`, `atol=0`, not a signed-zero bit-pattern assertion). Covered: FP32/FP16/BF16, odd widths, 3-D inputs, empty inputs, non-contiguous input, and an odd last dimension (rejected). - The fused path joins `test_grouped_gemm_paths_match_the_for_loop` for every form. - New `test_matches_separate_with_skewed_and_empty_experts`, for every form, with an empty expert and unaligned rows: - output and `w1`/`w2`/`w3` gradients match exactly in the tested shapes; - the original relative-L2 FP64 error comparison is retained with the corrected represented-weight reference described above. - Not part of this change: the 2 failures are `TestMiniMaxM3Block::test_replaced_block_matches_hf_block`. That preset requires Transformers >= 5.15.0, and the image has 5.12.0. - The shipped `GroupedExperts`, one H100, Qwen3-30B-A3B per-rank expert shape (131,072 routed rows, hidden 2048, expert hidden 768, 16 local experts, BF16): | Per call (ms) | Forward, no grad | Forward | Backward | Total | |---|---:|---:|---:|---:| | `"separate"` | 2.149 | 2.292 | 4.857 | 9.298 | | `"fused"` | 2.224 | 2.236 | 4.321 | 8.781 | Output and all weight gradients matched exactly; the input gradient differed by 2.9e-3 relative. - End to end, Qwen3-30B-A3B on 8 H100s: - Setup: expert parallel size 8, micro-batch 4, 16 accumulation steps, sequence length 4096, DeepEP, reentrant activation checkpointing, BF16 parameters and Adam state, 30 steps with 10 warmup. - Measured with an equivalent benchmark patch on our benchmark base: the same concatenation and the same packed SwiGLU kernel. - Separate 22,018.8 / 22,018.5 ms against fused 21,601.2 / 21,608.4 ms, in two pairs in the same job: **-414 ms (-1.9%)**. - Step-1 loss identical. Later losses agree within same-arm repeat noise. - Peak allocated memory unchanged; reserved memory about 0.4 GB higher. ### Review follow-up validation - Current-master targeted validation on NVIDIA H100 80GB HBM3, torch 2.10/CUDA 13.0, Triton 3.6.0 and transformers 5.12.0: **83 passed, 0 skipped** (the full 63-case SwiGLU file plus 20 `TestFusedGateUp` cases). The SwiGLU file also passed **63/63 with no skips under `pytest --forked`**. Cross-device cases used two GPUs; the lifecycle cases use the actual native grouped GEMM. - With only the same new tests applied to the prior runtime, all four focused guards fail as intended: both FP16/BF16 cross-device cases reproduce Triton's inaccessible-pointer error, and both FP32/BF16 higher-order cases fail because `create_graph=True` was not rejected. This verifies the actual root causes, not an unrelated collection or test-API failure. - Regressions cover strided packed inputs and upstream gradients, two-device/current-stream selection and caller restoration, explicit higher-order rejection, the corrected FP64 reference, and two SGD updates with two accumulation micro-batches per step. All four activations run with and without reentrant recomputation; a model/optimizer state-dict round trip between updates keeps the separate `w1`/`w2`/`w3` names and shapes and checks the next update against the independent separate path. - The SwiGLU test file imports its runtime in the worker, not during collection, so its Triton driver probe is deferred until after `pytest --forked` starts the worker. - The follow-up changes launch-device selection and unsupported-autograd handling, not the GEMM implementation, kernel arithmetic, launch configuration, default or parameter layout. The earlier local and E2E performance evidence is reused without retiming; the historical percentages are not new measurements of these guards. ### AI review follow-up - Changed tests and documentation only; the runtime is byte-identical to the previous validated head. - The numerical-contract comment is addressed by the reference/tolerance policy above, plus a regression that accepts a single BF16 representable-step change but rejects a clearly incorrect value. On NVIDIA H100 80GB HBM3, torch 2.10/CUDA 13.0, Triton 3.6.0 and transformers 5.12.0, the focused fused-gate-up/grouped-GEMM suite passed **33/33 with no skips**. - A synthetic one-ULP perturbation of the fused down-GEMM output passed all **12/12** reference/lifecycle cases with the contract-level checks; the old exact checks failed all **12/12** under the same perturbation. An independently injected 25% output/gradient error was rejected by all **four** reference cases. These are small test controls, not new GEMM-backend or performance measurements; they demonstrate that legitimate rounding is allowed without accepting an incorrect computation. - The DCO comment was checked against the actual GitHub PR commit list and local ancestry: every PR commit has the configured author's `Signed-off-by` trailer, and DCO passes. The SHA cited in that comment is not in the PR and is not found by GitHub's commit API, so no commit amendment or history rewrite is needed. ### CI and scope DCO, test collection and changed-files pre-commit passed on the published follow-up. Full GPU CI and infrastructure-smoke jobs were skipped on the PR event; they are not reported as full GPU CI acceptance. The targeted H100 runs above provide the device evidence for this follow-up. This branch inherits recent upstream `master`. The backward-overlap child (deepspeedai#8739) remains a separate stacked PR and is not included here; no recompute-skip, buffer-slice, row-weighting-activation or communication optimization is added by this PR. The new optimizer round trip is a focused SGD test, not an independent qualification of every optimizer, expert-TP/universal-checkpoint combination, ZeRO stage or accelerator. The E2E evidence covers one model/configuration on H100, not long-run convergence or a universal speedup. --------- Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Preserve merged gate-up, checkpoint replay and exact-size buffer behavior. Enforce ZeRO stages 1/2 and forward the actual engine stage into validation; keep the disabled default unchanged. Restore SM carveout and wait for the async combine even if dW fails, reject incomplete higher-order derivatives, and retain input-device standalone backward launches. Read accumulated parameter gradients before clipping in the real stage-1/2 regression rather than accepting released param.grad as missing coverage. The human author reviewed this exact staged tree before signoff; targeted H100 acceptance and CPU controls passed without a new performance sweep. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>

chatgpt-codex-connector · 2026-10-13T03:10:25Z

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-13T03:19:57.663917Z	`b016772`	Draft marked ready

ℹ️ 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.

Copilot

🔵 Needs a closer look

The distributed asynchronous autograd scheduling and private PyTorch SM-carveout integration warrant final human review despite extensive targeted testing.

0 open findings

What changed in this PR

Adds opt-in overlap of AutoEP expert weight-gradient GEMMs with DeepEP’s backward combine, building on the merged fused gate/up work while keeping activation weighting separate.

Changes:

Adds deferred expert weight-gradient computation with async DeepEP combine and SM carveout.
Validates supported DeepEP, grouped-GEMM, ZeRO, and PyTorch configurations.
Adds activation backward helpers, documentation, and CPU/H100 coverage.

File	Description
`deepspeed/moe/ep_experts.py`	Implements deferred expert weight gradients.
`deepspeed/module_inject/auto_ep_comm.py`	Overlaps combine with gradient GEMMs.
`deepspeed/module_inject/auto_ep_config.py`	Validates overlap requirements.
`deepspeed/module_inject/auto_ep_layer.py`	Connects overlap to AutoEP routing.
`deepspeed/module_inject/auto_ep_presets/base.py`	Adds the configuration option.
`deepspeed/ops/triton_ops/swiglu_triton.py`	Exposes standalone SwiGLU backward helpers.
`deepspeed/runtime/engine.py`	Passes the actual ZeRO stage.
`docs/code-docs/source/autoep.rst`	Documents behavior and performance.
`docs/_pages/config-json.md`	Documents the configuration field.
`tests/unit/module_inject/test_auto_ep_comm.py`	Tests scheduling, recovery, and checkpointing.
`tests/unit/v1/moe/test_autoep_deepep_parity.py`	Tests real DeepEP gradient parity.
`tests/unit/v1/moe/test_autoep_expert_activation.py`	Tests deferral and validation.
`tests/unit/v1/ops/triton_ops/test_swiglu_triton.py`	Tests backward helper parity and device selection.

🧠 Review effort: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

yh0903 · 2026-10-13T03:29:11Z

The documented contract here says overlap_weight_grad is supported with ZeRO stage 1 or 2, but the validation only rejects stage 3, so ZeRO-0 currently passes through.

Could we either explicitly reject zero_stage not in (1, 2), or add ZeRO-0 correctness coverage and document it as supported? Since this path reroutes expert dW through the dispatch autograd node, I'd prefer failing closed for an unvalidated optimizer path.

Thanks @0z5a. Overlap now explicitly rejects ZeRO 0/3 and accepts only stages 1/2, with the engine forwarding its actual stage to validation. Error paths restore the carveout and wait for the pending async combine!

tohtana

Hi @yh0903,
Thank you for submitting this PR!

I wasn't able to run this with real DeepEP setting, but the code change overall looks good.
One small change I want to request: the engine now passes zero_stage to the AutoEP config validator, but an existing validator test func rejects that keyword. A minimal signed test-only repair would be like commit e452d350.

Retain the merged clamped activation kernels and tests while resolving the worker-fixture imports as a union with the standalone backward helpers. Keep both sets of imports inside the worker, preserving fork-safe collection. Accept and record zero_stage in the existing sequence-parallel validator test double and assert the actual stage is forwarded. This fixes the reproduced unexpected-keyword failure without changing the product runtime or weakening existing assertions. The human author reviewed the exact staged tree before signoff. Targeted CPU configuration, deferred-gradient, engine and comm/replay checks and changed-files pre-commit passed. Reuse earlier device evidence without claiming a fresh GPU or end-to-end run. Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

yh0903 · 2026-10-27T22:35:45Z

Thanks, Masahiro! Applied the minimal test-double repair you suggested in 894a116. record_validate now accepts and records zero_stage, and the existing sequence-parallel test explicitly asserts that the engine forwards stage 0. Existing sequence-parallel assertions and ZeRO stage 0/1/2/3 coverage are preserved.

Also synchronized with master, resolving the SwiGLU fixture conflict by retaining both the standalone backward and clamped activation helpers inside the worker-only fixture. This follow-up introduces no manual product-runtime changes.

yh0903 and others added 3 commits October 3, 2026 01:11

Run the overlap tests' CPU cases where Triton is installed, and enabl…

56dac2a

…e grouped GEMM in the DeepEP case Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>

yh0903 mentioned this pull request Oct 5, 2026

AutoEP: opt-in skip of the recomputed DeepEP combine #8745

Open

0z5a reviewed Oct 6, 2026

View reviewed changes

yh0903 mentioned this pull request Oct 10, 2026

Add an opt-in fused gate+up grouped GEMM for AutoEP experts #8733

Merged

yh0903 marked this pull request as ready for review October 13, 2026 03:10

yh0903 requested review from hwchen2017, loadams, tjruwase and tohtana as code owners October 13, 2026 03:10

yh0903 mentioned this pull request Oct 13, 2026

AutoEP: opt-in routing weights inside the expert activation #8748

Draft

hwchen2017 requested a balanced review from Copilot October 13, 2026 03:11

Copilot started reviewing on behalf of hwchen2017 October 13, 2026 03:12

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

tohtana requested changes Oct 14, 2026

View reviewed changes

0z5a mentioned this pull request Oct 14, 2026

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

Open

pengdurice mentioned this pull request Oct 16, 2026

AutoEP: support combine_impl="fused_weighted_sum" with AutoEP+AutoTP folding #8847

Draft

Navigation Menu

Repository navigation

Overlap AutoEP's expert weight gradients with the DeepEP backward combine - #8739

Overlap AutoEP's expert weight gradients with the DeepEP backward combine#8739
yh0903 wants to merge 5 commits into
deepspeedai:masterfrom
yh0903:yh0903/autoep-overlap-weight-grad

yh0903 commented Oct 4, 2026 •

edited

Loading

Uh oh!

0z5a left a comment

Uh oh!

chatgpt-codex-connector Bot commented Oct 13, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

yh0903 commented Oct 13, 2026

Uh oh!

tohtana left a comment

Uh oh!

yh0903 commented Oct 27, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Arm	Mean step
Serial backward	21,569.5 ms
Overlap without the carveout	21,418.5 ms (-0.70%)
Overlap with the carveout	21,097.3 ms (-2.19%)

Conversation

yh0903 commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Change

Tests

Review follow-up

Latest review follow-up

Uh oh!

0z5a left a comment

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

yh0903 commented Oct 13, 2026

Uh oh!

tohtana left a comment

Choose a reason for hiding this comment

Uh oh!

yh0903 commented Oct 27, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yh0903 commented Oct 4, 2026 •

edited

Loading

chatgpt-codex-connector Bot commented Oct 13, 2026 •

edited

Loading