github.com-microsoft-DeepSpeed

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>

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>

Ulysses SP for HF Integration (#7268) This is the Deepspeed counterpart of https://github.com/snowflakedb/ArcticTraining/pull/45 - as the new feature(s) require changes on both sides. For PR reviewers: Readiness status: - [x] Code - [x] Tests - [ ] Docs - working on it Features: - [x] add support for delaying grad addition via `param.ds_grad_is_ready` flag (used when performing tiled compute in an autograd function) - [x] add light sp-only mpu version (Jeff Rasley) - [x] improved debug - [x] added `all_gather_object` to `dist` - [x] `UlyssesSPAttentionHF` (port of UlyssesAttention from Megatron-Deepspeed plus modern MHA-variations) - [x] `UlyssesSPDataLoaderAdapter` - DL adapter to shard the normal DL batches to be used by `UlyssesSPAttentionHF` - [x] `SequenceTiledCompute` - generic autograd function to perform compute after tiling on the sequence dimension - [x] `TiledMLP` - a specific autograd function to perform tiled MLP (it's much easier to understand before trying to grok `SequenceTiledCompute`) - [x] added a differentiable `_DimZeroAllToAll` (Samyam Rajbhandari) - [x] torch-dist-check now allows `torch.distributed.nn` (which is needed since deepspeed's dist is not up to date with `torch.distributed.nn`) --------- Signed-off-by: Stas Bekman <stas.bekman@snowflake.com> Signed-off-by: Stas Bekman <stas@stason.org> Co-authored-by: Stas Bekman <stas.bekman@snowflake.com> Co-authored-by: Jeff Rasley <jerasley@microsoft.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.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>

CI: prefer bf16 over fp16 (#7304) these days fp16 is barely ever used, so we should be testing bf16 instead of fp16 where possible. had to fix a bunch of tests to adapt to this change. a few bugs as well on the way. --------- Signed-off-by: Stas Bekman <stas.bekman@snowflake.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Stas Bekman <stas.bekman@snowflake.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 DeepCompile+Z3 on PyTorch v2.9/2.10 (#7951) DeepCompile+Z3 didn't work with PyTorch v2.9/2.10 because: - PyTorch v2.9+ started enforcing stricter TorchDynamo parameter tensor-match guards. During DeepCompile tracing, some ZeRO-3 parameters were temporarily all-gathered, so Dynamo recorded full sizes such as 4096 - By the time guard evaluation ran, DeepSpeed had already released those params back to the normal ZeRO-3 partitioned representation, where `param.data` is `empty(0)`. That produced guard failures like `expected 4096, actual 0`. This PR resolves the issue by: - Leep full-shape dummy tensors for symbolic tracing - Override guard size/stride metadata for ZeRO-3 params to the stable released representation instead of transient gathered sizes This PR also includes fixes of these bugs: - For v2.7 and v2.8, the compiled backward graph could hoist `end_backward` ahead of the real `reduce_grad` calls. - Selective unsharding pass can overcount the persistence memory budget. Note: DeepCompile is still incompatible with v2.11. It will be addressed by another PR. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Update GH org references (#6998) Signed-off-by: Olatunji Ruwase <olruwase@microsoft.com> Signed-off-by: Logan Adams <loadams@microsoft.com> Signed-off-by: Fabien Dupont <fdupont@redhat.com> Co-authored-by: Fabien Dupont <fabiendupont@fabiendupont.fr>

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>

Enable torch.autocast with ZeRO (#6993) DeepSpeed supports mixed precision training, but the behavior is different from `torch.autocast`. DeepSpeed maintains parameters and gradients both in FP32 and a lower precision (FP16/BF16) (NVIDIA Apex AMP style) and computes all modules in the lower precision while `torch.autocast` maintains parameters in FP32 but computes only certain operators in the lower precision. This leads to differences in: - performance: `torch.autocast` needs downcast in forward/backward - memory usage: DeepSpeed needs more memory to keep copies of parameters and gradients in lower precision - accuracy: `torch.autocast` has a list of modules that can safely be computed in lower precision. Some precision-sensitive operators (e.g. softmax) are computed in FP32. To align DeepSpeed's behavior with `torch.autocast` when necessary, this PR adds the integration with `torch.autocast` with ZeRO. Here is an examples of the configuration. ```json "torch_autocast": { "enabled": true, "dtype": "bfloat16", "lower_precision_safe_modules": ["torch.nn.Linear", "torch.nn.Conv2d"] } ``` Each configuration works as follows: - `enabled`: Enable the integration with `torch.autocast` if this is set to `True`. You don't need to call `torch.autocast` in your code. The grad scaler is also applied in the DeepSpeed optimizer. - `dtype`: lower precision dtype passed to `torch.autocast`. Gradients for allreduce (reduce-scatter) and parameters for allgather (only for ZeRO3) of `lower_precision_safe_modules` are also downcasted to this dtype. - `lower_precision_safe_modules`: Downcast for allreduce (reduce-scatter) and allgather (ZeRO3) are applied only to modules specified in this list. (The precision for PyTorch operators in forward/backward follows `torch.autocast`'s policy, not this list.) You can set names of classes with their packages. If you don't set this item, DeepSpeed uses the default list: `[torch.nn.Linear, torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d]`. Note that we only maintain FP32 parameters with this feature enabled. For consistency, you cannot enable `fp16` or `bf16` in DeepSpeed config. --------- Signed-off-by: Masahiro Tanaka <mtanaka@microsoft.com> Signed-off-by: Fabien Dupont <fdupont@redhat.com> Signed-off-by: Olatunji Ruwase <olruwase@microsoft.com> Signed-off-by: Logan Adams <loadams@microsoft.com> Signed-off-by: inkcherry <mingzhi.liu@intel.com> Signed-off-by: Omar Elayan <oelayan@habana.ai> Signed-off-by: Roman Fitzjalen <romaactor@gmail.com> Signed-off-by: Hongwei <hongweichen@microsoft.com> Signed-off-by: shaomin <wukon1992@gmail.com> Signed-off-by: Stas Bekman <stas@stason.org> Signed-off-by: siqi <siqi@tecorigin.com> Signed-off-by: Wei Wu <wuwei211x@gmail.com> Signed-off-by: ShellyNR <shelly.nahir@live.biu.ac.il> Signed-off-by: Lai, Yejing <yejing.lai@intel.com> Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com> Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Fabien Dupont <fabiendupont@fabiendupont.fr> Co-authored-by: Liangliang Ma <1906710196@qq.com> Co-authored-by: inkcherry <mingzhi.liu@intel.com> Co-authored-by: Omar Elayan <142979319+oelayan7@users.noreply.github.com> Co-authored-by: Stas Bekman <stas00@users.noreply.github.com> Co-authored-by: Roman Fitzjalen <romaactor@gmail.com> Co-authored-by: Ramya Ramineni <62723901+rraminen@users.noreply.github.com> Co-authored-by: Guanhua Wang <alexwgh333@gmail.com> Co-authored-by: root <root@ftqtmec25000000.taxzvufipdhelhupulxcbvr15f.ux.internal.cloudapp.net> Co-authored-by: Hongwei Chen <33092912+hwchen2017@users.noreply.github.com> Co-authored-by: Joe Mayer <114769929+jomayeri@users.noreply.github.com> Co-authored-by: wukong1992 <wukong1992@users.noreply.github.com> Co-authored-by: shaomin <wukon1992@gmail.com> Co-authored-by: loadams <loadams@users.noreply.github.com> Co-authored-by: siqi654321 <siqi202311@163.com> Co-authored-by: siqi <siqi@tecorigin.com> Co-authored-by: Wei Wu <45323446+U-rara@users.noreply.github.com> Co-authored-by: Shelly Nahir <73890534+ShellyNR@users.noreply.github.com> Co-authored-by: snahir <snahir@habana.ai> Co-authored-by: Yejing-Lai <yejing.lai@intel.com> Co-authored-by: Siddharth Singh <siddharth9820@gmail.com> Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>

Add HuggingFace tp_plan support for AutoTP (#7901) ## Summary Adds automatic detection and use of HuggingFace's built-in `base_model_tp_plan` for AutoTP, addressing the HuggingFace tp_plan support item from #7861. Models that ship with a `tp_plan` (e.g. Llama, Qwen, Gemma2) now work with AutoTP out of the box — no `preset_model` or `partition_config` needed, just set `autotp_size`. ## Changes **Runtime** - `engine.py`: Added tp_plan fallback in `_apply_autotp_partitioning`. Priority order: `partition_config` > HF `tp_plan` > AutoTP heuristics. - `config.py`: Added `_get_hf_tp_plan(model)` to extract tp_plan from `model._tp_plan` or `model.config.base_model_tp_plan`. - `tp_plan_converter.py`: New file. `TPPlanConverter` converts HF tp_plan entries (`colwise`/`rowwise`) to DeepSpeed `TPLayerSpec`. Other HF partition types (`colwise_rep`, `local_colwise`, etc.) are not yet supported (documented with TODO). **Tests** (11 files, 17 CPU + 5 GPU tests) - `test_tp_plan_converter.py`: Unit tests for the converter (alternate prefixes, projection names, unsupported types, etc.) - `test_tp_plan_extraction.py`: Unit tests for `_get_hf_tp_plan` with mock models. - `test_tp_plan_e2e.py`: GPU e2e tests with ZeRO 0/1/2 (requires 2 GPUs). - `test_tp_plan_real_models.py`: GPU tests with Qwen2 and custom models (requires 2 GPUs). **Documentation** - Tutorial: New "HuggingFace tp_plan Support" section in `autotp-training.md`. - Config reference: Added tp_plan paragraph in `config-json.md`. - API docs: Added tp_plan subsection in `training.rst`. - Blog: Updated ongoing work in `blogs/huggingface-tp/README.md`. ## Limitations - Only `colwise` and `rowwise` partition types are supported. Extended types (`colwise_rep`, `local_colwise`, `local_rowwise`, `local_packed_rowwise`, `gather`, `sequence_parallel`) are deferred. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: Ma, Guokai <guokai.ma@gmail.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.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>