[yh0903]

Summary

Add opt-in expert_parallel.skip_recompute_combine for AutoEP's DeepEP route. In supported Hugging Face per-decoder-layer reentrant checkpointing, the recomputed MoE output is used only by the residual add and is discarded by the checkpoint. The route records post-expert row weighting and combine backward without executing those forward operations; dispatch, expert compute and all required backward transfers still run. Default behavior is unchanged.

This is the existing skip-combine PR, not a duplicate implementation. It remains distinct from the non-reentrant layout replay merged in #8731. The branch preserves that replay, the fused gate+up defaults from #8733, and the independent exact-size buffer checks from #8747. AI assistance was used for finalization, with author review before submission.

Supported contract

• comm_backend="deepep", autoep_size > 1, BF16 dispatch; eager and fused row weighting, at their unchanged score_apply boundary.
• Supported decoder presets: deepseek_v2, deepseek_v3, minimax_m3, mixtral, qwen3_5_moe, qwen3_moe, checked against Transformers 5.12.0 modeling code.
• The recompute guard checks the attached decoder layer's checkpoint configuration: gradient_checkpointing=True and the standard Torch checkpoint partial with use_reentrant=True. Transformers 5 defaults to non-reentrant, so explicitly use gradient_checkpointing_kwargs={"use_reentrant": True}.
• Any non-null moe_layer_pattern is rejected during configuration validation when the skip is enabled, even with an explicit or automatically selected supported preset, a pattern matching the preset, or an empty string. Leave the pattern unset or disable the skip; custom-pattern support is unchanged with the option off.
• Other checkpoint APIs, non-reentrant checkpointing with this flag and unsupported presets are outside the supported contract. Without recomputation, the ordinary path runs.
• The skip requires that the MoE output has no value-dependent consumer in the recompute. The combined sentinel is NaN; an invalid consumer can produce non-finite values or gradients, but the sentinel is not a general runtime proof that arbitrary wrappers, hooks or patched decoder graphs are safe.

Finalization

Resolved the upstream conflicts without dropping replay or buffer behavior. Both normal and recorded combines now share their backward-context setup, including transferring the pending dispatch entry from the exchange to the consuming autograd node. Otherwise the recorded path could retain a stale entry and attempt to replay the previous step's routing in a later reentrant backward.

The CPU reproducer fails before the context fix with activation checkpointing recomputed a different routing. The new regression performs three successive steps with changing routes and compares outputs, input/router/expert gradients and required transfers for both pre/post score application.

The single CUDA weighting test now runs in a fresh DistributedTest worker, retaining its actual fused/recorded-backward comparison rather than allocating CUDA from an ordinary forked unittest. The existing real-DeepEP skip parity cases run two steps per model. The one-step test helper's defaults and original tolerances remain unchanged.

Review follow-up

Following the review, reject pattern overrides before model conversion. Applying an override preserves the supported preset name but can select a different decoder parent, so a preset-name whitelist cannot establish the residual-only output-consumer contract. The fix is a four-line configuration guard, not a decoder-graph validator or a runtime fallback.

Six parametrized cases cover automatic and explicit preset selection, a custom parent, the preset's own pattern and an empty pattern. Each also checks that disabling the skip or leaving the pattern unset remains accepted. Both configuration references document the restriction. The kernels, recorded backward, supported execution paths and numerical thresholds are unchanged.

Testing

Post-review CPU validation: 139 passed, 1 CUDA test skipped, 52 subtests passed, 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. This covers the complete comm/replay suite, configuration tests and supported decoder attachment:

DS_ACCELERATOR=cpu PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONDONTWRITEBYTECODE=1 \
  python -m pytest -q -p no:cacheprovider \
  tests/unit/module_inject/test_auto_ep_comm.py \
  tests/unit/v1/moe/test_autoep_unit.py::TestAutoEPConfig \
  tests/unit/v1/moe/test_autoep_unit.py::TestModelDetectionAndReplacement::test_skip_recompute_combine_attaches_each_layer_to_its_decoder_layer

The skipped case requires CUDA and Triton for recorded fused weighting; its existing device evidence is listed below. The six added pattern regressions also passed under isolated forked workers (6 passed, 0 skipped) on the same CPU setup:

DS_ACCELERATOR=cpu PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONDONTWRITEBYTECODE=1 \
  python -m pytest -q -p no:cacheprovider -p pytest_forked --forked \
  tests/unit/v1/moe/test_autoep_unit.py::TestAutoEPConfig::test_skip_recompute_combine_rejects_pattern_overrides_even_with_supported_presets

Existing product validation is bound to the original PR commit 29303c378: 275 passed, 1 skipped on four H100s. All four eager/fused and uniform/skewed skip-parity cases ran, as did the recorded fused-weighting backward test. The skipped case is a MiniMax-M3 router-logit test requiring Transformers >=5.15.0; the validation environment has 5.12.0. This is reused evidence for the original implementation, not a claim that the entire suite was rerun after rebase.

Pre-review focused CPU validation: 75 passed, 1 CUDA test skipped, under both ordinary pytest and pytest --forked, on macOS with Torch 2.10.0 and Transformers 5.12.0. This covers the comm/replay suite, three skip config/conversion cases and the one-step helper defaults. Ordinary pytest also reports 52 subtests passed.

Most recent focused device validation, before the configuration-only review follow-up: 9 passed, 0 skipped with pytest --forked, on 4x NVIDIA H100 80GB HBM3, driver 580.159.04, Torch 2.10.0.2+cu130 / CUDA 13.0, Triton 3.6.0, Transformers 5.12.0, using the previously validated DeepEP runtime:

• All four eager/fused and uniform/skewed skip-parity cases run two training steps on each model. They retain the original numerical bounds, check last-step output/input/router/parameter gradients and accumulated parameter deltas, and verify one forward combine per layer per step rather than two.
• The four upstream non-reentrant/nested-reentrant replay cases still pass, confirming that the shared context preserves normal replay behavior.
• The recorded fused-weighting backward comparison runs in a fresh worker. The original ordinary unittest reproduces Cannot re-initialize CUDA in forked subprocess on the same setup; the new worker-based case passes.

GPU reproduction selector, from tests/:

PYTHONPATH=.. DEEPSPEED_RUN_H100_TESTS=1 pytest --forked \
  unit/module_inject/test_auto_ep_comm.py unit/v1/moe/test_autoep_deepep_parity.py \
  -k 'skipping_the_recomputed_combine or recorded_fused_row_weighting or non_reentrant_checkpointing_replays'

Changed-files pre-commit passed. Existing device evidence is reused for unchanged supported paths; GPU tests, full E2E and the existing performance matrix were not restarted for the configuration-only follow-up. The focused tests do not establish every ZeRO/offload/TP/PP combination, other hardware, arbitrary checkpoint composition or long-run convergence.

The static CI selector includes the v1 AutoEP unit and parity files. Real DeepEP cases remain gated on the opt-in H100 test flag, four GPUs and a GIN-capable DeepEP build, so test collection alone does not prove they ran. The standalone recorded-weighting CUDA comparison is outside Modal's v1 test scope and was explicitly included in the focused hardware command above; its recorded path is also exercised by the v1 model-parity cases. Full GPU CI follows the repository's merge-queue workflow.

Existing performance evidence

Qwen3-30B-A3B, 8 H100s, EP8, micro-batch 4, gradient accumulation 16, sequence 4096, DeepEP 28 SMs, reentrant activation checkpointing; 30 steps with 10 warmup. The equivalent benchmark prototype records the same weighting/combine backward but returns uninitialized output where the product fills NaN. Both arms inherit concatenated gate+up and backward weight-gradient overlap; this is the incremental skip effect, not their combined benefit.

One allocation, order normal / skip / skip / normal:

Recompute	Run medians	Mean
Normal	20,575.1 / 20,574.2 ms	20,574.7 ms
Combine + post-weighting forward skipped	19,945.7 / 19,928.2 ms	19,936.9 ms (-3.10%)

Peak allocated memory stayed at 58.69 GB; peak reserved memory was 0.25 GB lower. Loss differences over these 30 steps were within same-arm repeat noise (maximum 9.4e-4). The 8-GPU residual-block gate had numerical equality under torch.equal for output/input/router gradients and expert-gradient differences around 5e-5 relative L2, within two normal recomputes' variation. torch.equal is not a claim about signed-zero bit patterns.

The measurements are historical prototype evidence, not a retiming of the rebased product or a universal benefit across models/hardware. Skipping the forward combine removes that recompute transfer; it does not add overlap or remove the cached dispatch and weight-gradient communication required by backward.

[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-12T23:46:05.746814Z	42243d7	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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 42243d788f

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

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Reject custom MoE patterns before enabling combine skipping

With a supported preset_model (or auto-detected model_type) plus a moe_layer_pattern override, resolve_preset_candidates() applies the override while preserving the supported preset name, so this test accepts a custom parent even though docs/code-docs/source/autoep.rst:423-426 says custom patterns are rejected. If that checkpointed parent does anything value-dependent with the MoE output before the residual, _RecordedDeepEPCombine feeds it NaNs and can corrupt gradients. Reject the override at validation/conversion or validate the actual parent contract rather than only model_family.

Useful? React with 👍 / 👎.

[Copilot]

🟡 Changes recommended

Custom MoE layer-pattern overrides can bypass the preset safety gate and enable combine skipping for unvalidated decoder graphs.

1 open finding

• Reject custom layer patterns when skip_recompute_combine is enabled · New

What changed in this PR

Adds an opt-in AutoEP optimization that skips DeepEP combine work during supported reentrant checkpoint recomputation while preserving backward behavior.

Changes:

• Adds configuration, validation, decoder attachment, and recorded combine/weighting paths.
• Adds CPU and H100-focused correctness and replay regressions.
• Documents supported presets, constraints, and performance evidence.

File	Description
deepspeed/​module_inject/​auto_ep_config.py	Parses and validates the new option.
deepspeed/​module_inject/​auto_ep_presets/​base.py	Adds the configuration field.
deepspeed/​module_inject/​auto_ep.py	Attaches replacement MoE layers to decoder parents.
deepspeed/​module_inject/​auto_ep_layer.py	Detects supported recomputation and records skipped operations.
deepspeed/​module_inject/​auto_ep_comm.py	Implements recorded combine with shared backward context.
deepspeed/​ops/​triton_ops/​autoep_fused_token_ops.py	Adds recorded fused row weighting.
tests/​unit/​module_inject/​test_auto_ep_comm.py	Covers recorded backward behavior and repeated steps.
tests/​unit/​v1/​moe/​test_autoep_unit.py	Tests configuration and decoder attachment.
tests/​unit/​v1/​moe/​test_autoep_deepep_parity.py	Adds real DeepEP parity and combine-count tests.
docs/​code-docs/​source/​autoep.rst	Documents behavior, safety constraints, and results.
docs/​_pages/​config-json.md	Adds the configuration reference.

🧠 Review effort: Balanced

---

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

[tohtana]

Thank you for this awesome optimization!

One small issue I found is that skip_recompute_combine still accepts
a custom moe_layer_pattern together with a supported preset. The pattern
override changes the decoder parent while the preset name remains supported,
so the later whitelist does not prove that the custom parent only performs the
validated residual add. A value-dependent wrapper can consume the recorded
combine's NaNs during recomputation and corrupt gradients.

The minimal fix could be to reject a non-null moe_layer_pattern whenever skip_recompute_combine is enabled and add a supported-preset-plus-custom-pattern regression. See this commit.

[yh0903]

Thank you for this awesome optimization!

One small issue I found is that skip_recompute_combine still accepts a custom moe_layer_pattern together with a supported preset. The pattern override changes the decoder parent while the preset name remains supported, so the later whitelist does not prove that the custom parent only performs the validated residual add. A value-dependent wrapper can consume the recorded combine's NaNs during recomputation and corrupt gradients.

The minimal fix could be to reject a non-null moe_layer_pattern whenever skip_recompute_combine is enabled and add a supported-preset-plus-custom-pattern regression. See this commit.

Thanks @tohtana ! Applied the minimal validation guard you suggested in d1776a8c0. Any non-null moe_layer_pattern is now rejected when skip_recompute_combine is enabled, even with a supported explicit or auto-detected preset. Retaining the preset name does not establish the custom parent’s output-consumer contract. This also addresses the matching Codex and Copilot findings.
Added six parametrized regressions covering automatic/explicit preset selection and custom, preset-matching, and empty pattern overrides. Each also verifies that custom patterns remain accepted with the skip disabled, and that leaving the pattern unset preserves supported skip configurations. Both configuration references are updated; kernels, backward behavior, defaults, and numerical thresholds are unchanged.