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>

Split DeepCompile ZeRO-3 memory scheduler (#8233) Splits the DeepCompile ZeRO-3 scheduling and gather/release lifetime changes out of #8169. The exact candidate passed 10/10 tests on a single node with 2 H100 GPUs. Signed-off-by: Masahiro Tanaka <tanaka.masahiro@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 shared loss gradient accumulation (#8245) ## Problem Fixes #8224. With gradient accumulation enabled, a loss can combine a forward pass through the DeepSpeed engine with a forward pass through the model inside the engine. The engine-output hook divides only the gradient from the engine forward by the accumulation count. The gradient from the forward pass on the model inside the engine remains unscaled, producing an incorrect gradient for that parameter. ## Approach Apply gradient-accumulation scaling to the complete loss passed through managed `engine.backward`. During that call, mark the loss graph as already scaled so engine-output hooks do not apply the scaling again, and restore all managed-backward state on every exit. If a managed backward is interrupted, preserve the existing ZeRO retry behavior by not running a direct-backward epilogue over incomplete reduction state. Direct tensor backward, pipeline output hooks, and `scale_wrt_gas=False` otherwise retain their existing behavior. ## Testing - `pytest -q tests/unit/runtime/zero/test_zero_shared_loss_gradient.py` - Verified the existing ZeRO-3 exception-retry lifecycle behavior. 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 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>

Partition AutoEP expert parameters per layer under ZeRO-3 (#8424) ### Problem When a model created under `zero.Init` is passed to `deepspeed.initialize`, AutoEP creates new local-expert tensors while replacing MoE layers. ZeRO-3 previously converted those tensors only after every layer had been replaced, so full replacement-expert storage accumulated across layers and increased the initialization peak. This addresses the initialization-memory behavior described as problem 1 in #8353. It does not change the independent debug-map retention behavior covered by #8356. ### Approach After each AutoEP layer is repacked, copied, tagged, and installed, resolve its expert-data-parallel group and immediately pass only its newly created expert parameters through the existing ZeRO conversion helper. Optimizer classification and grouping use the experts' ZeRO logical shape and size after partitioning. Existing source gathers, router/shared/dense placement, and batched output-recorder retargeting remain unchanged. The later ZeRO setup already skips parameters converted by this step. --------- Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>