Masahiro Tanaka

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>

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>

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>

Fix DeepCompile ZeRO-3 gathered parameter ownership (#8157) ## Problem DeepCompile inserts ZeRO-3 parameter all-gather and release operations into compiled graphs. When Dynamo skips a frame because of a graph break, however, that frame executes eagerly and does not run those graph operations. The eager fallback introduced in #8059 handles this case by all-gathering a partitioned parameter when the skipped frame accesses it through `ZeROOrderedDict`. The fallback is enabled around `DeepSpeedEngine.forward()`. Dynamo guard evaluation occurs inside that outer forward context and also resolves parameters through `ZeROOrderedDict`, while `torch.compiler.is_compiling()` is false. The fallback could therefore mistake a guard lookup for actual eager execution and unnecessarily all-gather the parameter. Parameters gathered by the fallback are normally partitioned after backward, but that cleanup does not run when backward is skipped. A fallback-gathered parameter may also be passed to an explicit `GatheredParameters` context, which must keep the full tensor available until the context exits. ## Why it matters These cases require different behavior: - Dynamo guard evaluation should not trigger an all-gather. - A parameter gathered for an eagerly executed frame must remain available through backward and then be partitioned. - If backward does not run, a leftover full parameter must be partitioned before the next outermost forward. - A parameter covered by `GatheredParameters` must remain fully gathered until that context exits. Without distinguishing these cases, a full parameter can remain allocated into a later forward, or fallback cleanup can partition it while a `GatheredParameters` block is still using it. ## Solution This PR: - detects parameter access during Dynamo guard evaluation and skips the eager fallback all-gather; - partitions leftover nonpersistent full parameters before the next outermost forward when the normal post-backward cleanup did not run; - removes a parameter from fallback cleanup when it is passed to `GatheredParameters`, so that context alone partitions it on exit; - restores the `GatheredParameters` state even when context exit raises; and - rejects nested `GatheredParameters` contexts that overlap on the same parameter, while continuing to allow nesting over disjoint parameter sets. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix DeepCompile ZeRO-1 grad target lifetime (#8036) DeepCompile ZeRO-1 kept compile-time reduce target buffers alive into the optimizer step, causing backward gradient storage to overlap with optimizer temporaries. This PR fixes the issue by making DeepCompile ZeRO-1 reduce targets follow the normal step-local ZeRO partition gradient-buffer lifetime, instead of preserving cloned target storage from compile setup. The actual code changes are: - During compile initialization, register empty DeepCompile ZeRO-1 gradient targets, then bind them to the step-local flat ZeRO partition gradient buffer and per-parameter views when gradients are ready to synchronize. - After ZeRO-1 builds the optimizer-facing fp32 gradient partition, release the DeepCompile registry references and clear reduce bucket storage after backward synchronization. | | Step 10-30 avg sec | Peak alloc GiB | | --- | --- | ---: | | Without this PR | 0.858 | 43.594 | | Without this PR | 0.859 | 39.366 | Fine-tuning style training (8xH100, Qwen3-8B random weights, bs/GPU=1, seq=4096, GAS=1) showed only finite value losses for 1000 steps. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix DeepCompile profiling memory cleanup (#8106) This PR reduces DeepCompile ZeRO-3 memory pressure in profiling and compiled backward by fixing three related issues: - Let the partitioner keep using its min-cut rematerialization policy for non-parameter activations, while still forcing ZeRO-3 parameter aliases/casts to be recomputed. - Compute selective-gather persistence capacity from profiled transient headroom instead of current allocator availability. - Avoid retaining temporary profiling outputs across warmup/measured profiling iterations, and clear gathered ZeRO-3 parameters when profiling exits through an exception. - Mark incomplete memory profiles explicitly, synchronize the completion status across ranks, and skip profile-dependent prefetch/selective-gather decisions when profiling data is incomplete. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix AutoEP ZeRO-1/2 universal conversion (#8198) Fixes #8147 ## Problem AutoEP checkpoints saved with ZeRO stages 1 and 2 store expert optimizer state in ZeRO-sharded fragments, but Universal Checkpoint conversion currently merges those expert fragments through the generic path before AutoEP consolidation. With source EP size greater than one, conversion can combine different EP ranks and fail while reshaping the result. The fallback path can also produce expert tensors without repartition metadata and use mixed-precision model weights where FP32 optimizer masters are required. ## Approach - Identify fused AutoEP expert parameters from checkpoint metadata before the generic merge, then reconstruct their FP32 masters and Adam states from the ZeRO fragments in EP-rank order. - Use the saved topology setting to handle both expert-before-data and data-before-expert rank layouts. - Save the existing expert metadata with every expert state so Universal Checkpoint load can repartition to a different EP size. - Restrict pipeline layer discovery to complete pipeline checkpoint filenames so per-expert files are not misclassified. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>

Release GIL in blocking AIO bindings (#8196) The blocking parallel AIO bindings can call the handle's internal `wait()` while retaining the Python GIL. An AIO worker may need that GIL while completing PyTorch tensor cleanup. This causes a deadlock: the Python caller waits for the worker, while the worker waits for the GIL held by the caller. This PR releases the GIL at the `pread`, `pwrite`, `sync_pread`, and `sync_pwrite` pybind entrypoints to avoid the deadlock. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Enable ZeRO-3 linear wrapper for existing models (#8189) Passing an already-constructed model to `deepspeed.initialize()` with ZeRO-3 and `memory_efficient_linear=true` does not install the ZeRO-3 Linear wrapper. The wrapper is currently installed only when the model is constructed inside a `deepspeed.zero.Init()` context. Without the wrapper, the standard Linear implementation can retain the gathered weight storage until backward completes, significantly increasing memory usage. This PR activates the existing ZeRO-3 Linear wrapper for the `deepspeed.initialize(model=...)` path when `memory_efficient_linear=true`, without requiring a `deepspeed.zero.Init()` context. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix DeepCompile profile metadata backfill (#8094) DeepCompile ZeRO-3 scheduling can crash with KeyError or AssertionError on missing device_time when graph profiling fails partway and leaves some nodes without timing metadata. This PR backfills default timing and memory metadata after profiling cleanup so the profiler guarantees downstream passes see complete metadata on every node. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix DeepCompile all-gather scheduler candidate selection (#8033) This PR fixes issues with the heuristic in DeepCompile's scheduler: - Fix a candidate-selection bug in `fast_free_schedule()`: the scheduler computed the zero-`free_acc_mem` candidate subset, but then sorted the full runnable set instead of that subset. - Keep the existing local scheduling heuristic, but rank candidates with graph-local all-gather pressure metrics before release-side cost when a low-live release path is available. - Add deterministic CPU-only FX scheduler regressions for the zero-free filter, pressure ordering, fallback candidate ordering, and single-all-gather ordering. ## Rationale `fast_free_schedule()` is a local heuristic for reducing gathered-parameter live ranges. This patch keeps that model, but fixes a general selection inconsistency: when at least one runnable candidate can reach release without additional all-gathers, the scheduler should choose from that zero-`free_acc_mem` subset. The previous code used the subset only as a branch condition, then ranked all runnable candidates by `free_cost`, so it could select a candidate that still required additional all-gathers before release. After preserving the zero-`free_acc_mem` filter, the ordering uses only workload-independent graph pressure signals already available to the scheduler: scheduled all-gather count, all-gather byte pressure, release-side cost, and a stable node-name tie breaker. In the fallback path, where every candidate still requires additional all-gathers, `free_acc_mem` remains the primary selector and the scheduler preserves the previous boundary of scheduling only through `schedule_until_ag`; this avoids making a memory-budget decision without tracking already-live gathered parameters. ## Testing - `python -m pytest tests/unit/compile/test_list_schedule.py -q` - `pre-commit run --all-files` --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Normalize ZeRO-3 DeepCompile grad dtype before reduction (#8038) Some backward kernels produce gradients in their computation dtype, not necessarily in the parameter storage dtype. For example, if a backward path accumulates or promotes math in fp32, a parameter stored as bf16 can still receive an fp32 raw gradient from that backward computation. In normal PyTorch execution, that raw gradient reaches the leaf-gradient accumulation step, which stores it according to the tensor's expected grad dtype. ZeRO-3 DeepCompile intercepts the raw compiled-backward gradient before that leaf accumulation boundary. The reducer was assuming the raw gradient dtype was already the expected leaf grad dtype, so it could select an fp32 communication bucket even when the ZeRO grad partition storage was bf16. To address this, this PR changes `dc.reduce_grad`'s behavior to match PyTorch's leaf-gradient dtype contract. ZeRO-3 registration now records the expected grad dtype for each parameter, and `reduce_grad` normalizes raw compiled-backward gradients to that dtype before selecting the communication bucket. This follows the documented `grad_dtype` behavior, including preserving explicit `grad_dtype=None` opt-outs: https://docs.pytorch.org/docs/main/generated/torch.sparse.semi_structured.SparseSemiStructuredTensorCUSPARSELT.html#torch.sparse.semi_structured.SparseSemiStructuredTensorCUSPARSELT.grad_dtype Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix DeepCompile ZeRO-3 release parameter lifetime (#8032) PR #7489 made ZeRO-3 all-gather allocate a padded base buffer for uneven shards and return a true-shape view into that buffer. That means the registry tensor and the tensor returned to the compiled graph no longer necessarily share the same `TensorImpl`, although they still share the same underlying storage. The existing release path only did `set_data(empty)` on the registry tensor before unregistering it. With the new base/view relationship, that clears the registry-side tensor metadata but does not resize the shared `StorageImpl` still referenced by returned views. As a result, the padded gathered allocation can remain live after the final `release_param`. This patch keeps the release graph ordering unchanged and makes final non-persistent release resize the registered gathered storage to 0 bytes before unregistering it. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Junjie Mao <junjie.mao@linux.alibaba.com>

Add configurable torch-latest dependency versions (#8016) ## Summary The modal test now shows the error because of the combination of PyTorch v2.7 and Transformer `main` branch, and it is blocking PRs. To address it, we improve our test workflows as follows. - Add manual dependency version inputs for the torch-latest CI workflows and default the torch-latest family to PyTorch 2.10 plus Transformers git `main`. - Let CPU and AWS full torch-latest runs select either released Transformers package versions or an explicit Transformers git ref for manual validation. - Let Modal torch-latest runs select supported PyTorch/CUDA image presets and an optional Transformers git ref, defaulting to `2.10.0-cuda12.8` and Transformers git `main`. ## Known follow-up - The AWS full real CI lane for PyTorch 2.10 plus Transformers main reached `Unit tests (parallel)` but failed with 33 failures. Some of these may overlap with fixes in #8015; I am opening this PR now so the workflow/input changes can be reviewed while those failures are handled separately. - CPU and Modal real CI validation for the requested tuple passed. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix ZeRO-3 autocast gather with mixed parameter dtypes (#8113) This PR allows ZeRO-3 coalesced all-gather buckets to contain parameters with different original dtypes when they share the same autocast communication dtype. Motivation: PEFT-style LoRA adapters can remain FP32 trainable parameters while the base model parameters are BF16. When DeepSpeed autocast marks both sets of parameters for BF16 communication, the existing assertion still checks their original parameter dtypes and fails on the BF16/FP32 mismatch. With this PR, LoRA-style FP32 adapter parameters can participate in BF16 autocast communication without requiring applications to pre-cast the adapters to BF16. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Fix DeepCompile AOT kwargs patching for PyTorch >= v2.11 (#8024) DeepCompiles breaks for PyTorch >= v2.11 because these versions can construct the AOT Autograd backend without a bw_compiler kwarg, while DeepSpeed's Inductor patch assumes that key is always present. This PR fixes DeepCompile's AOT Autograd patch so unrelated AOT backend registrations can pass through unchanged. `TestDeepCompile` passes with this fix. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Preserve tensor learning rates in OneCycle (#8205) Follow up #8202: When an optimizer starts with a tensor learning rate, OneCycle initialization replaces it with a Python scalar while applying `cycle_min_lr`. This loses the caller's tensor identity, shape, and dtype before later scheduler updates can preserve them. This PR initializes OneCycle learning rates through the existing tensor-aware update helper, matching the path used by subsequent scheduler steps. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Keep required CI checks visible for ignored paths (#8019) The motivation for this PR is that some checks we want to make required are produced by workflows that currently use trigger-level paths-ignore. When GitHub skips an entire workflow before it starts, such as on a docs-only PR, it never creates the corresponding check run, so branch protection can remain stuck waiting for a required status that will never be reported. This PR keeps those workflows starting consistently and moves the path filtering inside the workflow, following the existing check-paths / skip pattern used by the AWS workflows, so ignored-path PRs can still report successful required checks without running the expensive CI work. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>

Fix full CI test isolation for ZeRO chmod and NVMe quantization tests (#8008) ## Summary This PR fixes two intermittent full-CI test isolation [failures](https://github.com/deepspeedai/DeepSpeed/actions/runs/25789145638/job/75749943219) observed in the scheduled `aws-torch-latest-full` workflow. - Avoid TCP `env://` rendezvous port collisions in `TestZeRONonDistributed::test_chmod_exception_handling`. - Give the NVMe int4 quantization tests per-test offload directories instead of sharing `~/tmp_offload_dir`. ## Root Cause - The ZeRO chmod test sets `world_size = 1`, but it disabled the distributed test harness initialization while still calling `deepspeed.initialize()`. In the full CI `pytest-xdist -n 8` environment, this could fall back to TCP rendezvous and collide on the selected `MASTER_PORT`. - The NVMe quantization tests both used the same `~/tmp_offload_dir`. When the post-init NVMe test and the quantized-initialization NVMe test ran concurrently on different xdist workers, one worker could remove or recreate rank-local swap files while the other worker was still reading them. ## Changes - Let `TestZeRONonDistributed` use the existing file-store distributed test harness initialization. - Add an optional `nvme_path` argument to the NVMe quantization helpers. - Pass a `tmpdir`-scoped `nvme_offload` path from each NVMe test. The full workflow passed with this PR branch: https://github.com/deepspeedai/DeepSpeed/actions/runs/25842039450 Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>