## 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
Adds a native host-memory pinning backend, selectable via the
`DS_PIN_MEMORY_BACKEND` environment variable (defaults to `torch`). When
set to `native`, CPU memory is page-locked through the standalone
DeepSpeed `pin_memory` op (`PinMemoryBuilder` / `pin_handle`,
`posix_memalign` + `mlock`) instead of `torch.pin_memory()`.
Stacked on #8236 (standalone `pin_memory` op, now on `master`). Native
allocations go through `pin_handle`, so DeepNVMe I/O handles recognize
them via the process-wide manager and skip bounce buffers — without
requiring libaio / AIO worker threads.
- **New `deepspeed/utils/pin_memory.py`**: a process-wide shared
`NativePinnedMemory` manager that pins CPU memory, tracks pinned pointer
ranges (so slices/views report as pinned), tags buffers with
`.ds_pinned`, supports `make_copy`/`match_shape`, and frees on unpin. It
fails early with a clear error if the `pin_memory` op cannot be built
(no silent torch fallback). Native pins also use a `weakref` finalizer
so GC releases mlocked pages when tensors are dropped without an
explicit unpin.
- **Accelerator owns dispatch**: `pin_memory` drops `align_bytes` and
gains `make_copy`/`match_shape`; `is_pinned` is FakeTensor/meta-tensor
safe; new `unpin_memory` (native frees, torch no-op). Subclasses retain
only the device-specific `_torch_pin_memory`/`_torch_is_pinned`
primitives. Preserves master's `track_pinned_memory` accounting (CPU
torch no-op still bypasses it).
- **Consolidation**: XPU's bespoke `align_bytes=0` path is folded into
the shared native backend.
- **Callers**: `compile` paths route through `get_accelerator()`.
Swap-tensor buffers continue to allocate via I/O handles; with the
shared manager they interoperate with native-pinned tensors. ZeRO /
ZenFlow `destroy()` explicitly unpins optimizer-owned CPU-offload
buffers under the native backend.
- **Docs**: Host Memory Pinning section under RTD Memory Usage
(`docs/code-docs/source/memory.rst`).
- **Tests**: unit tests for the native manager, accelerator pinning
APIs, destroy-path unpin, and cross-op recognition with AIO.
## Test plan
- [x] Rebased onto `master` after #8236 merge; retargeted
`NativePinnedMemory` from `AsyncIOBuilder` → `PinMemoryBuilder`.
- [x] `pre-commit` on changed files.
- [x] Focused UTs on GPU (`tunji-h200-n1g2-ds2-0`, job
`20260809T183311Z`): `tests/unit/v1/pin_memory/` +
`tests/unit/v1/accelerator/test_accelerator.py` +
`tests/unit/v1/nvme/test_pinned_manager.py` — **30 passed**.
- [x] Bounce-buffer / cross-op smoke: under
`DS_PIN_MEMORY_BACKEND=native`, a `pin_handle` buffer is `is_pinned` on
a separate AIO handle.
Made with [Cursor](https://cursor.com)
---------
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.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>
## Problem
ZeRO-3's gradient GPU→CPU offload in `partition_grads()` used a blocking
`copy_()` without `non_blocking`, and its destination buffer defaulted
to
pageable host memory. This forced the offload onto a synchronous,
low-bandwidth (staged pageable) path with no overlap against backward
compute, even though the copy already runs on the dedicated
`reduce_and_partition_stream`.
ZeRO stage 1/2 already issues this copy with `non_blocking=True`
(`stage_1_and_2.py:1530`); stage 3 is the inconsistent one.
## Changes
- `offload_config.py`: default
`offload_optimizer/offload_param.pin_memory`
to `True`. Pinned (page-locked) host memory is required for async,
full-bandwidth DMA; the prior `False` silently selected the slow staged
pageable copy. Disable only on hosts with tight `ulimit -l` memlock.
- `stage3.py`: issue the grad offload copy with `non_blocking=True`.
Stays
on `reduce_and_partition_stream`.
- `stage3.py` (2nd commit): remove an orphaned helper
(`async_inplace_copy_grad_to_fp32_buffer_from_gpu`) that referenced an
uninitialized attribute and had no callers. The live stage 1/2 version
is untouched.
## Validation
4× RTX 4080-SUPER, autotp=2, `offload_optimizer`, `cpu_adam`, per-rank
CPU affinity:
| model | baseline BWD | fixed BWD | Δ BWD |
|----------------|--------------|-----------|-------|
| Qwen2.5-1.5B | 1810 ms | 1308 ms | -28% |
| Qwen2.5-3B | 3097 ms | 2540 ms | -18% |
Memory footprint unchanged; FWD/STEP unchanged.
---------
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
## Summary
- Extends unmanaged gradient accumulation
(`managed_gradient_accumulation=false`) to **ZeRO stage 3**
(non-offload). Stacked follow-up to #8203 (now merged).
- Stage 3 partitions gradients into `__param_id_to_grad_partition` on
every `backward()` (as in managed mode); `step()` only marks the
boundary via `finalize_gradient_accumulation_boundary()`.
- Since stage 3 is now permitted, also reject ZeRO **parameter** offload
(previously unreachable behind the stage-3 guard). ZeRO optimizer-state
offload, pipeline parallelism, DeepCompile, and Apex AMP remain
unsupported pending follow-ups.
## Test plan
Validated on a 2-GPU node:
- [x] Full `-k Unmanaged` suite (29 passed), including
`test_unmanaged_matches_managed[3]` and
`test_unmanaged_varying_backward_count[3]`
- [x] Stage-3 unmanaged feature subset (10 passed)
- [x] Core ZeRO-3 regression (`test_zero_grad_clip`,
`test_zero_leaf_module`, `test_zero_context`,
`test_zero_tensor_fragment`): 158 passed, 62 skipped (expected)
Docs (`config-json.md`, `training.rst`) updated; previewable on
`rtd-staging`.
Made with [Cursor](https://cursor.com)
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
## Summary
Extends unmanaged gradient accumulation
(`managed_gradient_accumulation=false`) to **ZeRO stage 2**. Stacked on
top of the stage 0/1 foundation in #8184 (base branch
`sfc-gh-truwase/gas_mgmt_zero01`); review that PR first.
Unlike stage 0/1 (where `backward()` accumulates locally and `step()`
performs the reduction), ZeRO stage 2 must reduce/partition gradients on
**every** `backward()` to preserve its memory characteristics. This is
compatible with unmanaged mode because reduce-scatter is linear:
accumulating the reduced partitions across N caller-controlled backwards
is equivalent to reducing once at the boundary. Micro-step tracking
stays disabled and the caller still owns the boundary; only the
`averaged_gradients` finalization is deferred to `step()`.
- `ZeROOptimizer.finalize_gradient_accumulation_boundary()` (stage 1/2)
builds `averaged_gradients` from the accumulated `all_grad_tensors` at
`step()`.
- Validation relaxed to allow stage 2 (`not partition_weights`), still
rejecting stage 3 and ZeRO offload (follow-up PRs).
- **`overlap_comm` is now supported for stage 2**: its async reduction
is confined to the per-backward path and the epilogue synchronizes
before finalizing, so it behaves exactly as in managed mode. It remains
rejected for stage 0/1 (where reduction is deferred to `step()`).
## Test plan
Validated on a 2-GPU node (full `-k Unmanaged` suite, 22 passed):
- [x] `test_unmanaged_matches_managed[2]` — unmanaged stage-2 matches
managed stage-2
- [x] `test_unmanaged_varying_backward_count[2]` — variable backward
count per step, stage 2
- [x] `test_unmanaged_matches_managed_overlap_comm` — unmanaged stage-2
with `overlap_comm=True` matches managed reference
- [x] `test_unmanaged_rejects_stage3`,
`test_unmanaged_rejects_zero_offload` — stage 3 / offload rejected
- [x] `test_unmanaged_rejects_overlap_comm[0,1]` — overlap_comm still
rejected for stage 0/1
- [x] `pre-commit` (yapf/flake8/codespell) clean
Docs (`config-json.md`, `training.rst`) updated to describe stage-2
behavior and the overlap_comm support; previewable on `rtd-staging`.
Made with [Cursor](https://cursor.com)
---------
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
## Summary
DeepSpeed's activation checkpointing wrapper currently accepts
positional
arguments only, while `torch.utils.checkpoint.checkpoint` also supports
keyword
arguments. This change allows callers to pass keyword arguments through
`deepspeed.checkpointing.checkpoint`.
Keyword names and non-Tensor values are retained for reconstruction
during the
forward and recompute passes. Tensor keyword values are flattened into
the
inputs passed to `CheckpointFunction`, so autograd tracks them and
returns their
gradients correctly.
The activation checkpointing documentation now describes keyword
argument
support, and the regression test covers both Tensor and non-Tensor
keyword
arguments as well as gradient propagation.
## Validation
- Activation checkpointing unit tests: 27 passed
- Pre-commit checks for all changed files: passed
- 1-GPU CUDA correctness smoke: direct and checkpointed execution
matched, with
zero maximum gradient error
- 2-GPU distributed CUDA correctness smoke: each rank matched direct
execution,
with zero maximum gradient error
- DCO sign-off is included in the commit
Fixes#7038
Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
When managed_gradient_accumulation=false, disable micro-step tracking
and treat each engine.step() as the accumulation boundary: reduce
locally accumulated grads then apply the optimizer update. Stage 2/3 and
pipeline remain unsupported in this change.
Part fix for #8183
---------
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>
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>
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>
## Summary
Adds `engine.coalesce_grad_reduction()`, an opt-in context manager that
defers ZeRO 1/2/3 gradient reduction across multiple `engine.backward()`
calls inside one optimizer step. On context exit, a single reduction
pass populates `averaged_gradients` for the next `engine.step()`.
This is the third step of the multi-backward feature:
- #7665 made `set_gradient_accumulation_boundary()` plus manual
`engine.backward()` (PyTorch-style backward) a first-class API.
- #7981 fixed a silent gradient-loss bug on that path for ZeRO-1/2 with
`cpu_offload`. Chunks 1 through N-1 were dropped at `ga_steps=1 + N>1`
because of an outer gate in `copy_grads_in_partition`.
This PR makes the path efficient. N reduce-scatters collapse into 1
across all ZeRO stages, removing the communication bottleneck that
remained after the correctness fix.
## Motivation
`engine.no_sync()` (engine.py) explicitly asserts ZeRO 2/3 are
incompatible with the no_sync context manager. ZeRO needs per-backward
reduction to partition gradients. The assert enforces that, but it
blocks patterns where multiple `engine.backward()` calls per step are
intentional:
- **Cached contrastive learning** (GradCache): sentence-transformers
`CachedMultipleNegativesRankingLoss`, `CachedGISTEmbedLoss`,
`CachedMultipleNegativesSymmetricRankingLoss` all call
`engine.backward()` once per cached chunk.
- **Custom autograd Functions** that invoke `torch.autograd.backward()`
inside their forward
Both rely on the PyTorch-style backward API from #7665. With that API,
the user (or a custom autograd Function) issues N `engine.backward()`
calls per `engine.step()` and toggles
`set_gradient_accumulation_boundary()` to mark the last one. Without
this PR, the pattern issues N reduce-scatters per step on ZeRO 2/3 even
when the math only needs 1.
## What changed
`deepspeed/runtime/engine.py`
- New `engine.coalesce_grad_reduction()` context manager.
- Stage-aware flush helpers (`_flush_coalesced_reduction_zero{12,3}`).
Iterates params explicitly instead of calling `reduce_gradients()` to
bypass the `overlap_comm` short-circuit and the `contiguous_gradients`
setup_buckets dependency.
`deepspeed/runtime/zero/stage_1_and_2.py` and
`deepspeed/runtime/zero/stage3.py`
- `_coalesce_grad_reduction = False` init plus a 2-line guard at the top
of the per-param reducer entry point. No existing function bodies
modified.
## Compatibility matrix (all bit-exact vs. baseline multi-backward)
All four (contiguous_gradients, overlap_comm) combinations bit-exact vs.
baseline multi-backward:
| Stage | (F, F) | (T, F) | (F, T) | (T, T) default |
|:-:|:-:|:-:|:-:|:-:|
| ZeRO-1 | OK | OK | OK | OK |
| ZeRO-2 | OK | OK | OK | OK |
| ZeRO-3 | OK | OK | OK | OK |
Additional verified:
- CPU offload (offload_optimizer Z1/Z2/Z3, offload_param Z3).
- BF16 with `gradient_accumulation_dtype=fp32` (Z2 directly, Z1 with
offload via the `use_grad_accum_attribute=True` path).
- FP16 with dynamic loss scaling (Z1/Z2/Z3).
- Multi-bucket flush (small `reduce_bucket_size`).
- MoE smoke (ep_size=1, Z1/Z2). MoE ep_size=2 test included but requires
world_size=4.
- Gradient clipping, multi-step state hygiene.
- N=4 deferred backward issues strictly fewer cross-rank collectives
than baseline (`TestCoalesceCollectiveCount`, patches
`dist.all_reduce`/`reduce`/`reduce_scatter_fn`).
- ZeRO-3 `optimizer.micro_step_id` invariant. Stays 0 at flush across
multiple steps, so `partition_grads` always takes the `copy_` branch
instead of the stale-buffer `add_` branch
(`TestCoalesceZero3MicroStepInvariant`).
## Unsupported (NotImplementedError)
- ZeRO stage 0.
- BF16_Optimizer / FP16_Optimizer wrappers. BF16_Optimizer dispatches
only for ZeRO-1 with bf16, `grad_accum_dtype=fp32`, and no offload.
Users on this combo can switch to ZeRO-2.
- PipelineModule (pipeline parallelism schedules its own reductions).
- Reentry / nesting with `engine.no_sync()`.
---------
Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
# AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based
Sequence Parallelism
## Overview
AutoSP is a compiler optimization pass that shards inputs along the
sequence dimension and enables Ulysses styled sequence parallelism while
preventing graph breaks during `torch.compile()`. All the passes operate
at the Torch IR on the forward graph.
## API Design
### User-Facing Entry Point: `prepare_autosp_inputs()`
Users must explicitly call this function to prepare inputs for AutoSP
compilation:
```python
def prepare_autosp_inputs(
input_id: torch.Tensor,
label_id: torch.Tensor,
position_id: torch.Tensor = None,
attention_mask: torch.Tensor = None,
seq_dim: int = 1
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
```
**Purpose**: Symbolize sequence dimension and annotate tensors for
identification.
**Operations**:
1. Mark sequence dimension as dynamic using
`torch._dynamo.decorators.mark_dynamic()`
2. Attach metadata tags for tensor identification for auto-sharding:
- `input_id.tag = constants.INPUT_ID_KEY`
- `label_id.tag = constants.LABEL_ID_KEY`
- `position_id.tag = constants.POSITION_ID_KEY` (if provided)
**Rationale**: PyTorch's FX graph tracer requires explicit annotation of
data-dependent dimensions. Marking the sequence dimension as dynamic
prevents symbolic shape propagation from losing dimension information
through reshape/view operations.
## Compilation Passes
### Pass 1: `pass_shard_seq_dim()`
**Objective**: Propagate sharded sequence dimension to all consumers.
**Algorithm**:
1. Extract symbolic sequence dimension from `input_id` shape metadata
2. Locate the symbolic dimension node in the FX graph
3. Create a floor-divide node: `seq_dim / world_size`
4. Perform worklist-based graph traversal to find all direct and
indirect consumers of input node, label node and position id node.
5. Replace symbolic dimension references with sharded dimension in
consumer nodes
**Rationale**: Reshapes and views that consume the sequence dimension as
an argument do not get updated during propagation of symbolic shapes.
This pass explicitly rewires the computation graph to use sharded
dimensions, enabling proper shape inference downstream.
### Pass 2: `pass_shard_input_ids()` / `pass_shard_label_ids()` /
`pass_shard_position_ids()`
**Objective**: Insert slicing operations after input tensors.
**Implementation**: Call `shard_tensor_node()` utility which inserts
slice operations. Each rank retains only the portion of the tensor
corresponding to its sequence partition and drops the remaining buffer.
**Note on `attention_mask`**: Not sharded because it applies to the full
sequence length, not the partitioned dimension.
### Pass 3: `pass_insert_attention_all_to_all()`
**Objective**: Insert all-to-all collectives around attention (Ulysses
styled) to avoid graph breaks during compilation.
**Algorithm**:
1. Identify all SDPA (Scaled Dot-Product Attention) nodes in the graph
2. For each SDPA node with inputs Q, K, V, after each of Q, K, V: insert
A2A scatter heads (dim=1), gather sequence (dim=2)
3. Insert A2A after thre attention output O: scatter sequence (dim=2),
gather heads (dim=1)
**Graph Rewrite Example**:
```
Q [B, N, S/P, H] --A2A(scatter_heads,gather_seq)--> [B, N/P, S, H]
K [B, N, S/P, H] --A2A(scatter_heads,gather_seq)--> [B, N/P, S, H]
V [B, N, S/P, H] --A2A(scatter_heads,gather_seq)--> [B, N/P, S, H]
|
SDPA
|
O [B, N/P, S, H] --A2A(scatter_seq,gather_heads)--> [B, N, S/P, H]
```
**Current support**: Currently only supports
`torch.nn.functional.scaled_dot_product_attention()`. Composite
attention patterns require additional pattern matching logic.
### Pass 4: `pass_propagate_shapes()`
**Objective**: Compute static shapes for all nodes using fake tensor
execution.
**Implementation**:
1. Create `ShapeEnv` for symbolic dimension tracking
2. Construct `FakeTensorMode` with the shape environment
3. Execute `FakeTensorProp.propagate()` to compute shape metadata
### Pass 5: `pass_canonicalize()`
**Objective**: Finalize graph representation.
**Operations**:
1. `eliminate_dead_code()`: Remove unused operations
2. `lint()`: Validate graph structure
3. `recompile()`: Regenerate compiled representation
## Execution Order
```
prepare_autosp_inputs()
↓
pass_shard_seq_dim
↓
pass_shard_input_ids
↓
pass_shard_label_ids
↓
pass_shard_position_ids
↓
pass_insert_attention_all_to_all
↓
pass_propagate_shapes
↓
pass_canonicalize
↓
pass_selective_activation_checkpointing
```
## Memory savings
AutoSP adds some heuristics to torch.compile's partitioniner which
splits the joint graph into the forward and backward graph. Matmul and
related ops are not checkpointed since recomputing them is much cheaper
compared to the attention op, while reducing the peak active memory.
## Reducing gradients across ranks
AutoSP requires an all-reduce to reduce the gradients across ranks. This
is automatically called by DeepSpeed's engine
[here](https://github.com/deepspeedai/DeepSpeed/blob/93524c8931799a7631a2321d7ef4afaff6b6e54b/deepspeed/runtime/engine.py#L2433)
## Known Limitations
1. **Attention Pattern Matching**: Only
`torch.nn.functional.scaled_dot_product_attention()` is supported. Fused
attention implementations require pattern-specific handling.
2. **No Graph Break Requirement**: AutoSP will fail if there are graph
breaks because use-def chains are lost and it becomes tricky to
propagate auto-sharding information across graph modules.
## Example
DeepSpeedExample PR:
https://github.com/deepspeedai/DeepSpeedExamples/pull/999
---------
Signed-off-by: Neel Dani <neeldani98@gmail.com>
Signed-off-by: Ahan Gupta <ahangupta.96@gmail.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Ahan Gupta <ahangupta.96@gmail.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Co-authored-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.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>
Add a document section clarifying the behavior of nesting autocast and
why/when we need it.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.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>
DeepSpeed optimizer always creates fp32 master
params/gradients/optimizer states.
However, we sometimes want to keep them lower precision given
[torch.autocast
support](https://deepspeed.readthedocs.io/en/latest/training.html#mixed-precision-training).
This PR allows lower precision master params/grads/optimizer states with
bf16/fp16 enabled.
DeepSpeed currently accepts `fp16_master_weights_and_gradients` option
under `fp16` section (not documented) with ZeRO1/2. This PR extends this
for bf16 and also ZeRO3.
In `bf16` section, we can have new items `bf16_master_weights_and_grads`
and `bf16_optimizer_states`.
Similary to `fp16_master_weights_and_grads`,
`bf16_master_weights_and_grads` keeps master parameters in bf16.
`bf16_optimizer_states` keeps optimizer states also in bf16. Here is an
example configuration:
```json
"bf16": {
"enabled": true,
"bf16_master_weights_and_grads": true,
"bf16_optimizer_states": true
}
```
Note that `bf16_master_weights_and_grads==True` and
`bf16_optimizer_states==False` is supported only with cpu offloading.
Also, we don't have `fp16_optimizer_states` as it won't be practical.
More details are described in
[`config-json.md`](https://github.com/tohtana/DeepSpeed/blob/88e0bbdfba89c4712d815980ddb28353d6da5b2e/docs/_pages/config-json.md)
Previously, `torch.autocast` support (`torch_autocast` section in
config) was not compatible with `bf16` `fp16` enabled, but we now accept
the combination.
This PR also adds some test cases for the configurations as well as the
combination with `torch.autocast`.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
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>
This PR improves the usability of the leaf module feature.
Here are the changes:
- Allow enabling the leaf module via both the DeepSpeed config and APIs.
- Relax matching criteria to support class-based matching.
- Support multiple ways of specifying the target module: class, class
name (with or without package name), module name, or suffix.
- Add documentation to the training guide, including config snippets and
explanations of default behavior.
- Add default classes (e.g., Mixtral, Qwen2/Qwen3) that automatically
enable the leaf module feature. (Welcoming requests to add more classes)
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Support training multiple models, such as in
[HF](https://huggingface.co/docs/accelerate/en/usage_guides/deepspeed_multiple_model)
Here is some update on supporting multiple DS engines with single
loss.backward(). The main message is that I think we can support this.
First, some context. Backward pass in ZeRO is complicated because the
optimizations/features require special handling of gradients, such as:
1. Gradient partitioning
2. Overlapping backward and reduction
3. Upcasting for fp32 grad accumulation
So, we created engine.backward(loss) as a wrapper function to provide us
fine-grained control over backward as below
```python
def backward(loss):
backward_prologue() # setup logic for special gradient handling
loss.backward()
backward_epilogue() # cleanup/teardown logic
```
As demonstrated by @muellerzr, this approach breaks down when loss
originates from multiple DS engines. Our proposed solution is to use
backward hooks on the module to launch backward_prologue() and
backward_epilogue() . Specifically,
1. backward pre hook on engine.module to launch backward_prologue()
before any module gradient is created.
2. backward post hook on engine.module to launch backward_epilogue()
after all module gradients are created.
We plan for this solution to preserve BC, i.e., engine.backward() will
remain correct for single engine scenarios.
The current status is that (1) is completed, while (2) is in progress.
To unblock e2e testing for multi-engine scenarios, since there are
probably other issues, we have a temporarily added
engine._backward_prologue() . You can try this out via the following
artifacts.
1. Simple multi-engine test code:
https://gist.github.com/tjruwase/f1adccf087b8fa269ffce2ab91c4f1c6#file-multi_engine-py
2. DS branch:
https://github.com/microsoft/DeepSpeed/tree/olruwase/zero_multi_models
---------
Signed-off-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
This PR adds an API `deepspeed.runtime.zero.offload_states
get_state_devices`, which gets devices of offload states as suggested in
this
[comment](https://github.com/microsoft/DeepSpeed/pull/6011#issuecomment-2358068777).
We could lift this up to `deepspeed.utils` but would need to resolve a
circular import: User code -> `deepspeed.utils` ->
`deepspeed.utils.offload_states` -> `deepspeed.runtime.zero` ->
`deepspeed.runtime.zero.partition_parameters` -> `deepspeed.utils`
This will require a significant refactoring as long as we have
`OffloadStateTypeEnum` in `deepspeed.runtime.zero`.
---------
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
This PR introduces a new monitoring option - `CometMonitor` which comes
up as an official integration with
[CometML](https://www.comet.com/site/).
The new monitor is covered with unit tests.
Notes:
* We've updated `docs/code-docs/source/monitor.rst` but it doesn't look
used anymore
* We've updated the "Monitoring Module" section name in `config-json.md`
to be generic so the next integration won't require updating it.
---------
Co-authored-by: Boris Feld <lothiraldan@gmail.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
The DeepSpeed currently supports a set of debugging APIs to
[get](https://deepspeed.readthedocs.io/en/latest/zero3.html#debugging)
and
[set](https://deepspeed.readthedocs.io/en/latest/zero3.html#modifying-partitioned-states)
the **full** model states (parameters, gradients, and optimizer states).
However, in some scenarios, only **local states** are needed, for
example, when pruning some model layers based on a local criterion.
After calling `model_engine.step()`, we need to apply the local mask to
the partitioned parameters owned by each process. Therefore, I am
submitting this PR to introduce some new APIs for `get` and `set` ZeRO-3
partial model states.
### APIs intro
```python
def safe_get_local_fp32_param(param):
"""Get the fp32 partitioned parameter."""
def safe_get_local_grad(param):
"""Get the fp32 gradient of a partitioned parameter."""
def safe_get_local_optimizer_state(param, optim_state_key):
"""Get the fp32 optimizer state of a partitioned parameter."""
def safe_set_local_fp32_param(param, value):
"""Update the partitioned fp32 parameter."""
def safe_set_local_optimizer_state(param, value, optim_state_key):
"""Update the fp32 optimizer state of a partitioned parameter."""
```
### Usage
```python
# local API
from deepspeed.utils import (
safe_get_local_fp32_param,
safe_get_local_grad,
safe_get_local_optimizer_state,
safe_set_local_fp32_param,
safe_set_local_optimizer_state
)
```
### TODO
- [x] Add local APIs
- [x] Add UTs
- [x] Update Docs
@tjruwase
---------
Signed-off-by: yliu <test@do_not_reply@neuralstudio.intel.com>
Co-authored-by: yliu <test@do_not_reply@neuralstudio.intel.com>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
* include mics config and optimizer
* change private vars to public vars
so the child class can initialize these vars
* Port the init function from stage3
* adding a model test file for mics
* adopt to get_acceleartor api and fp16 group defrag
* WIP: porting mics modification to ms master
* WIP: included gradient all-reduce among replication groups
* WIP: ported hierarchical all gather part
did basic loss test on a simple MLP model
* [Bug fix] using the comm group attached on the param
* torch2.0 support
* remove print
* delegate wait op
* [Bug] fix naming
* adding doc string
* resolving recursive import
* fix formating, typo and license
* fix license and unit test error
---------
Co-authored-by: Ubuntu <ubuntu@ip-172-31-14-191.us-west-2.compute.internal>
Co-authored-by: Ubuntu <ubuntu@ip-172-31-7-70.us-west-2.compute.internal>
Co-authored-by: Zhen Zhang <zhzhn@amazon.com>
Co-authored-by: zhzhn <zhzhn@ip-10-2-57-114.us-west-2.compute.internal>
Olatunji Ruwase