DeepSpeed — Engineering Performance

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>

Support AutoEP with ZeRO-3 zero.Init source modules (#8060) This PR enables ZeRO3 support for AutoEP-managed MoE layers by partitioning expert parameters over expert replica groups while router and replicated parameters use the global data-parallel group. With ZeRO3 enable, AutoEP preserves global data-parallel gradient averaging for AutoEP expert parameters while reducing them over expert replica groups. ZeRO parameters are gathered before AutoEP reads router or expert tensors when replacing MoE modules created under `deepspeed.zero.Init()`. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

[AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP (#8168) AutoTP + ZeRO-3 silently produced incomplete checkpoints: both export paths handled only the ZeRO data-parallel dimension and dropped the tensor-parallel shards. - ds_to_universal.py: stage3 conversion recovers the (tp,dp) grid from checkpoint file names, extracts shards under the real tp_index, and reuses the stage<=2 TP-aware merge when tp_degree>1 (DP-only path preserved for tp_degree==1 -> no regression for plain ZeRO-3). - engine.py: _zero3_consolidated_16bit_state_dict nests GatherReplacedLayerParams inside GatheredParameters so save_16bit_model gathers both DP and TP; remove the blanket autotp+zero3 training block now that checkpoint consolidation is implemented. - stage3.py: load_hp_checkpoint_state resolves the TP shard before the ZeRO-DP partition, so universal checkpoint restore round-trips. Add end-to-end universal conversion tests and update existing tests for the refactored merge_tp_slices / extract_zero_shards_stage3 signatures. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Support custom partitioning patterns for AutoTP (#7806) This PR introduces a flexible, configuration-driven API for AutoTP (Automatic Tensor Parallelism) that allows users to define custom layer partitioning patterns for training. @inkcherry @delock ## Motivation Previously, AutoTP relied on hardcoded layer detection logic that was difficult to customize for new model architectures. This PR enables: 1. **Custom models**: Users can define exact regex patterns to match their model's parameter names 2. **Fused layers**: Support for fused QKV, gate_up_proj, and other packed weight matrices with unequal sub-parameter sizes (e.g., GQA with different Q/K/V dimensions) 3. **Extensibility**: Easy to add new model presets or customize existing ones Here is an example of a config including custom partitioning patterns: ```json { "tensor_parallel": { "autotp_size": 4, "partition_config": { "use_default_specs": false, "layer_specs": [ { "patterns": [".*\\.o_proj\\.weight$", ".*\\.down_proj\\.weight$"], "partition_type": "row" }, { "patterns": [".*\\.[qkv]_proj\\.weight$"], "partition_type": "column" }, { "patterns": [".*\\.gate_up_proj\\.weight$"], "partition_type": "column", "shape": [2, -1], "partition_dim": 0 } ] } } } ``` Refer to the [document](https://github.com/tohtana/DeepSpeed/blob/tohtana/autotp_custom_patterns/docs/code-docs/source/training.rst) for more details (including preset models and how to define partitioning for fused models). We also opened a new [PR](https://github.com/deepspeedai/DeepSpeedExamples/pull/998) to show the usage. ## Simplified initialization step AutoTP previously required calling ``set_autotp_mode(training=True)`` and ``deepspeed.tp_model_init`` before ``deepspeed.initialize``. Now we can include all the necessary configurations in the DeepSpeed config. We still support the traditional initialization path for backward compatibility. When you use both (i.e. calling ``set_autotp_mode(training=True)`` and ``deepspeed.tp_model_init`` and passing the config to ``deepspeed.initialize``), we will merge the settings at initialization. When we have conflicting settings, we will error out. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Non-reentrant activation checkpoint CPU offload (#8282) ## Summary - Continues [@winglian](https://github.com/winglian)'s work from #8181: a `saved_tensors_hooks` context manager that offloads non-reentrant checkpoint hidden-state inputs to a pinned CPU buffer pool on a side stream (`use_reentrant=False`). - The first two commits are **authored and signed off by Wing Lian** (`wing@axolotl.ai`); they are the original #8181 patches, rebased onto current `master`. This follow-up commit addresses review without rewriting those commits. - Review follow-up: restore `GradientCheckpointingLayer.__call__` when no manager is active (HybridEngine train/rollout), skip offloading the last checkpoint input (`keep_last_count=1`), rename `*_size` knobs to `*_bytes` / `*_count`, and add tests for the HF signature contract and keep-last behavior. - Wires the async offload into DeepSpeed native `cpu_checkpointing`: the copy machinery is factored into a reusable `_ActivationOffloadEngine`, which both the HF hooks class and native `CheckpointFunction` / `non_reentrant_checkpoint` share. Also fixes two pre-existing `non_reentrant_checkpoint` + `cpu_checkpointing` bugs (inputs emptied before forward; `saved_data` never restored during recompute). Original upstreaming context: axolotl-ai-cloud/axolotl#3776, requested in #8181. ## Test plan - [x] `pytest tests/unit/runtime/activation_checkpointing/test_offload_activations.py` on H200 — HF `saved_tensors_hooks` path (26 passed) - [x] `pytest tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py` on H200 — native reentrant + new `cpu_checkpointing` offload tests (30 passed) - [x] `pytest tests/unit/runtime/activation_checkpointing/test_activation_checkpointing_non_reentrant.py` on H200 — native non-reentrant + new `cpu_checkpointing` offload test (49 passed) - [x] H200 microbenchmark (activation-dominant): async CPU offload matches blocking's 58% peak-memory reduction at ~5.7% step-time overhead (vs ~6.8x for blocking), i.e. 6.4x faster than blocking offload - [ ] CI unit tests for activation checkpointing Made with [Cursor](https://cursor.com) --------- Signed-off-by: Wing Lian <wing@axolotl.ai> Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Wing Lian <wing@axolotl.ai> Co-authored-by: Cursor <cursoragent@cursor.com>

Z3: Support for activation ckpt with frozen params (#8148) Build on #8130 --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Run pull request code only inside a Modal Sandbox (#8170) GitHub recently changed `actions/checkout` to reject fork pull request checkouts by default in privileged `pull_request_target` workflows. This protection was introduced in `actions/checkout` v7 and backported to supported floating major versions, including `actions/checkout@v4`. As a result, the current workflow's checkout of `${{ github.event.pull_request.head.sha }}` is now rejected for fork pull requests. Setting `allow-unsafe-pr-checkout: true` would restore the previous behavior, but it would deliberately bypass this protection. A `pull_request_target` workflow runs trusted default-branch code with access to the base repository's token and secrets. Because this workflow needs `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET` to start GPU CI, the code that handles those credentials must remain trusted code from `master`; pull request code must not execute in the same GitHub runner context. Modal documents Sandboxes as secure containers for executing untrusted code, including checking out a repository and running its test suite. This change follows that model: GitHub runs only the trusted selector and Modal controller, while the pull request is treated as data during test selection. The controller passes only a validated public repository name, exact commit SHA, and validated test selection to a Modal Sandbox. The Sandbox receives no GitHub, Modal, or Hugging Face credentials, and performs the candidate checkout, dependency installation, DeepSpeed installation, and pytest execution itself. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

[CI] diff driven test selection (#8077) TLDR: Analyze PR's diff and filter out tests that aren't exercising what has changed, potentially cutting down runtime and expense by 95-99% most of the time. Detailed: Deepspeed's CI takes forever - most of the time burning $$ and wastes dev time for no reason as most changes require just a few tests to run. HF Transformers has a system to select which tests to run based on the diff of the PR - Sylvain Gugger wrote it many years ago since that repo has now probably thousands of tests. Deepspeed's CI isn't too bad but can easily take hours. So I asked Claude Opus 4.8 to replicate the system for Deepspeed. Please have a look. It looks super complicated, so I'm not sure how easy it'd be to maintain/operate unless we always use AI to continue maintaining it. I asked Claude to leave a detailed state file so that it or another model could pick it up where it left. I started with just the slowest costliest workload `.github/workflows/modal-torch-latest.yml` to see if it works well. If you're happy we can replicate it to the rest of the workloads. CC: @loadams, @tjruwase - please tag others if you think they would be helpful to discuss this. --------- Signed-off-by: Stas Bekman <stas@stason.org>

PyTorch-compatible backward API (#7665) Currently DeepSpeed's backward API has more constraints compared to PyTorch's normal backward API. Here is the usage as described in the documentation: ```python loss = model_engine(batch) model_engine.backward(loss) ``` In this example, 1. Only accepts a (scalar) loss value 1. Need to call engine's backward API In contrast, in standard PyTorch, you can do: ```python output = model(batch) output.backward(out_grad) ``` There are several use cases that rely on this flexibility. For example, combining multiple models or using loss functions defined separately from the main model. If you attempt the same pattern with a DeepSpeed engine, some preprocessing and postprocessing steps will be silently skipped, which can lead to incorrect results. The [document](https://deepspeed.readthedocs.io/en/latest/training.html#jointly-training-models-with-shared-loss) explains we can call `_backward_epilogue` manually (possibly `backward_prologue` as well). However, it's easy for users to miss these calls, and passing a non-scalar gradient is still not supported. This PR introduces the same `.backward()` behavior as PyTorch, allowing .backward() to be called directly on tensors and supporting non-scalar outputs. To implement post-backward hooks, we had to use some torch internal APIs. See [comments](https://github.com/deepspeedai/DeepSpeed/blob/73f7ff1aab9d1387eb7dd4eca7453a25024533f4/deepspeed/runtime/engine.py#L424) for more details. When the internal APIs are not available, DeepSpeed engine only accepts the traditional way `model_engine.backward(loss)`. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Scope DeepCompile compiler state to graph and engine lifecycles (#8159) DeepCompile uses several process-global PyTorch compiler mechanisms while building graph-specific state. Scheduled recompilation also replaces previously compiled graphs. These lifetimes did not have explicit ownership boundaries. - Forward inputs could be consumed by a different graph - Dynamo frame IDs could collide across engines - Cleanup could either leave graph-specific patches installed or release shared compiler state that another engine still needed This PR gives those existing mechanisms explicit graph and engine lifecycle ownership. Fixes: - Keep forward inputs in a graph-local one-shot queue and `InputStorage`. - Qualify compiled-backward frames with an owner token and track them per engine. - Keep the `torch.autograd.Function` patch active until the last owned backward frame is released, then clear captured backward inputs. - Reference-count the ZeRO-3 Dynamo configuration overrides across engines and restore the original values after the last owner exits. - Limit the graph-specific `AotAutograd.__init__` patch to the Inductor compilation call and restore it in `finally`. - Release engine-owned compiler state during scheduled recompilation, deactivation, and destruction. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>