## Summary
This PR is the first step (phase 0) of enabling Apple Silicon support
for DeepSpeed: make single-device training work end to end with
pure-PyTorch ops.
[**To-Do in phase 1**] Metal kernels will come later and plug into the
`op_builder/mps` classes added here.
The MPS accelerator was a stub: memory queries returned `None`, no
communication backend was set, `fp16/bf16` were reported unsupported,
and every op builder resolved to `NotImplementedBuilder`.
`deepspeed.initialize` + one training step failed for every ZeRO stage
on an Apple Silicon machine. This PR aims on enabling capabilities.
### Changes
- **`accelerator/mps_accelerator.py`** — real `torch.mps` memory stats,
fp16/bf16 support (bf16 gated on macOS 14+), `torch.mps.Event`, `gloo`
as the comm backend, and `is_synchronized_device() = True` (PyTorch's
MPS backend effectively exposes a single in-order execution stream and
currently provides no public CUDA-style stream API or `record_stream`
mechanism.). Unified memory makes `pin_memory` a no-op (torch's
`pin_memory()` also raises under MPS).
- **`deepspeed/comm/torch.py`** — gloo cannot operate on MPS tensors
(even at world size 1), so collectives stage MPS tensors through CPU
copies via a `stage_on_cpu` decorator; async ops copy back on `wait()`.
- **`accelerator/abstract_accelerator.py` + `runtime/zero`** — MPS has
no fp64. Gradient-norm accumulation now picks its dtype via a new
concrete `is_fp64_supported()` (default `True`) and `get_norm_dtype()`
instead of hard-coded `.double()`.
- **`op_builder/mps/`** — new backend package (`MPSOpBuilder`,
`NotImplementedBuilder`, `FusedAdamBuilder`). `FusedAdam` is implemented
with `torch._foreach_*` ops and mirrors the math in
`csrc/adam/multi_tensor_adam.cu`, following the HPU precedent of
Python-backed builders.
- **`tests/unit/common.py`** — MPS must use `spawn` (Metal's compiler
service is lost in `forkserver` children, which hangs the harness) and
reports its device count via the accelerator.
- **`tests/unit/ops/adam/test_adamw.py`** —
`test_fused_adam_matches_torch` checks `FusedAdam` against
`torch.optim.Adam/AdamW` on the active accelerator (fp32/bf16 ×
Adam/AdamW), so it also guards the CUDA kernel.
### Verified on an M5 Max (macOS 26.3, torch 2.13.0)
- ZeRO 1/2/3 × fp32/bf16/fp16 train end to end with
`deepspeed.initialize` (single process). Also tested on ZeRO stage 0
which disables ZeRO completely and falling back to standard data
parallelism.
- MPS `FusedAdam` matches `torch.optim` to 2e-7 in fp32.
- `DS_ACCELERATOR=mps pytest unit/runtime/test_ds_config_dict.py
unit/runtime/test_ds_initialize.py
unit/runtime/half_precision/test_fp16.py
unit/runtime/half_precision/test_dynamic_loss_scale.py
unit/runtime/zero/test_zero_grad_clip.py
unit/runtime/zero/test_zero_context.py
unit/checkpoint/test_zero_optimizer.py`: 132 passed, 0 failed, 126
skipped (multi-device tests; `device_count() == 1`).
### Known limitations / follow-ups
- bf16 `FusedAdam` differs from the CUDA kernel by ~1 bf16 ulp (CUDA
computes in fp32 and stores bf16; the `_foreach` path rounds in bf16).
- The CPU-staged gloo path is only exercised at world size 1 here;
multi-Mac runs are untested.
- Follow-ups: macOS arm64 CI workflow, arm64 build of CPU Adam for
ZeRO-Offload, Metal kernels via `torch.mps.compile_shader`, and an Apple
Silicon tutorial page.
---------
Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
## 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>
Follow up #8146.
## Summary
This pull request introduces support for uneven sub-parameter sharding
in DeepSpeed's universal checkpoint conversion, updating the universal
checkpoint format to version 0.4. The changes ensure that partitioned
parameters with sub-parameters of varying sizes are correctly handled
during checkpoint conversion, merging, and restoration. Additionally,
the PR adds validation to prevent conversion of unsupported checkpoint
layouts and improves error handling and metadata validation.
Key updates by theme:
**Universal Checkpoint Format and Metadata:**
- Bumped the universal checkpoint version to 0.4 and introduced the
`SUB_PARAM_SHARD_WIDTHS` field to record per-rank widths for each
sub-parameter, enabling correct handling of uneven sub-parameter
layouts. (`deepspeed/checkpoint/constants.py`,
`deepspeed/checkpoint/ds_to_universal.py`)
[[1]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eL61-R64)
[[2]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eR93-R97)
[[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R35)
[[4]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R302-R303)
[[5]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L335-R441)
- Added `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` to checkpoint metadata
and implemented validation to prevent conversion if unsupported
parameter patterns are present. (`deepspeed/checkpoint/constants.py`,
`deepspeed/checkpoint/ds_to_universal.py`)
[[1]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eL61-R64)
[[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R47)
[[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R950-R957)
[[4]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R1134-R1139)
**Parameter Merging and Sharding Logic:**
- Enhanced the merging logic to correctly handle missing fragments for
ranks with no data in uneven parameter sharding, ensuring proper
alignment of slices and placeholder insertion.
(`deepspeed/checkpoint/ds_to_universal.py`)
[[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R239-R248)
[[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R273-R285)
- Refactored the sub-parameter merging code to use the new shard widths
metadata, supporting both legacy (even) and new (uneven) layouts. Added
logic to reconstruct logical shapes with placeholder dimensions.
(`deepspeed/checkpoint/ds_to_universal.py`)
[[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L302-R327)
[[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L335-R441)
[[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R463-R482)
**Validation and Error Handling:**
- Added early validation for unsupported AutoTP conversions to fail fast
before expensive extraction steps.
(`deepspeed/checkpoint/ds_to_universal.py`)
[[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R950-R957)
[[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R1134-R1139)
- Improved shape consistency checks for pipeline-parallel parameters to
ensure all replicas agree on shape.
(`deepspeed/checkpoint/ds_to_universal.py`)
**Restoration Logic:**
- Updated the restoration logic to use the new shard widths metadata,
ensuring correct reconstruction of sub-parameters during model loading.
(`deepspeed/checkpoint/universal_checkpoint.py`)
[[1]](diffhunk://#diff-22d7b9e3b6eac1dc6e989cd8582946321c36f4fc80b527648eeb2d77a8fd3ee2L13-R13)
[[2]](diffhunk://#diff-22d7b9e3b6eac1dc6e989cd8582946321c36f4fc80b527648eeb2d77a8fd3ee2R34-R73)
These changes collectively improve the robustness and flexibility of
DeepSpeed's checkpoint conversion, especially for advanced tensor
parallelism scenarios.
## Testing
* Added coverage for uneven vocabulary, GQA projections, checkpoint
conversion/restore, and PP + TP tied parameters.
* loss curve https://github.com/deepspeedai/DeepSpeedExamples/pull/1008
## Limitations
tp_size > num_kv need to further be optimized
---------
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Ma,Guokai <guokai.ma@intel.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
## Summary
- Extract host page-locking (`posix_memalign`/`mlock`) into a standalone
`deepspeed.ops.pin_memory` / `PinMemoryBuilder` that does not require
libaio or AIO worker threads.
- Compile the pin manager only in the `pin_memory` op and share one
process-wide manager with `async_io`/`gds` (via exported symbol +
`RTLD_GLOBAL`) so DeepNVMe bounce-buffer skipping and `is_pinned` stay
consistent.
- Keep `new_cpu_locked_tensor` / `free_cpu_locked_tensor` / `is_pinned`
on `aio_handle`/`gds_handle` as thin wrappers; point XPU `align_bytes=0`
at `pin_handle`.
## Test plan
- [x] `tests/unit/v1/pin_memory/test_pin_memory_op.py` (pin without
async_io)
- [x] `tests/unit/v1/nvme/test_pinned_manager.py` cross-op recognition
(`pin_handle` ↔ `aio_handle`)
- [x] Confirm AIO I/O tests still pass with shared manager
(`tests/unit/v1/nvme/` including `test_aio.py` / `test_gds.py` — 131
passed on `tunji-h200-n1g2-ds2-0`, job `20260809T123310Z`, HEAD
`a6f6ab6e`)
- [x] `ds_report` shows `pin_memory` as compatible without libaio-dev
(`pin_memory ... [OKAY]`; with `io_submit`/libaio mocked missing,
`pin_memory` stays compatible while `async_io` does not — job
`20260809T124859Z`)
## Follow-up
Native backend (`DS_PIN_MEMORY_BACKEND=native`, #8211) will be stacked
on this PR once it lands.
Made with [Cursor](https://cursor.com)
---------
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
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>
### Changes
Enable HF `colwise_gather_output` in `tp_plan` through AutoTP, allowing
vocab parallelism for an untied `lm_head`.
### Known Limitations
- Only supports an untied `lm_head`. Tied `lm_head` falls back to the
legacy implementation.
- Uneven tensor parallelism is not supported in this PR.
---------
Signed-off-by: iLeGend <824040212@qq.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Replace the unconditional zero_optimization_stage() <= 2 assert in
_configure_tensor_parallel_states with a guard that only blocks AutoTP +
ZeRO-3 when an optimizer is present. The ZeRO-Inference path (no
optimizer → DummyOptim → DeepSpeedZeRoOffload) is now permitted; the
training path (optimizer → DeepSpeedZeroOptimizer_Stage3) raises
NotImplementedError.
AutoTP with ZeRO-3 checkpoint saving should be tracked seperatly in a
seperate PR.
---------
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
As the extensions share the build folder building them in parallel can
cause failures or wrong results due to extensions overwriting the files
of other extensions.
Closes#949 which is one instance of an actual failure: transformer_op
AND the stochastic_transformer_op compile the same file in the same
folder with different options in parallel.
Other issues include mangled ninja build files (created by PyTorch):
```
ninja: error: build.ninja:31: expected '=', got lexing error
on3.12/site-packages/torch/include/torch/csrc/api/include -I/software/...
^ near here
```
Signed-off-by: Alexander Grund <alexander.grund@tu-dresden.de>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
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>
This PR adds AutoEP (Automatic Expert Parallelism) to DeepSpeed training
for HuggingFace MoE models.
AutoEP detects MoE blocks during `deepspeed.initialize()`, builds the
required EP/EDP process groups, and replaces supported MoE blocks with
an EP-enabled execution path, so expert parallelism can be enabled with
DeepSpeed config only and without model code changes.
Current scope in this PR is the base AutoEP feature:
- ZeRO stages 0, 1, and 2 support
- checkpoint save/load support
- universal checkpoint conversion support
ZeRO-3 extensions are intentionally left as follow-up work (#7928 should
be merged for this work)
Supported presets in this PR:
- Mixtral
- Qwen3-MoE
- DeepSeek-V2
- DeepSeek-V3
For end-to-end benchmarking and testing, an AutoEP example is available
in DeepSpeedExamples:
-
<https://github.com/tohtana/DeepSpeedExamples/tree/tohtana/add_auto_ep/training/expert_parallel>
## Attribution
This implementation substantially builds on TorchTitan's MoE /
expert-parallel implementation, and we want to explicitly acknowledge
that prior work.
The TorchTitan-derived pieces in this PR are primarily:
- `deepspeed/moe/ep_router.py`: adapted from TorchTitan's
`TokenChoiceTopKRouter`
- `deepspeed/moe/ep_experts.py`: adapted from TorchTitan's
`GroupedExperts` and grouped-GEMM expert execution path
- `deepspeed/moe/ep_kernels.py`: adapted from TorchTitan's
`TokenReorderer`, `generate_permute_indices`, Triton fill-indices
kernel, and token-group alignment / padding helpers
- `deepspeed/module_inject/auto_ep_layer.py`: adapts the same router ->
reorder -> dispatch -> local expert compute -> combine structure used in
TorchTitan's MoE / EP flow
Relevant TorchTitan sources:
-
<https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/moe.py>
-
<https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/kernels.py>
-
<https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/utils.py>
-
<https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/expert_parallel.py>
The DeepSpeed-specific work in this PR is the AutoEP integration layer
around those building blocks:
- HuggingFace MoE detection and structural validation
- model-family presets and custom-config path
- weight repacking from HF expert layouts into grouped expert tensors
- DeepSpeed runtime group setup and module replacement
- DeepSpeed checkpoint save/load and universal checkpoint support
- DeepSpeed docs and tests
## Design
The implementation is split into a few layers:
- `deepspeed/module_inject/auto_ep_config.py`
- user config parsing
- built-in model presets
- validation for EP topology and per-model constraints
- `deepspeed/module_inject/auto_ep.py`
- scans the model for MoE blocks
- validates the detected structure
- builds a `MoELayerSpec` for each supported MoE layer
- replaces the original HF block with `AutoEPMoELayer`
- `deepspeed/module_inject/auto_ep_layer.py`
- the drop-in execution wrapper for a detected MoE block
- implements router execution, token reorder, EP dispatch/combine, local
expert compute, and shared-expert merge
- `deepspeed/moe/ep_router.py`, `deepspeed/moe/ep_experts.py`,
`deepspeed/moe/ep_kernels.py`
- reusable MoE runtime pieces for routing, grouped expert compute, token
permutation, and aligned grouped-GEMM execution
- `deepspeed/moe/ep_repack.py`
- converts HF expert weights into the grouped expert layout expected by
the runtime
- `deepspeed/runtime/engine.py` and checkpoint conversion code
- wires AutoEP into `deepspeed.initialize()`
- handles checkpoint save/load metadata and universal checkpoint
integration
At runtime, the execution path is:
1. detect and replace supported HF MoE blocks during initialization
2. route tokens with the EP router
3. reorder tokens by expert assignment
4. perform all-to-all dispatch across the EP group when `autoep_size >
1`
5. run local grouped expert compute
6. all-to-all combine and restore the original token order
7. merge shared experts if the model has them
## Adding new model support
There are two supported ways to extend AutoEP to a new MoE model family.
1. Add a preset in `PRESET_MODELS`.
This is the preferred path for a model family we want to support out of
the box. A preset defines:
- MoE layer pattern
- router child name
- experts child name
- expert weight names / layout
- `num_experts` and `top_k` config attributes
- routing defaults
- optional shared-expert structure
2. Use the custom config path.
For models that are not yet built into DeepSpeed, AutoEP can be driven
from config with:
- `moe_layer_pattern`
- `router_pattern`
- `expert_pattern`
- `expert_w1`, `expert_w2`, `expert_w3`
- `num_experts_attr`
- `top_k_attr`
- optional shared-expert fields
Once detection can produce a valid `MoELayerSpec`, the replacement,
execution, and checkpoint paths are shared.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Co-authored-by: Guokai Ma <guokai.ma@intel.com>
DS4Sci EvoformerAttention currently depends on CUTLASS, but requiring
users to manually set `CUTLASS_PATH` creates unnecessary friction for an
otherwise standard extension build flow. This change makes CUTLASS
discovery automatic while preserving `CUTLASS_PATH` as the explicit
override.
The discovery approach is based on PyTorch's CUDA detection pattern in
`torch.utils.cpp_extension`: honor the explicit environment variable
first, then infer from installed packages and conventional filesystem
locations, and only fail with an actionable message when discovery
cannot succeed.
This improves first-run usability, CI behavior, editable installs, and
package-based environments where CUTLASS may already be installed in a
discoverable location. It also reduces setup divergence between users
who clone CUTLASS manually and users who install NVIDIA's
`nvidia-cutlass` package.
DeepSpeed should already have had this because EvoformerAttention is
part of DeepSpeed's extension-builder system, and extension builders
should locate common build dependencies using predictable heuristics
instead of requiring users to export paths manually. CUDA itself is not
treated as "you must always set CUDA_HOME"; PyTorch attempts discovery
first and uses the env var as a fallback. CUTLASS should follow the same
principle here.
---------
Signed-off-by: Max Tretikov <max@tretikov.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
## 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>
Hi DeepSpeed team — thanks for your time reviewing this PR.
## Summary
Add Universal Checkpoint (UC) metadata support for DeepSpeed AutoTP to
enable saving and resuming from Universal Checkpoints.
## Motivation
AutoTP partitions parameters across TP ranks. To make checkpoints
portable and restorable, we need a stable UC metadata representation
that can be collected at save time and consumed at restore time.
## What’s in this PR
- Collect AutoTP-specific Universal Checkpoint metadata for
TP-partitioned parameters.
- Provide restore/merge helpers that normalize shapes and correctly
interpret the saved conversion/partition view.
- Keep existing (non-AutoTP / non-UC) checkpoint paths unchanged (no
behavior change expected for other users).
## Testing
- `pytest -q
tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py`
- `pytest -q tests/unit/checkpoint/test_autotp_universal_checkpoint.py`
## Request for feedback
Could you please take a look at the UC metadata schema and let me know
if you’d prefer any changes to naming, field placement, or compatibility
expectations? I’m happy to iterate quickly based on your guidance.
## References
- Refs: #7861 (Q2 2026 roadmap — AutoTP Universal Checkpoint support)
---------
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: nathon-lee <248585198+nathon-lee@users.noreply.github.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Fixes#7863
Replaces #7872
@Flamefire
Issue #7863 reports order-dependent failures in Evoformer when building
for mixed CUDA architectures. The guard-only approach prevents some bad
outputs but does not solve multi-generation packaging requirements.
This PR takes the root-cause direction: produce a correct multi-arch
binary that can run on pre-Ampere and Ampere+ and select the right
kernel family at runtime.
With TORCH_CUDA_ARCH_LIST='7.0;8.0':
1. Build is no longer pinned by -DGPU_ARCH; it uses runtime arch
dispatch (evoformer_attn.py:33, gemm_kernel_utils.h:53).
1. Runtime chooses implementation by device CC:
- CC >= 80 -> Sm80 (Ampere+ path)
- CC >= 75 -> Sm75
- CC >= 70 -> Sm70
1. So pre-Ampere uses pre-Ampere kernels, and Ampere+ uses the
Ampere-family kernel path.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
With Intel Extension for PyTorch retiring, XPU device would be supported
by PyTorch 2.8+ and dependency to Intel Extension for PyTorch would not
be needed.
This PR removed IPEX dependency, adapt to builder protocol in PyTorch
for XPU, and updated documents and tests accordingly.
Note after this update, DeepSpeed will not work with previous
PyTorch+IPEX on XPU devices. Suggest user to upgrade to latest PyTorch
to get latest XPU features on XPU devices.
Come with this PR is removal of InferenceBuilder, the kernel needed by
InferenceBuilder is supported through Intel Extension for PyTorch.
---------
Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
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>
### Description
This PR addresses Issue #7672.
When sequence_parallel_size is smaller than world_size (e.g., sp_size=2
on 4 GPUs) with PyTorch < 2.3, using
torch.distributed.nn.functional.all_gather for loss aggregation triggers
an IndexError: tuple index out of range during the backward pass. This
is due to a known PyTorch issue where the backward hook accesses the
global rank instead of the group rank.
### Solution
1. Regression Test & Workaround: Updated the regression test
TestUlyssesLossBackward to implement a Weighted All-Reduce pattern.
- Before: all_gather -> manual sum (Vulnerable to rank indexing mismatch
on older PyTorch).
- After: all_reduce(weighted_loss) / all_reduce(total_weight) (Robust
and supports weighted averaging).
2. Runtime Warning: Added a version check (required_torch_version) in
DeepSpeedEngine. It now logs a warning if Sequence Parallelism is
enabled on PyTorch < 2.3, providing a link to the workaround test case.
3. Documentation: Updated ulysses-alst-sequence-parallelism.md with a
note regarding legacy PyTorch versions and the recommended workaround.
### Verification
Added and verified the regression test
tests/unit/sequence_parallelism/test_ulysses.py which now validates the
weighted averaging logic.
**1. Reproduction (Before Fix)**
Confirmed IndexError crash on Rank 2/3 with sp_size=2 on a 4-GPU setup.
<img width="1370" height="860" alt="Screenshot 2026-01-23 at 23 53 42"
src="https://github.com/user-attachments/assets/f4005c02-ff6c-46ea-a1a7-caac2093128b"
/>
**2. Verification (After Fix)**
Verified the fix using the regression test logic on 4x RTX A6000. The
backward pass now completes successfully on all ranks without error.
<img width="1192" height="605" alt="Screenshot 2026-01-23 at 23 52 54"
src="https://github.com/user-attachments/assets/c14cd093-67b7-42b0-ae15-65555c129082"
/>
---------
Signed-off-by: vensen <vensenmu@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Two spelling errors in docs/_tutorials/accelerator-setup-guide.md:
Line 50: comma-seperated-dash-range → comma-separated-dash-range
Line 97: optimzied → optimized
Both typos are in the Intel Architecture CPU section of the accelerator
setup guide.
Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>
`EvoformerAttnBuilder` has some problems which preclude compiling the
extension on several scenarios (e.g., [isolated conda environment with
cuda toolchain](https://github.com/aqlaboratory/openfold-3/pull/34),
lack of hardware in the system) and breaks some standard DeepSpeed
configuration of target capabilities.
*Changes*
- Fix evoformer CUTLASS detection:
- Allow to skip it, useful when CUTLASS is already correctly setup
(e.g., in a conda environment with CUTLASS and the CUDA toolchain)
- Fix misleading use of deprecated nvidia-cutlass pypi package by
actually using the provided bindings but discouraging this route as
[these bindings are not maintained
anymore](https://github.com/NVIDIA/cutlass/discussions/2119)
- Fix evoformer compilation with no GPU is present:
- this is taken care correctly and more generally by
builder.compute_capability_args
- allow for cross-compilation in systems without GPU
- allows for compilation against all available virtual architectures and
binary outputs
- see e.g., https://github.com/deepspeedai/DeepSpeed/issues/5308
- Make all these changes configurable and explicit through documented
environment variables
Tested in all scenarios.
---------
Signed-off-by: Santi Villalba <sdvillal@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
docs: update Megatron-DeepSpeed tutorial to match current repo structure
- Update outdated file paths and script names in
`docs/_tutorials/megatron.md`.
- Replace `scripts/` with `examples/` for training scripts.
- Replace `pretrain_gpt2.py` with `pretrain_gpt.py`.
- Correct locations for `arguments.py` and `utils.py` to `megatron/`.
- Ensure tutorial instructions align with the latest Megatron-DeepSpeed
repository layout.
Resolves#7757
---------
Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
As I was integrating ALST/Ulysses SP into HF Accelerate/Trainer I
noticed that the initial
`UlyssesSPAttentionHF.register_with_transformers` API was a bit
inflexible/confusing wrt variable seqlen.
This PR deprecates the misleading `max_length` arg name, replaces it
with `seq_length` and makes the latter optional if
`seq_length_is_variable` is True.
Updated tests and docs.
Signed-off-by: Stas Bekman <stas@stason.org>
We are a team at Argonne National Laboratory working on low-overhead
asynchronous checkpointing approaches for LLMs and transformers. As part
of these efforts, we have developed DataStates-LLM, a library that we
would like to contribute to the DeepSpeed community:
https://github.com/datastates/datastates-llm
The key idea we leverage is to allow non-blocking tensor copies during
the forward and backward pass from the GPU to the host. Only if these
copies do not finish until the update phase, then we block. Meanwhile,
from the host memory, the tensors are flushed asynchronously to durable
storage (parallel file systems, local SSDs, etc).
To enable this capability, our initial implementation makes the
scheduler aware of checkpointing, calling a ckpt.wait() primitive before
starting the update phase. We illustrated this with the pipeline
scheduler. We are also considering a scheduler-independent solution that
integrates with DeepSpeed/Megatron and provides a hook for the start of
the update phase, which we can leverage to run ckpt.wait().
We appreciate your feedback and look forward to a collaboration in this
space.
---------
Signed-off-by: amaurya <amaurya@anl.gov>
Co-authored-by: amaurya <amaurya@anl.gov>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
MoE tutorial fixes:
1. cifar example has been moved - fix the url
2. fixing text and improving markup
---------
Signed-off-by: Stas Bekman <stas@stason.org>
This PR adds ZenFlow, a importance-aware offloaded training framework
for DeepSpeed ZeRO. ZenFlow enables multi-step overlap between
computation and communication during offloaded training, improving GPU
utilization and reducing stalls.
Highlights:
- New ZenFlow optimizers (ZenFlowCPUAdam, ZenFlowSelectiveAdamW)
- ZenFlowZeroOptimizer for ZeRO Stage 1/2 integration
- Configurable via ZenFlowConfig, integrated with DeepSpeedZeroConfig
- Unit tests and documentation included
Note: This PR focuses on Stage 1 and 2 integration. Stage 3 support will
be introduced in a follow-up PR.
---------
Signed-off-by: Tingfeng Lan <erc8gx@virginia.edu>
Signed-off-by: Yusen Wu <xrn4ub@virginia.edu>
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Yusen Wu <xrn4ub@virginia.edu>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>
Co-authored-by: Guokai Ma <guokai.ma@gmail.com>
In ZeRO offload, significant time is spent on CPUAdam, which is CPU
code. Thus use `--bind_cores_to_rank` in deepspeed launch command would
help improve the performance of ZeRO offload. This PR add this command
to ZeRO offload tutorial to increase user awareness.
For Qwen2.5-3B finetuning on 2 A100-40B cards, running on CPU host with
128 CPU cores, the average step time is as follow, near 1.3x performance
improvement:
without `--bind_cores_to_rank`: 3084.44ms per step
with `--bind_cores_to_rank`: 2383.16ms per step
---------
Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>
The project has been renamed at the last moment, so this PR is adapting
to that change.
There are no code changes in this PR, just docs.
---------
Signed-off-by: Stas Bekman <stas@stason.org>
Improve or fix some minor indentation, typo, and list numbering issues
of the Ulysses Plus tutorial.
---------
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
As title says.
Default behavior of arctic model produces shape issues with AutoTP due
to the MLP layer performing `w2 * act(w1*w3)`. However, method provided
to fix Mixtral-7x8b in #5257 does not work since the MLP for Arctic is
also used within a ModuleList for the MoE. This results in MLP weights
hiding behind individual experts as layers `#.w#`, which is not caught
by the fix in #5257. This adds the check directly within replace, where
it can check for actual layer names for the `w2` key in the model to
patch with `all_reduce`.
---------
Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Restoring the functionality of the cpu locked tensor in the AIO library.
Make async_io operator available for CPU accelerator, i.e., CPU only
environment.
---------
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Zhipeng Wang