github-featured/deepspeedai--deepspeed
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity
3,325 Commits 864 Branches 127 Tags
master
Commit Graph

3325 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Zhipeng Wang 64fcec6ba7 Enable DeepSpeed support on Apple Silicon (MPS) with ZeRO Stage 1-3 (#8293) 
## 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>
 2026-08-23 09:13:43 +00:00
Vedant Chauhan ebb75d7e81 Fix repeated FlopsProfiler metric accumulation (#8246) 
## Summary

  - restore `F.scaled_dot_product_attention` after profiling
  - restore `Tensor.__matmul__` after profiling
  - restore `torch.bmm` from its correct saved implementation
- add a CPU regression test covering repeated profiling sessions and
operation restoration
  - scope the existing FP16 skip to the test that actually requires FP16

  ## Problem

`FlopsProfiler` temporarily replaces PyTorch operations with
FLOP-counting wrappers. Its cleanup path did not restore
`F.scaled_dot_product_attention` or
`Tensor.__matmul__`, causing wrappers to accumulate across profiling
sessions. As a result, identical model executions reported progressively
increasing FLOPs and
  MACs.

  The `torch.bmm` cleanup was also incorrect:

  ```python
  torch.bmm = old_functions[torch.matmul.__str__]
 ```
Because torch.matmul had already been restored, this rebound torch.bmm
to torch.matmul for the rest of the process. This changed normal PyTorch
behavior after
profiling: torch.bmm began accepting inputs supported by broadcasting
matmul but invalid for bmm.

  ## Fix

Restore every patched operation from its matching saved original
function. The patch and cleanup paths now cover the same set of
operations with matching PyTorch
  version guards.

  ## Testing

A CPU regression test profiles the same scaled-dot-product-attention
operation three times and verifies:

  - every session reports identical FLOP and MAC totals
- F.scaled_dot_product_attention, Tensor.__matmul__, and torch.bmm are
restored to their original function objects after every session

  Observed totals:

  Before:
  [(65536, 32768), (131072, 65536), (196608, 98304)]

  After:
  [(65536, 32768), (65536, 32768), (65536, 32768)]

  All pre-commit checks pass.

  Fixes #7413

Signed-off-by: Vedant Chauhan <staranonymous1011@gmail.com>
 2026-08-23 07:42:18 +00:00
Ehsan f5671753cb fix: stop DeepSpeedConfig writing max_grad_norm back into the caller's config dict (#8289) 
## What happens

`deepspeed.initialize()` writes into the dict the caller passed as
`config`. When
`optimizer.params.max_grad_norm` is set to a positive value,
`DeepSpeedConfig._do_warning_check`
assigns `0.0` into `self.optimizer_params`, and that is the caller's own
`config["optimizer"]["params"]` object rather than a copy.

Measured in a clean `python:3.11-slim` container at HEAD `11b518a00`,
torch `2.13.0+cpu`,
deepspeed installed with `pip install -e .` from the checkout
(`deepspeed.__file__ = /src/deepspeed/__init__.py`,
`deepspeed.__version__ = 0.19.6+unknown`):

```python
import os, json, copy
os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT="29517",
                  RANK="0", LOCAL_RANK="0", WORLD_SIZE="1")
import torch, deepspeed

cfg = {"train_micro_batch_size_per_gpu": 1,
       "optimizer": {"type": "AdamW", "params": {"lr": 1e-3, "max_grad_norm": 1.0}}}
model = torch.nn.Linear(4, 4)
client_opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

print("BEFORE:", json.dumps(cfg["optimizer"]["params"]))
engine, *_ = deepspeed.initialize(model=model, optimizer=client_opt, config=cfg)
print("AFTER :", json.dumps(cfg["optimizer"]["params"]))
print("gradient_clipping in force:", engine.gradient_clipping())
```

Observed:

```
BEFORE: {"lr": 0.001, "max_grad_norm": 1.0}
[WARNING] [config.py:1068:_do_warning_check] DeepSpeedConfig: In FP32 mode, DeepSpeed does not permit MAX_GRAD_NORM (1.0) > 0, setting to zero
AFTER : {"lr": 0.001, "max_grad_norm": 0.0}
gradient_clipping in force: 1.0
```

Expected: `initialize` leaves the caller's dict as it found it.

Passing a client optimizer is what makes this visible, because
`_configure_basic_optimizer` is
then never called and the usual `ValueError` never fires, so
initialization succeeds with the
caller's config quietly rewritten. The same run without a client
optimizer still raises the
`ValueError`, and still leaves `0.0` behind in the caller's dict,
because the zeroing happens
during config construction and the engine's check tests for the key's
presence rather than its
value.

The warning is also no longer accurate. The value it claims to zero is
not read by anything:
`get_optimizer_gradient_clipping` (`config.py:458`) is its only reader
and has no callers
anywhere in `deepspeed/` or `tests/` (checked with an AST scan for
`Call` nodes, not a text
search). The clipping actually applied comes from `gradient_clipping`,
which is why the run
above reports `1.0`. The engine side of this behaviour was removed in
`abe2204d` (#232, 2020)
and replaced by the hard `ValueError`; the config side predates that
change and was not
revisited with it.

## Why the fix looks like this

`_configure_basic_optimizer` already declares the invariant this line
breaks, at
`engine.py:2088`, added three months ago in `3c337b542` (#8010):

```python
# Copy so the pop() calls below (torch_adam, adam_w_mode, fp32_optimizer_states) do not
# mutate the shared config dict returned by optimizer_params().
optimizer_parameters = dict(self.optimizer_params() or {})
```

Enumerating every writer of that dict across `deepspeed/` by AST
(subscript assignment plus
`update`/`pop`/`setdefault`/`clear`/`popitem` calls):

| | writers |
|---|---|
| before | 4: `config.py:1071` on the caller's dict, plus
`engine.py:2096`, `2097`, `2115` on the copy |
| after | 3: `engine.py:2096`, `2097`, `2115`, all on the copy |

The FP16 and FP32 branches were a pair with the zeroing, so removing it
leaves them without a
distinction to draw. Nothing passes `max_grad_norm` to an FP16 wrapper
today: the only
consumers of a `max_grad_norm` param group are the Lamb and OneBit
optimizers, which take it
as a constructor argument. The two branches therefore collapse into one
warning carrying the
same remedy that `_configure_basic_optimizer` already raises.

If you would rather keep both original messages and drop only the
assignment, or split the
message change into its own PR, say so and I will rework it.

## Tests


`tests/unit/runtime/test_ds_config_dict.py::test_max_grad_norm_leaves_caller_config_untouched`
pins the caller's dict directly. In the same container, on master it
fails with
`assert 0.0 == 1.0`; with this change it passes.

The rest of that file is unaffected: 27 passed, 5 skipped. `TestArgs`
needs `--shm-size` above
the Docker default and fails with `OSError: [Errno 28] No space left on
device` without it, on
master and on this branch alike.

`yapf --style .style.yapf --diff` and `flake8 --config .flake8` are both
clean on the two
changed files.

One limit worth stating: the unit test covers config construction, which
is where the write
happens. The full `deepspeed.initialize` path is covered by the
container run above rather
than by a unit test, since it needs a built comm extension.

Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
Co-authored-by: Guokai Ma <guokai.ma@intel.com>
 2026-08-23 07:40:23 +00:00
Ehsan ad1c516eb6 Fix ZeroDivisionError in compute_elastic_config return_microbatch on non-0.2 elasticity (#8286) 
## Root cause

With `return_microbatch=True`, `compute_elastic_config` takes the `else`
at `elasticity.py:372` for any elasticity version other than `0.2`, and
that branch divides `final_batch_size` by `world_size`. In practice that
means version `0.1`: `0.3` is rejected at line 301 and any other value
raises `NotImplementedError` at line 348, both before this point. The
branch is only reachable when `world_size` is unset, because the `if
world_size > 0:` block above it returns for every positive value. So on
master the division is always by zero and the caller gets a bare
`ZeroDivisionError` at line 375 rather than a configuration error.

Version `0.2` avoids this by resolving `world_size` from the
`WORLD_SIZE` environment variable at lines 321-334, and raising
`ElasticityConfigError` naming that variable when it cannot. Version
`0.1` never reads the environment, so a caller who follows the `0.2`
message's own advice, "set it as an environment variable", still
crashes:

| config | `WORLD_SIZE` in env | master |
|---|---|---|
| 0.2 | yes | returns `(9792, [...], 17)` |
| 0.2 | no | `ElasticityConfigError` naming `WORLD_SIZE` |
| 0.1 | yes | `ZeroDivisionError` |
| 0.1 | no | `ZeroDivisionError` |

## Fix

Resolve `world_size` from `WORLD_SIZE` in the non-`0.2` branch the way
`0.2` already does, and raise `ElasticityConfigError` with the same
guidance when it cannot be resolved. Then check the resolved value
against `valid_gpus` before dividing, matching the sibling block at
lines 352-355; without that check an out-of-range `WORLD_SIZE` would
reach the loop and fail on the `micro_batch_size is not None` assertion
instead of `ElasticityIncompatibleWorldSize`.

Both divisions by `world_size` in this function now run only on a value
that is positive and a member of `valid_gpus`. Nothing that works today
changes: on master this branch raised `ZeroDivisionError` for every
input, and the `0.2` path is untouched.

## Verification

- Three new tests in `tests/unit/elasticity/test_elastic.py` cover the
unset case, resolution from `WORLD_SIZE`, and an out-of-range
`WORLD_SIZE`. All three fail on master with `ZeroDivisionError` at
`elasticity.py:375` and pass here.
- `pytest unit/elasticity/` gives 23 passed, 3 skipped, on Python 3.12
with `torch==2.10.0+cpu` to match the `cpu-torch-latest` leg. The 3
skips are the `DistributedTest` classes that need `FusedLambBuilder`,
and they skip on master too.
- `pre-commit run --files` passes on both changed files, yapf, flake8,
codespell and `check-torchdist` included.
- Not checked: the GPU legs, and the `0.2` `return_microbatch` return at
line 371, which no test in the repo reaches either before or after this
change.

#8162 proposed the same resolution in July, and its author closed it
unmerged on 2026-08-12 without a review.

Fixes #8156

Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
 2026-08-23 07:18:06 +00:00
Olatunji Ruwase 96ffef2cca docs: note async cpu_checkpointing perf and expandable_segments (#8287) 
## Summary

Follow-up to #8282. The docs note for `cpu_checkpointing` landed after
that PR was merged, so this cherry-picks it onto `master`.

Adds a short note under the `cpu_checkpointing` config entry covering:
- The Qwen3-8B single-H200 result: async side-stream copy matches a
blocking offload's peak reduction (up to ~14% at 32K) while staying
within ~2% of no-offload step time.
- Recommending `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` for
very long sequences to avoid allocator fragmentation from the
offload/restore cycle.

## Test plan

- [x] `pre-commit run --files docs/_pages/config-json.md`
- [ ] Docs preview of the Activation Checkpointing `cpu_checkpointing`
section

Made with [Cursor](https://cursor.com)

Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
 2026-08-23 06:57:26 +00:00
Castronaut edaa722126 Fix ZeRO-1/2 with zero-sized parameters (#8280) 
## Summary
- skip zero-sized parameters before ZeRO-1/2 gradient reduction
- add regression coverage for stages 1 and 2

## Testing
- `pytest -q tests/unit/runtime/zero/test_zero_empty_param.py -s`
- original #8279 reproducer on A6000 with PyTorch 2.13.0+cu126 for
ZeRO-1 and ZeRO-2
- `pre-commit run --files deepspeed/runtime/zero/stage_1_and_2.py
tests/unit/runtime/zero/test_zero_empty_param.py`

Fixes #8279

Signed-off-by: Cao Yuhang <caoyuhang@fwerkor.com>
 2026-08-22 12:19:04 +00:00
iLeGend 11b518a00f [typo] Add raise for RuntimeError (#8281) 
Fix missing raise statement for RuntimeError in fusedqkv_utils.py

Signed-off-by: iLeGend <824040212@qq.com>
 2026-08-21 16:36:58 +00:00
pengdurice 9a297cfa74 Enable activation offloading (#8255) 
# Activation offloading for DeepCompile: plan against a floor the run
has actually reached

## What this does

Enables the activation-offload pass to DeepCompile init_z3. The forward
pass saves tensors the backward needs;
they sit in device memory for the whole step. This pass copies chosen
ones to pinned host memory
during the forward and brings each back before the backward reads it.

Enable with `compile.offload_activation: true`. It replaces the prefetch
and selective-gather passes
and is mutually exclusive with `offload_parameters` and
`offload_opt_states` — each plans against the
whole budget on its own.

## Results

Where a job is too tight to run at all, offloading everything rescues it
but costs 40%+ in step
time. Planning which activations to move recovers most of that.

Qwen3-14B, mb4, 8xH200, ZeRO-3, inductor, bf16 optimizer states,
`expandable_segments`. All arms use
DeepCompile; the question is what the offload pass adds.

| seq | baseline (no offload) | forced-all (blind) | ours (planned) |
ours moved | vs forced-all |
| --- | --- | --- | --- | --- | --- |
| 2048 | 1.76s | 2.50s | **1.85s** | 0GB — declines | run fits; pass
correctly does nothing |
| 3072 | 5.46s | 7.81s | **5.41s** | 0GB — declines | run fits; pass
correctly does nothing |
| 3328 | **OOM** | 8.67s | **7.38s** | 11.7GB of 22.0 | **15% faster** |
| 3584 | **DIED** (watchdog) | 10.80s | **9.76s** | 20.5GB of 23.7 |
**9.6% faster** |
| 4096 | **OOM** | 11.17s | 11.65s | 26.8GB of 27.0 | declines to plan;
safe |

Three properties, each measured:

- **It resolves the out-of-memory case.** At seq3328/3584/4096 the job
does not run without it.
- **It beats moving everything, where there is room to plan.** 15% at
seq3328 and 9.6% at seq3584 --
the two lengths that need offloading at all -- moving roughly half the
bytes. Copy cost is linear
  at 0.108s per GB (R²=0.96), which is where the time comes from.
- **It does nothing when nothing is needed** — at seq2048/3072. It keeps
every activation resident
  and costs nothing measurable, while blind offloading costs +42%/+43%.

At seq4096 there is only ~2.3GiB of real headroom, so it correctly
declines and matches forced-all
rather than dying.

The place this pass would pay is where recompute is expensive and PCIe
is idle -- long sequences
where attention recompute scales O(s^2) while a copy stays O(s), or
composed with recompute rather
than against it. Neither is measured here.

---------

Signed-off-by: pengdurice <pengduhit@gmail.com>
 2026-08-21 15:17:33 +00:00
Olatunji Ruwase 858e91eea0 Non-reentrant activation checkpoint CPU offload (#8282) 
## 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>
 2026-08-21 15:06:45 +00:00
pengdurice 35c8b03cd6 Make pass contracts cover the passes DeepCompile actually schedules; Renaming files (#8251) 
# [DeepCompile] Make pass contracts cover the passes DeepCompile
actually schedules

Follow-up to #8139. Three gaps kept `PassContract` from firing on real
schedules:

1. `conflicts_with` was ignored whenever the other pass had no
registered contract.
2. Only four pass modules were registered, so `offload_parameters` and
the ZeRO-1/2 reduce passes
   could not be validated at all.
3. The one incompatibility DeepSpeed enforces — one offload target per
run — was checked inside
`if schedule is None:` in `init_z3.py`, so a user-supplied schedule
combining both was accepted.

### The conflict fix

`validate_schedule` matches conflicts against `applied`, but an
uncontracted pass hit `continue`
before `applied.append(name)`, so a conflict naming it was missed in
*both* orderings. Resolving a
missing contract to a shared empty one instead of branching around it
fixes that and removes two
lines of control flow:

```python
contract = _pass_contracts.get(name, _UNCONSTRAINED)
```

The pass stays unconstrained; it is simply visible to conflicts others
declare. `test_conflict_is_symmetric`
missed this because it registers the second pass with an *empty*
contract, which still lands in the registry.

### Contracts

| Registered name | Callable | Contract |
|---|---|---|
| `zero1_compile` | `add_z1_reduce` | conflicts with `zero3_compile` |
| `zero2_compile` | `add_z2_reduce` | conflicts with `zero3_compile`,
`zero1_compile` |
| `offload_parameters` | `offload_parameter_fwd` | requires
`z3_gather_release` |
| `offload_adam_states` | `move_opt_states` | requires
`opt_states_evicted`, conflicts with the three below |
| `offload_adam_states_sync` | `move_opt_states_sync` | conflicts with
the three below, and `offload_adam_states` |
| `offload_adam_states_for_init` | `offload_adam_states_for_init` |
provides `opt_states_evicted`, conflicts with the three below |

The three shared conflicts are `offload_parameters`, `zero1_compile`,
`zero2_compile`.

- `offload_parameter_fwd` rewrites the `dc.allgather_param` nodes
`zero3_compile` inserts; without
  it the pass silently matches nothing.
- The optimizer-state passes dereference a module global set only from
`init_z3.py`, so naming one
  in a ZeRO-1/2 schedule validates today and then dies on `None`.
- `opt_states_evicted` records a dependency previously implicit in
schedule order: `move_opt_states`
plans from profiled peaks, which only describe the run once
`offload_adam_states_for_init` has
taken the optimizer state off the accelerator. `move_opt_states_sync`
reads neither
  `profiling_results` nor `mem_budget` and so takes no such requirement.

Conflicts are declared on one side only, which the fix above is what
makes reliable.

### Naming

Two modules serve stages 1 and 2 but were named for stage 1, matching
neither
`deepspeed/runtime/zero/stage_1_and_2.py` nor the `zero2_compile`
registration:
`zero1_compile.py` → `zero_1_and_2_compile.py`, and `init_z1.py` →
`init_z1_and_2.py`
(`init_z1()` → `init_z1_and_2()`). Pure renames, recorded by git as
such. Constants become
`NAME_Z1` / `CONTRACT_Z1` beside the existing `_Z2` pair; registered
names stay per-stage.

### Testing

`tests/unit/compile/test_pass_contract.py`, 11 tests to 20, all
CPU-only. Covers the conflict fix
in both orderings, each new rejection, schedules written with callables
rather than names, and that
every schedule `init_z1_and_2` and `init_z3` build still validates. One
guard test asserts every
name in a built-in `conflicts_with` is itself registered, since those
are string literals.

`BUILTIN_PASSES` mirrors the engine's registration block by hand —
reaching the real one needs a
constructed engine — so a pass added to the engine and not the test
would go untested.

### Compatibility

Registering the remaining passes widens what a schedule may name and
narrows nothing. The new
conflicts reject only combinations that already failed at runtime, later
and less legibly.

One change genuinely narrows: `move_opt_states` without
`offload_adam_states_for_init` is now a
`PassContractError` where it previously ran. Such a run did not fail
outright — it planned its
offloading from peaks inflated by the resident optimizer state. Turning
that into a schedule-time
error is the intent, but it is the one thing here that rejects something
that used to execute.

---------

Signed-off-by: pengdurice <pengduhit@gmail.com>
 2026-08-21 05:33:42 +00:00
Hongwei Chen 24a41afaff Add fused triton kernel for swiglu (#8244) 
**Performance**:

SwiGLU Benchmark

**Configuration:** `hidden=1024`, `dtype=bf16`

| Tokens | Variant | Forward (µs) | Forward + Backward (µs) | Backward
(µs) |

|-------:|:--------|-------------:|------------------------:|--------------:|
| 8,192  | `triton_fused` | 13.05 | 36.29  | 23.24  |
| 8,192  | `eager`        | 22.53 | 61.47  | 38.94  |
| 16,384 | `triton_fused` | 26.25 | 67.65  | 41.40  |
| 16,384 | `eager`        | 44.97 | 120.95 | 75.98  |
| 65,536 | `triton_fused` | 94.71 | 251.60 | 156.89 |
| 65,536 | `eager`        | 163.12 | 444.64 | 281.52 |

---------

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
 2026-08-21 01:28:51 +00:00
Olatunji Ruwase 310eb8cb48 Route FPDT and checkpoint writer pins through accelerator pin_memory (#8257) 
## Summary
- Route FPDT `SequenceChunk` and backward zero-chunk pins through
`get_accelerator().pin_memory()`, pinning only when the chunk is
on-accelerator (avoid pin-then-discard on CPU inputs).
- Route FastFileWriter AIO buffer through accelerator pin with
`make_copy=False`, and fall back to `Tensor.pin_memory()` when the CPU
accelerator torch path no-ops so DeepNVMe can still skip bounce buffers.

## Test plan
- [x] `pre-commit run --files` on touched paths (already run locally)
- [x] FPDT path smoke: `SequenceChunk` GPU pin + CPU reuse under native
and torch; `TestFPDTAttention` combo `[32-4-128-2048-4]` **PASSED**
- [x] Checkpoint FastFileWriter / AIO write with default torch backend:
`test_fast_file_writer_fd_close.py` **3 passed**; writer pin pattern
`is_pinned=True`
- [x] `DS_PIN_MEMORY_BACKEND=native` writer buffer: `is_pinned=True`,
`aio.is_pinned(buf)=True`, unpin OK; `test_pinned_manager.py` **5
passed**

Evidence: H200 autorun `job-20260819T184219Z` + `job-20260819T184528Z`
on `13131752` / `tjruwase/pin-memory-route-fpdt-writer` (follow-up EXIT
0). GitHub CI green.

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
 2026-08-20 15:59:03 +00:00
iLeGend aa3914df82 [AutoTP] Complete uneven sharding and universal checkpoint support (#8185) 
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>
 2026-08-20 07:32:56 +00:00
Aditya Singh 6098f799ee Keep parameter dtype through ZeRO-3 weight quantization (#8215) 
Fixes #7775

The quantizer op is fp16-only in both directions. quantize_kernel in
csrc/quantization/pt_binding.cpp casts input_vals.data_ptr() to __half*
whatever the tensor dtype actually is, and dequantize is bound as
dequantize<__half> so it always allocates an fp16 output. CUDAQuantizer
passed parameters straight through, so with bf16 enabled and
zero_quantized_weights set, ZeRO-3 quantized bf16 bits reinterpreted as
fp16 and then restored param.data as fp16. Training fails on the
resulting dtype mismatch, which is the BERT failure in the issue. The
bit reinterpretation on the way in is the quieter half of the bug:
values are wrong before the dtype mismatch is ever noticed.

CUDAQuantizer.quantize now converts to fp16 on the way into the kernel
so the values are read correctly, and dequantize takes an optional dtype
so each caller can ask for the dtype its parameter actually has. The
five call sites in the gather paths pass the parameter dtype. Omitting
the argument keeps the previous fp16 return, so no other caller changes
behavior. Precision is not a concern here, since the values are being
quantized to int8 regardless.

Verification: added a parametrized test to
tests/unit/runtime/zero/test_zeropp.py that stands in for the compiled
op with a stub asserting the fp16 contract, and checks a bf16 and an
fp16 parameter both round trip in their own dtype. It passes and it
fails against the unmodified code on both halves of the fix. Run on CPU,
since the test does not need the compiled op. The root cause is verified
by reading csrc/quantization/pt_binding.cpp, not by running on a GPU.
yapf and flake8 are clean on the changed files.

---------

Signed-off-by: Aditya Singh <adisin650@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
 2026-08-19 22:26:14 +00:00
Dmitry 32d51a1840 Fix docstring Args entries that name a parameter the function does not take (#8223) 
Forty-one `Args:` entries name a parameter the function does not take.
Docstrings only — no signature, no behaviour, no test touched.

Twenty are renames where the docstring kept the old name. The ones that
stand out:

| Where | Documented | Actual |
|---|---|---|
| `RaggedUnembed.forward` | `raged_metadata` | `ragged_metadata` (typo)
|
| `FlopsProfiler` | `object` | `model` |
| `attn_out_in_features` | `in_features` | `out_features` |
| `RaggedTopKGating.__call__` | `expert_assignment`, `expert_offset` |
`assignments`, `offsets` |
| `MoEScatter.__call__` | `hidden_states` | `activations` |
| `RaggedEmbeddingKernel.__init__` | `fp_dtype` | `embed_dtype` |
| `InferenceEngineV2.query` | `n_tokens` | `max_request_tokens` |
| `InferenceEngineV2.serialize` | `path` | `save_path` |
| `BlockedKVCache.__init__` | `config`, `enable_offload` | `configs`,
`offload` |
| `MOELayer` | `expert` | `experts` |
| `DeepSpeedCPULion.__init__` | `full_precision_optimizer_states` |
`fp32_optimizer_states` |
| `BertSparseSelfAttention.forward` | `attn_mask` | `attention_mask` |
| `get_grad_norm_direct` (stage 3 and stages 1/2) | `parameters` |
`params` |
| `apply_to_tensors_only` | `functional` | `function` |
| `_create_model_parallel` | `model_parallel_size` |
`model_parallel_size_` |
| `prune_config` | `configs` | `config` |

Where the position in the `Args:` block lined up with the position in
the signature, that is what I used to decide the mapping — for example
in `RaggedTopKGating.__call__` the third and fourth documented names sit
against the third and fourth parameters.

Two needed a description rather than a rename, because the old text
described something that is gone:

- `PipelineEngine.load_module_state_dict` documented `state_dict (str,
None): unused`. The argument is `checkpoint`, and it is used on the very
next line.
- `ResourceManager.parse_results` documented `finished_experiments`,
which is `self.finished_experiments`. The argument is `metric`, the key
read out of each experiment's metrics file.

The remaining entries document something that is not an argument at all:
`layer_id` on `DeepSpeedTransformerInference`,
`DeepSpeedDiffusersAttention` and `DeepSpeedMoEInference` (it is a class
attribute); `beta` on `CUDARMSPreNorm.__call__`; `q_ratio` on
`BlockedRotaryEmbeddings.__init__`; `slack` and `blocks` on
`BlockedKVCache.__init__`; `block_size` on `DSStateManager.__init__`;
`key_padding_mask_mode` and `attn_mask_mode` on
`SparseSelfAttention.forward` (both are constructor arguments);
`seq_len` on `DenseSparsityConfig.__init__`, whose description was a
copy of the one above it; `num_global_blocks` on
`BSLongformerSparsityConfig.__init__`; `scale_factor` on
`DynamicLossScaler`; `scale` on the one-bit `Adam.step` and
`ZeroOneAdam.step`; `max_norm` on both `get_grad_norm_direct`; and
`param` and `param_id` on `_process_selected_fp32_groups_grad`.

One was a formatting slip rather than a wrong name:
`quantize_transformer_layer` had a `Note:` line indented inside its
`Args:` block, so Doxygen-style readers and tooling see a parameter
called `Note`. Moved out.

Every entry was opened and read against its signature. `yapf` produces
no diff and `flake8` is clean on all thirty-one files, using the pinned
`yapf==0.40.0` and `flake8==5.0.4` from `.pre-commit-config.yaml`.

Signed-off-by: darkdi <rantovov5@gmail.com>
Co-authored-by: darkdi <rantovov5@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
 2026-08-19 22:07:17 +00:00
Zhipeng Wang 80c8e5ba42 Add CUDA graph support for HybridEngine generation (#8271) 
## Motivation

Token generation dominates a HybridEngine RLHF iteration, and it is
bound by CPU work rather than by the GPU.

Measured on `facebook/opt-1.3b`, ZeRO-2, single H100, 256-token prompt,
128 new tokens:

| batch | ms / decode step | tok/s |
| ----: | ---------------: | ----: |
| 1     | 5.00             | 200   |
| 4     | 5.10             | 785   |
| 8     | 5.17             | 1549  |
| 16    | 5.58             | 2865  |
| 32    | 5.37             | 5956  |

32× the batch costs 1.07× the latency. The GPU is idle waiting on the
host. A single decode step issues roughly 1300 kernel launches, and of
the 5.22 ms step, 4.52 ms is the model forward while summed kernel time
is only ~2.75 ms.

Replaying a captured CUDA graph replaces those launches with one call.

## Design

Adds an opt-in `hybrid_engine.enable_cuda_graph` flag. Two properties of
the inference kernels shape the design.

**One graph per decode position.** The kernels read the current sequence
length from a host-side counter (`InferenceContext::current_tokens()` in
`csrc/transformer/inference/csrc/pt_binding.cpp`) and pass it to kernels
as a launch parameter. Capture freezes launch parameters, so a single
graph would keep reading and writing one position forever. Host code
*does* run during capture, so capturing one graph per position records
the correct sequence offsets into each.

**All-or-nothing per sequence.** That counter is advanced from host code
(`advance_tokens()`) and is not exposed to Python. Replay runs no host
code, so it leaves the counter behind, and an eager decode step after a
replay would use a stale sequence length and corrupt the KV cache. Eager
and replayed steps therefore must never be mixed within a sequence.
`begin_sequence()` makes the decision once, up front, from the pinned
generation length, before any decode step runs.

Capture is followed immediately by a replay, since capture records work
without executing it; the replay is what actually fills the KV cache for
that position.

## Safety

Graphs are refused, with a warning, wherever captured pointers would not
stay valid:

* **ZeRO stage 3** — parameters are gathered into fresh buffers for each
generate call, and the inference containers hold no persistent weights
(`attn_qkvw is None`). A graph would replay whichever buffers existed at
capture time, which is silently wrong rather than merely slow.
* **`release_inference_cache`** — frees the workspace buffers the graphs
write into.
* **`inference_tp_size > 1`** — untested here.
* **Unpinned generation length** — a sequence that outruns its captured
graphs cannot fall back to eager safely, so graphs engage only when
`min_new_tokens == max_new_tokens`.

Capture failures fall back to eager execution and disable graphs, rather
than failing the training job.

Weight updates were verified explicitly: `reset_params()` writes into
the same inference buffers in place, so the captured pointers stay valid
across optimizer steps. After a step that moved the logits by 11.75,
replay matched a fresh eager forward to within 0.09 and differed from
the pre-step result by the full 11.75 — i.e. graphs track updated
weights rather than replaying stale ones.

## Results

`facebook/opt-1.3b`, ZeRO-2, batch 8, 256-token prompt, 128 new tokens,
single H100, averaged over 4 measured iterations of an RLHF-style loop
(`eval` → `generate` → `train` → forward/backward/step):

| phase       | eager    | CUDA graph | speedup |
| ----------- | -------- | ---------- | ------- |
| generate    | 676.7 ms | 360.8 ms   | 1.88×   |
| train step  | 134.8 ms | 141.0 ms   | 0.97×   |
| **iteration** | **812.8 ms** | **503.5 ms** | **1.61×** |

* **Generated tokens are unchanged**: 1024/1024 token agreement with the
eager path, all 8 sequences identical end to end.
* **Peak memory**: 27.83 → 27.96 GiB (+132 MiB for 127 captured
positions).
* **One-time capture cost**: the first generation captures the graphs
and takes ~12 s; every later generation replays.

## Tests

`tests/unit/hybrid_engine/test_he_cuda_graph.py`:

* 13 CPU-only tests covering the generation-length gate, the ZeRO-3 /
`release_inference_cache` / `inference_tp_size` rejections, and the
dispatch state machine (unknown length, over-long length, prompt
forwards, and invalidation when the generation length changes).
* One GPU end-to-end test (`seq_inference` marker, `opt-125m`) asserting
that a graphed generation is token-identical to the same generation run
eagerly.

Docs: a Hybrid Engine section added to `docs/_pages/config-json.md`,
covering the config block and the new flag's requirements and
restrictions.

## Scope

Default is off, so nothing changes unless the flag is set; the eager
path measured identically before and after this change (676.7 ms
generate both ways). Only ZeRO-2 with `inference_tp_size=1` was
benchmarked, which is what the guards allow.

Signed-off-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>
 2026-08-19 14:22:51 +00:00
Zhipeng Wang 313ce47bd8 Fix ZeRO-3 crash in AutoTP universal-checkpoint metadata (#8270) 
## Problem

`DeepSpeedHybridEngine` cannot be initialized with ZeRO-3.
`deepspeed.initialize()` raises before training starts:

```
File "deepspeed/runtime/hybrid_engine.py", line 354, in create_inference_module
  self.create_inference_containers(self.module)
File "deepspeed/runtime/hybrid_engine.py", line 294, in create_inference_containers
  self._other_layers.append(self.inference_policies[child.__class__][0](module=child, ...))
File "deepspeed/module_inject/layers.py", line 739, in __init__
  self._mark_uc_metadata()
File "deepspeed/module_inject/layers.py", line 807, in _mark_uc_metadata
  original_weight_shape = (original_out_dim, self.weight.shape[1])
IndexError: tuple index out of range
```

ZeRO-3 + HybridEngine is the standard DeepSpeed-Chat actor
configuration, so this blocks the engine's primary use case.

## Cause

`_mark_uc_metadata()` was added in #7908 (universal checkpoint for
AutoTP) and runs unconditionally from `LinearLayer.__init__` /
`LinearAllreduce.__init__`. It reads `param.shape[1]` to reconstruct the
pre-TP parameter shape.

Under ZeRO-3 a partitioned parameter's local data is an empty 1-D
tensor, so that index is out of range. Observed on `facebook/opt-1.3b`,
ZeRO-3, single GPU:

```
LinearLayer weight: shape=(0,) numel=0 ds_id=0
                    ds_shape=torch.Size([50272, 2048])
                    ds_status=ZeroParamStatus.NOT_AVAILABLE  tp_world=1
```

HybridEngine constructs a `LinearLayer` for every non-transformer layer
(embeddings, `lm_head`, final norm) regardless of TP size, so it hits
this path even with `inference_tp_size=1`.

## Fix

Read the pre-partition shape from `ds_shape` instead of the local
`.shape`.

ZeRO-3 sets `param.ds_shape = param.shape` in `partition_parameters.py`
*before* flattening the parameter's data, so `ds_shape` is exactly the
value these call sites were already trying to read. When ZeRO-3 is not
in use the attribute is absent and `param.shape` is used, so the
existing AutoTP path is unchanged.

The helper lives on `TensorParallel_Layer` and is used by both the
column-parallel (`LinearLayer`) and row-parallel (`LinearAllreduce`)
implementations, which had the same bug.

`SubParamLinearLayer` / `SubParamLinearAllreduce` take their shapes from
precomputed `_logical_shape` / `_orig_weight_shape` attributes rather
than indexing `param.shape`, so they are not affected and are left
alone.

## Tests

Two regression tests added to
`tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py`,
covering both the column-parallel and row-parallel paths. They construct
a parameter in the state ZeRO-3 leaves it in and assert the recorded
metadata carries the pre-partition shape.

Both fail on master with the exact `IndexError` (at lines 707 and 807)
and pass with this change. The 10 existing tests in that file continue
to pass.

## Verification

End-to-end RLHF-style loop (`eval()` → `generate()` → `train()` →
`forward`/`backward`/`step()`), `facebook/opt-1.3b`, batch 8, 256-token
prompt, 128 new tokens, H100:

- **ZeRO-3** — fails on master at `deepspeed.initialize()`; runs to
completion with this change (2.63 s/iter steady state).
- **ZeRO-2** — unaffected (796 ms/iter before and after).

Signed-off-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>
 2026-08-19 07:08:36 +00:00
Stas Bekman 24402c3be7 NCCL and other backend PG timeout 30min -> 10min (#8253) 
Apparently nccl/pytorch switched to a 10min collective timeout (from the
original 30min) in 2024 (torch==2.2) but ds is at 30min still - should
we sync with the more realistic timeout? 30min is a way too long, no
collective takes more than 10min these days

This change applies to all backends, even though pytorch changed the
default for nccl only - since it's a sensible modern upgrade.
 2026-08-17 21:16:07 +00:00
Zhipeng Wang c5331bc191 Fix int32 overflow in Triton grouped-GEMM expert offset (#8261) 
## Summary

`group_gemm_triton.py` computed the expert base pointer as `selected *
stride_be` in **int32**. Since `stride_be` is `K*N` **elements**, any
expert index past `2**31 / (K*N)` wraps to a negative offset, so the
kernel reads out of bounds and faults with an illegal memory access.

```python
b_base = b_ptr + selected * stride_be        # int32 product -> wraps
```

A 64-expert layer with `K=4096, N=14336` places the last expert at
**3.70e9 elements**, past the int32 limit of **2.15e9**:

```
RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered
```

The fix casts the expert index to int64 before the multiply. The failure
is a hard fault rather than silent corruption, so no previously-produced
numerics are suspect.

## Reachability

This is not a synthetic configuration. `prefer_triton_grouped_mm()`
selects the Triton path on **all sm<90 devices**, and on **any device
when `torch._grouped_mm` is unavailable**
(`accelerator/cuda_accelerator.py:281`). The overflow needs only
`(E_local - 1) * K * N >= 2**31` — reached by fine-grained-MoE
configurations at Mixtral-scale FFN dimensions when enough experts land
on one rank (low EP degree or high expert count).

## Root-cause evidence

Bisection matches the `2**31` threshold exactly — the trigger is the
**product**, not the expert count alone:

| K | N | experts | `(E-1)*K*N` | vs int32 max | result |
|---|---|---|---|---|---|
| 4096 | 14336 | 64 | 3.70e9 | over | **FAIL** |
| 4096 | 14336 | 64 (32 rows/expert) | 3.70e9 | over | **FAIL** |
| 4096 | 14336 | 32 | 1.82e9 | under | OK |
| 1024 | 14336 | 64 | 9.24e8 | under | OK |
| 4096 | 4096 | 64 | 1.06e9 | under | OK |
| 4096 | 14336 | 8 | 4.11e8 | under | OK |

Rows-per-expert is irrelevant (row 2), which rules out an M-dimension
indexing issue and isolates the fault to the expert-stride term.

## Testing

**Regression test** added to the existing
`tests/unit/v1/moe/test_group_gemm_triton.py` (no new file). It
fails-before / passes-after:

```
# before the fix
FAILED test_expert_offset_exceeds_int32 - RuntimeError: CUDA error: an illegal memory access was encountered

# after the fix
1 passed
```

The test asserts `(num_experts - 1) * stride_be > 2**31 - 1` up front,
so it fails loudly rather than silently stopping exercising the overflow
if the sizes are ever tuned.

**Numerics** at the previously-faulting shape (`E=64, K=4096, N=14336`),
checked per-expert against a `torch` reference:

```
worst relative error across all 64 experts = 3.31e-03   (bf16, K=4096)
```

**Suite results** on `torch 2.8.0+cu128`, triton 3.4.0, H100 (sm90):

```
71 passed, 0 failed   (of 75 collected)
```

No failures, and no regressions attributable to this change. 

## Verification caveat

Verified on **H100 (sm90)**. The overflow is arch-independent integer
arithmetic, so the diagnosis and fix carry over — but note that on a
current PyTorch, sm90 does **not** select this kernel in production
(`prefer_triton_grouped_mm()` returns `False` there, confirmed on `torch
2.8.0`). The regression test calls `group_gemm_triton` directly, so it
exercises the overflow regardless. An Ampere confirmation would still be
a welcome extra check.

The regression test allocates ~4.5 GiB of expert weights. That is
inherent — the bug cannot manifest below `2**31` elements — and it skips
automatically when free device memory is short.

## Incidental observation (not addressed here)

While benchmarking this kernel on sm90, the Triton grouped GEMM measured
**slower than a per-expert `mm` loop** (8192 tokens, H=4096, FFN=14336,
torch 2.4.1 where sm90 still selected the Triton path):

| experts | for-loop | Triton grouped | ratio |
|---|---|---|---|
| 2 | 1.46 ms | 2.26 ms | 0.65x |
| 8 | 1.87 ms | 2.21 ms | 0.85x |
| 32 | 2.11 ms | 2.74 ms | 0.77x |

This is consistent with the module's documented sm80/sm86 scope and is
not a defect — on sm90 with a current PyTorch,
`prefer_triton_grouped_mm()` correctly returns `False` and the native
`torch._grouped_mm` is used instead. Flagging it only because on a torch
build without `torch._grouped_mm`, that helper returns `True` regardless
of architecture, which routes sm90 onto the slower path. Worth a
separate look if that combination is considered supported.

Signed-off-by: Zhipeng Wang <zhipengbayern@gmail.com>
 2026-08-17 06:02:47 +00:00
Olatunji Ruwase c952d92a51 Add compile.offload_activation_pin_memory for DeepCompile activation offload (#8258) 
## Summary
- Add `compile.offload_activation_pin_memory` (default `true`) on
`CompileConfig` and wire it into DeepCompile C++ `offloadTensor` via
`at::TensorOptions().pinned_memory(...)`.
- Document honestly: the offload pass is not in the default schedule
(`offload_activation` alone does nothing); pinning here uses ATen/Torch
host pin and does **not** consult `DS_PIN_MEMORY_BACKEND`.

## Test plan
- [ ] `pre-commit run --files` on touched paths (already run locally)
- [ ] Rebuild deepcompile extension and confirm default still pins host
offload buffers when the offload pass is scheduled
- [ ] With a custom schedule that includes activation offload,
`offload_activation_pin_memory: false` yields pageable host buffers

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
 2026-08-16 05:13:40 +00:00
Olatunji Ruwase aa0e91b950 Add native (DeepNVMe) host-memory pinning backend for accelerators (#8211) 
## 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>
 2026-08-15 13:50:59 +00:00
Zupeng Wang 1cae4920ee Add configurable sum gradient reduction (#8232) 
## Summary

- add a `gradient_allreduce_op` configuration with `"mean"` as the
default and `"sum"` as the new option
- support unscaled gradient sums for ZeRO stages 0, 1, and 2 across
reduce-scatter, allreduce, and non-contiguous fallback paths
- reject unsupported ZeRO stage 3, ZenFlow, and DeepCompile combinations
with clear configuration errors
- document the option and preserve existing mean-reduction behavior

Addresses #7107.

## Motivation

Some distributed objectives, including contrastive learning over
globally gathered embeddings, require summing data-parallel gradients
rather than averaging them. Today users need to rescale the loss
manually to cancel DeepSpeed's world-size normalization.

This change makes the reduction semantics explicit while keeping the
current behavior as the default.

## Validation

Tested on two NVIDIA GeForce RTX 3090 GPUs with PyTorch 2.13.0+cu130:

- `pytest --forked -q
tests/unit/v1/zero/test_zero.py::TestGradientAllreduceOp` — 18 passed
- targeted configuration tests in
`tests/unit/runtime/test_ds_config_dict.py` — 8 passed; the new ZeRO-1/2
DeepCompile cases both failed against the prior head because no error
was raised
- `pytest --forked -q
tests/unit/v1/zero/test_zero_coalesce_grad_reduction.py::TestCoalesceCombinations`
— 12 passed
- `pytest --forked -q
tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py` — 1
passed
- changed-file `pre-commit` hooks, including YAPF, flake8, codespell,
license, `check-torchdist`, and `check-torchcuda` — passed

The distributed test matrix covers ZeRO stages 0/1/2, mean and sum
reductions, reduce-scatter, gradient predivide, prescale, and
non-contiguous gradient fallback.

### Real-training equivalence

A deterministic two-rank, five-step `SimpleModel` regression compares
the default MEAN reduction with SUM while normalizing only the SUM
backward loss by `world_size`, so both modes provide identical gradients
to the optimizer.

- AdamW: ZeRO-0/1/2
- Muon: ZeRO-1/2
- targeted gradient-reduction tests: `23 passed`
- complete `tests/unit/v1/zero/test_zero.py`: `98 passed, 1 skipped`

Across all five optimizer/stage configurations and all five training
steps, the observed loss, full-gradient, and full-parameter differences
were zero.

<img width="2496" height="1572" alt="loss_comparison"
src="https://github.com/user-attachments/assets/3532ad11-5297-42db-baca-8a632bf793ad"
/>

## Limitations

`gradient_allreduce_op="sum"` is intentionally not supported with ZeRO
stage 3, ZenFlow, or DeepCompile. These combinations fail during
configuration instead of silently applying mean semantics.

---------

Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
 2026-08-15 05:25:28 +00:00
Masahiro Tanaka 9bd89f9df9 Fix DeepCompile last use for unconsumed waits (#8254) 
With this PR, DeepCompile treats an unconsumed `wait_allgather` alias as
its own last use to prevent `KeyError`, while preserving downstream
propagation for consumed aliases.
Add synthetic FX regression coverage for both orphan-only and
later-real-use cases.

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
 2026-08-14 23:33:47 +00:00
Naveenraj Kamalakannan 79046032e5 (1/2) Implementing Compiler Pass for AutoTP (#8204) 
Working on #8104 

- Added support for AutoTP. Added two primitives: copy_to_tp and
reduce_from_tp - the f and g nodes.
- Identifies column/row-parallel matmuls by the injected layer type in
nn_module_stack, reading back.
- Wrote a test to verify correctness of module injection and compiler
pass.

cc @tohtana

---------

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.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>
 2026-08-13 19:55:08 +00:00
Masahiro Tanaka 22969fa753 Upgrade pip in Python install-smoke jobs (#8247) 
## Problem

The Python install-smoke matrix inherits different pip releases from its
Python 3.10, 3.11, and 3.12 container images. The Python 3.10 image can
start with pip 23.0.1, whose [PEP 658 name-normalization bug rejects
valid dependency
metadata](https://github.com/deepspeedai/DeepSpeed/actions/runs/31651849867/job/94299944361?pr=8204)
before DeepSpeed installation, while the later Python images use
unaffected pip releases.

## Approach

Upgrade pip in the shared matrix path immediately before PyTorch
installation. Keeping the upgrade common to all three Python versions
avoids image-specific pins and makes every cell select the current
compatible pip release during the run.

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
 2026-08-13 01:17:33 +00:00
Vedant Chauhan 1d580d652f [DeepCompile] Fix staticmethod handling on Python 3.9 (#8240) 
## Description

DeepCompile retrieves PyTorch's compiled backward hooks directly from
the class namespace. On Python 3.9, this returns a `staticmethod`
descriptor that cannot be
  called directly, resulting in:

  ```text
  TypeError: 'staticmethod' object is not callable
  ```

This change unwraps the underlying function before DeepSpeed invokes it.

  The fix covers both supported implementations:

  - PyTorch 2.6: `_backward_prologue`
  - PyTorch 2.7+: `_backward_impl`

A regression test emulates Python 3.9's non-callable `staticmethod`
behavior so the failure remains covered when CI runs on newer Python
versions.

  ## Testing

- `pytest -q tests/unit/compile/test_backend.py
tests/unit/compile/test_zero3_grad_dtype.py`
- `pre-commit run --files deepspeed/compile/patch_compiled_func.py
tests/unit/compile/test_backend.py`

  Fixes #7433

Signed-off-by: Vedant Chauhan <staranonymous1011@gmail.com>
 2026-08-12 00:28:02 +00:00
Sung Hyun Cho 04b4b0cb4d [DeepCompile] Fix KeyError on frozen parameters in ZeRO-3 (#8214) 
Under ZeRO-3 with `"compile": {"deepcompile": true}`, `engine.compile()`
fails with a `KeyError` if the model has any parameter with
`requires_grad=False`. LoRA/PEFT and partial-freeze runs cannot use
DeepCompile at all.

`init_z3()` looks up a grad partition for every module parameter, but
the stage-3 optimizer builds that map only for the parameters it owns.
`_get_trainable_parameter_groups()` drops `requires_grad=False` params
(`stage3.py` L651), and the map is filled from the resulting
`fp16_groups` (`stage3.py` L706-720), so the first frozen parameter is
simply not there.

### Repro

```python
# repro.py, run with: torchrun --nproc_per_node=2 repro.py
import torch
import deepspeed


class Net(torch.nn.Module):

    def __init__(self, dim=128):
        super().__init__()
        self.frozen = torch.nn.Linear(dim, dim)
        self.trainable = torch.nn.Linear(dim, dim)
        self.frozen.requires_grad_(False)

    def forward(self, x):
        return self.trainable(self.frozen(x)).sum()


config = {
    "train_batch_size": 2,
    "train_micro_batch_size_per_gpu": 1,
    "optimizer": {"type": "Adam", "params": {"lr": 1e-4}},
    "zero_optimization": {"stage": 3},
    "bf16": {"enabled": True},
    "compile": {"deepcompile": True},
}

model = Net()
trainable = [p for p in model.parameters() if p.requires_grad]
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=trainable, config=config)
engine.compile()

x = torch.randn(1, 128, device=engine.device, dtype=torch.bfloat16)
loss = engine(x)
engine.backward(loss)
engine.step()
print("ok")
```

```
File "deepspeed/compile/init_z3.py", line 131, in init_z3
    grad_buffer = optimizer._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[p.ds_id]
KeyError: 0
```

`ds_id 0` is the frozen layer's weight. We originally hit this through
the HF Trainer path with a LoRA adapter, where the crash happens inside
`accelerator.prepare()` when accelerate calls `engine.compile()`. Also
reproduces on the released 0.19.2 and 0.19.3.

### Fix

Skip the lookup for frozen parameters and leave them with the empty
buffer that the optimizer-less path (`use_opt == False`) already uses
for every parameter.

Nothing reads that buffer for a frozen parameter.
`add_gather_and_reduce()` skips parameters whose grad node is `None`
(`passes/zero3_compile.py` L130), and the registered buffer is only
consumed by `flushReduceBucket` through those reduce ops
(`csrc/compile/z3.cpp` L261, L282, L311). `set_grad_buffer()` a few
lines below already applies the same `requires_grad` guard against the
same map. Frozen parameters are still registered with the native handle,
since they are partitioned and need gather/release ops in the forward
graph.

### Test

`TestDeepCompile::test_frozen_params` runs the existing
`SimpleFrozenModel` for 10 steps on 2 ranks with ZeRO-3 + deepcompile,
comparing loss and parameters against a ZeRO-0 eager baseline. 10 steps
so the run crosses the `WARMUP` boundary and the prefetch and
selective-gather passes also see the frozen parameters. `compare_loss()`
takes an optional `model_cls`, so the existing callers are unchanged.

Without the fix the test fails at `engine.compile()` with the same
`KeyError`.

### Notes

ZeRO-1/2 does not hit this, since `init_z1()` iterates
`optimizer.bit16_groups`, which contains only trainable parameters. I
have not checked whether frozen parameters work end to end there, so
this change is scoped to ZeRO-3.

`grad_partitions.get(p.ds_id, torch.Tensor())` would also cover a
trainable parameter that was never passed to the optimizer, but that
case is already broken in eager ZeRO-3, so I kept the narrower guard to
match `set_grad_buffer()`. Happy to switch if you prefer the wider one.

Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>
 2026-08-12 00:24:48 +00:00
Vineeth Sai Varikuntla 2cfebbda5a Keep the elastic batch size within max_train_batch_size (#8237) 
## Problem

`_get_compatible_gpus_v01` validates every micro batch against
`max_acceptable_batch_size`:

```python
if not all(mb <= max_acceptable_batch_size for mb in micro_batches):
    raise ValueError(...)
```

but the first heuristic scales the **LCM** of the micro batches, and the
LCM is never checked. It goes into `base_list` and reaches
`get_candidate_batch_sizes`, where the `base >=
max_acceptable_batch_size` branch appends it unscaled. So a batch size
the caller already said was too large becomes a candidate, and since the
LCM divides every micro batch it tends to win the most-valid-GPU-counts
vote in `get_best_candidates`.

The docstring says the heuristic produces "the largest batch size less
than the max_acceptable batch size", and `config-json.md` documents
`max_train_batch_size` as "Max acceptable batch size can be used in
training", so the returned value is not supposed to exceed it.

## Repro

```python
import deepspeed
from deepspeed.git_version_info import version as ds_version

ds_config = {"elasticity": {"enabled": True, "max_train_batch_size": 100,
                            "micro_batch_sizes": [8, 10, 12], "min_gpus": 1,
                            "max_gpus": 1500, "min_time": 20, "version": 0.1}}

print(deepspeed.elasticity.compute_elastic_config(ds_config=ds_config,
                                                  target_deepspeed_version=ds_version))
# (120, [...])    <- 120 against a declared max of 100
```

`DeepSpeedConfig.__init__` writes that return value straight into
`self._param_dict[TRAIN_BATCH_SIZE]`, so the job runs 20 percent over
the limit the user set, with the matching effect on the LR schedule and
step count.

It is not an exotic corner. `[8, 12]` with a cap of 16 gives 24, and a
brute-force sweep over micro batch sets of size 2 and 3 drawn from 1 to
32, against every cap up to 400, finds it in 946105 configurations.

The clearest evidence is in this repo:
`tests/unit/elasticity/test_elastic.py::test_proper_mbsz` sets
`max_train_batch_size` to 32 with micro batches `[1, 2, 3, 7]`, whose
LCM is 42, and gets 42 back today.

## Fix

Skip a base larger than the cap. Scaling one can only make it bigger, so
it can never yield a legal candidate, and every micro batch is already
validated against the cap, so the candidate list cannot end up empty.

## What this changes for the existing tests

`test_basic_10k` is unaffected: still 9792, still 23 valid GPU counts.

`test_proper_mbsz` needed one number changed, and I want to be upfront
about it rather than bury it. Its `world_size=7` was only reachable
because the batch size came back as 42, over its own cap of 32; at any
legal batch size for that config, 7 is not a valid GPU count. I changed
it to 4, where the batch per GPU is 6, so 7 is still correctly ruled out
and the assertion that 3 is chosen is unchanged. That keeps the test
doing what it was written to do, which is check the micro batch picked
for a given world size.

If you would rather keep `world_size=7` working, then the LCM overshoot
is load-bearing rather than a bug, and this PR is the wrong change; I
would want to hear that before going further. I could not find a config
for those micro batches that makes 7 valid without exceeding the cap.

`test_batch_size_within_max` is new and pins the actual contract.

## Test

```
                                                  before      after
test_basic_10k                                    PASS        PASS
test_proper_mbsz  (world_size=7, the old value)   PASS        ElasticityIncompatibleWorldSize
test_proper_mbsz  (world_size=4, the new value)   ElasticityIncompatibleWorldSize   PASS
test_batch_size_within_max  (new)                 FAIL: 120 exceeds 100             PASS
```

Run on CPU by driving the test bodies against the real
`compute_elastic_config`, once against `master` and once against this
branch; this path is pure Python and needs no GPU. `yapf --style
.style.yapf` and `flake8 --config .flake8` are clean on both changed
files, and clean on the unmodified tree as a control.

There is one other open PR touching this file, #8162, in
`compute_elastic_config`'s `return_microbatch` tail. It does not overlap
these lines.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
 2026-08-12 00:24:04 +00:00
Vineeth Sai Varikuntla 8cae9d28a5 Implement the documented per-param-group lists in OneCycle (#8201) 
## The bug

`OneCycle` documents four of its arguments as accepting a
per-param-group list:

```
cycle_min_lr (float or list): Initial learning rate which is the
    lower boundary in the cycle for each parameter group.
cycle_max_lr (float or list): Upper learning rate boundaries in the cycle
    for each parameter group.
cycle_min_mom (float or list): Initial momentum which is the
    lower boundary in the cycle for each parameter group.
cycle_max_mom (float or list): Upper momentum boundaries in the cycle
    for each parameter group.
```

`_initialize_lr` and `_initialize_momentum` only ever broadcast a
scalar:

```python
self.min_lrs = [cycle_min_lr] * len(optimizer.param_groups)
...
self.min_moms = [(cycle_min_mom, 0.99)] * len(optimizer.param_groups)
```

so the documented list is written whole into every param group, and the
optimizer is left holding a list where it expects a number:

```
param_group lrs after construction:   [[0.001, 0.002], [0.001, 0.002]]
param_group betas after construction: [([0.8, 0.85], 0.99), ([0.8, 0.85], 0.99)]

scheduler.step()  -> TypeError: unsupported operand type(s) for -: 'list' and 'list'
optimizer.step()  -> TypeError: unsupported operand type(s) for -: 'int' and 'list'
```

The second line matters: the optimizer is corrupt from construction, so
even a plain `optimizer.step()` fails before the scheduler is stepped at
all.

This is reachable from a plain JSON config, not just the Python API.
`engine.py:1550` does `scheduler(optimizer, **scheduler_params)`, so
`"cycle_min_lr": [0.001, 0.002]` in `ds_config` deserializes to a Python
list and lands directly in `OneCycle.__init__`.

A wrong-length list is also accepted silently, where the siblings raise:

```
OneCycle:     accepted 3 values for 2 param groups, no error
LRRangeTest:  ValueError expected 2 lr_range_test_min_lr, got 3
WarmupLR:     ValueError expected 2 value for min_lr, got [0.0, 0.1, 0.2]
```

## Why implement it rather than delete the docstring lines

Deleting the four "or list" claims would be a smaller diff, but the rest
of `OneCycle` is already per-group end to end: `_get_cycle_lr` zips
`min_lrs` with `max_lrs`, `_get_cycle_mom` zips `min_moms` with
`max_moms`, and `update_lr` walks the param groups. Only the two
initializers collapse the input. Both sibling schedulers in this file
implement the same documented contract, and the two most recent
multi-group fixes here (#7969 for `WarmupCosineLR`, #8171 for
`WarmupLR`) went in the same direction. This reads as an unfinished port
rather than a design decision.

## The fix

Reuse `_format_param`, which is how the siblings already honour this
contract. It was defined twice, identically: as a method on `WarmupLR`,
and again on `WarmupCosineLR` where nothing calls it (`_format_param`
appears in only two files repo-wide, and in the test file only inside a
comment). I promoted the single copy to module level next to `update_lr`
and `get_torch_optimizer`, dropped the dead one, and pointed `WarmupLR`
and `OneCycle` at it. Net result is 19 added, 22 removed, and one
implementation of this logic instead of two.

I chose promoting over leaving one-line delegate methods behind because
`_format_param` is private and has no callers outside this file, so a
delegate would be indirection with no consumer; happy to switch to
delegates if you would rather not remove the methods.

Three details worth calling out rather than leaving for review:

**The momentum call has to wrap the scalar, not the tuple.**
`_format_param` accepts tuples, and the default `cycle_min_mom` pairs
with `0.99` into a length-2 tuple, so wrapping the existing
`(cycle_min_mom, 0.99)` expression would raise at construction for 1 and
3 param groups, and for exactly 2 groups would silently write
`group['betas'] = 0.8` as a float and blow up later in `_get_cycle_mom`.
The correct form, which is what this PR uses, formats the scalar first:

```python
self.min_moms = [(mom, 0.99) for mom in _format_param(optimizer, cycle_min_mom, 'cycle_min_mom')]
```

**Both bounds are now validated before the optimizer is touched.**
`_initialize_lr` used to compute `min_lrs`, write `group['lr']`, and
only then look at `cycle_max_lr`, so a bad-length `cycle_max_lr` left
the param groups half updated. Moving the second `_format_param` call
above the mutation loop makes the constructor all-or-nothing:

```
before: lrs after a failed ctor = [[0.001, 0.002], [0.001, 0.002]]
after:  ValueError, lrs after a failed ctor = [0.1, 0.2]   (untouched)
```

**One token in `_format_param`'s error message.** Both copies
interpolate `FileNotFoundError(param_value)` where the wording promises
a count, so `WarmupLR` currently reports `expected 2 value for min_lr,
got [0.0, 0.1, 0.2]`. Since the two copies are collapsing into one
shared helper, I corrected it to `len(param_value)` rather than carry
the typo into the surviving copy. It is the only change to `WarmupLR`'s
behaviour and nothing asserts on that message (no `pytest.raises(...,
match=...)` anywhere in the file); say the word and I will drop it back
to verbatim.

**Not claiming this is strictly safer for momentum.** Because
`_format_param` accepts tuples, a betas-shaped `cycle_min_mom=(0.8,
0.999)` on a two-group optimizer goes from a loud `TypeError` to
silently training with per-group momenta. That hazard already exists
identically in `WarmupLR`, so I kept the behaviour symmetric rather than
diverging, but it is a real trade rather than a pure win.

## Tests

Added to `tests/unit/runtime/test_lr_schedulers.py` as module-level
functions, matching the existing plain tests there:

- `test_one_cycle_accepts_per_group_lr_and_momentum_lists`: two param
groups, per-group lists for all four arguments, asserting the
constructor sets each group's own lr and `betas[0]`, that the cycle peak
reaches each group's own `cycle_max_lr` with momentum at its own
`cycle_min_mom`, and that the bottom of the cycle returns each group to
its own `cycle_max_mom`.
- `test_one_cycle_rejects_wrong_length_per_group_lists`, parametrized
over all four arguments.

It uses `Adam` rather than `SGD` on purpose: `_initialize_momentum`
returns early when `'betas' not in optimizer.defaults`, so the momentum
half of the test would silently never run under SGD.

`pytest` cannot start on my machine (no GPU, and the `tests/unit`
conftest pulls in the distributed harness), so I ran the module-level
tests in this file directly against the real `lr_schedules.py`, with the
`DistributedTest` classes stripped and only `deepspeed.utils.logger`
stubbed. Three runs:

```
control     upstream lr_schedules.py + upstream tests    21 passed, 0 failed
before      upstream lr_schedules.py + these tests       21 passed, 5 failed
after       this branch                                  26 passed, 0 failed
```

All 5 failures before are the new tests, and the 21 pre-existing ones
are unchanged by this diff. The `DistributedTest` OneCycle coverage
(`TestOneCycle.test_lr`, `test_mom`) and the other scalar-momentum users
(`test_fp16.py`, `test_bf16.py`, `test_pipeline.py`,
`test_other_optimizer.py`) all pass scalars, which take the unchanged
broadcast path; I am relying on CI for those since they need a GPU.

Lint: `yapf` 0.40.0 with the repo's `.style.yapf` reports no diff on
both files, and `flake8` with the repo's `.flake8` is clean on both
(also confirmed clean on the unmodified files, so that is a real result
rather than a config that checks nothing).

## Prior art

No open or closed PR implements list support here. `--search` over
`lr_schedules`, `_format_param`, `OneCycle`, `cycle_min_lr` and `lr
scheduler list param groups` turns up #8151, #8166, #8171, #7969, #8179,
#1455 and #4563, all merged and none touching these two initializers. No
open issue covers it either; the only open `OneCycle` issue is #3492, a
request for `CosineAnnealingLR` support.

This follows #8179 in the same class, so to be upfront about it: that
one was about the cycle shape (`_initialize_cycle` and
`_get_scale_factor`), this one is about the two value initializers, and
I did not see it while in there. If you would rather batch further
`lr_schedules.py` work, tell me and I will hold the rest.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Co-authored-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
 2026-08-11 22:47:19 +00:00
Logan Adams 66f9a36401 Update version.txt after 0.19.5 release (#8243) 
**Auto-generated PR to update version.txt after a DeepSpeed release**
Released version - 0.19.5
Author           - @loadams

Co-authored-by: loadams <loadams@users.noreply.github.com>
 2026-08-11 22:41:00 +00:00
Vineeth Sai Varikuntla 5dceeeab53 Return a copy from OnebitLamb.get_lamb_coeffs (#8227) 
## Problem

`OnebitLamb.step()` opens by dropping the previous step's stats in
place:

```python
# remove the previous stats
del self.lamb_coeffs[:]
```

and `get_lamb_coeffs()` handed back that same list object:

```python
def get_lamb_coeffs(self):
    return self.lamb_coeffs
```

So a caller who reads the coefficients (to log them, or to watch the
trust ratios during warmup) is holding the optimizer's own list, and the
next `step()` empties it under them. The snapshot silently becomes `[]`
instead of the values that were read.

Running the two accessors as they exist today, with the one `step()`
statement that touches the list:

```
OnebitLamb  (deepspeed/runtime/fp16/onebit/lamb.py)
   source            : return self.lamb_coeffs
   before step()     : [tensor(0.5000), tensor(1.5000)]
   after  step()     : []
   caller's snapshot emptied by step(): True

FusedLamb   (deepspeed/ops/lamb/fused_lamb.py)
   source            : return lamb_coeffs
   before step()     : [0.5, 1.5]
   after  step()     : [0.5, 1.5]
   caller's snapshot emptied by step(): False
```

`FusedLamb` has the identical `del self.lamb_coeffs[:]` in its `step()`
and the identical accessor name, and is not affected only because its
version builds a new list on the way out. So the two optimizers disagree
today about whether the value they hand you survives the next step.

## Fix

Return a copy, so both optimizers behave the same way:

```python
return list(self.lamb_coeffs)
```

## One thing I deliberately did not change

`FusedLamb.get_lamb_coeffs` returns Python floats (`[c.item() for c in
self.lamb_coeffs]`) while this one returns the tensors. Making them
match would mean changing the element type this method has always
returned, which is a separate call from fixing the aliasing, so I left
it alone rather than folding a behaviour change into a bug fix. Happy to
align it here or in a follow-up if you would rather the two accessors
were identical.

## Test

`test_onebit_lamb_get_lamb_coeffs_returns_a_copy` in
`tests/unit/runtime/half_precision/onebit/test_onebit.py`. It needs no
accelerator and no distributed backend: the accessor is pure Python, and
since `OnebitLamb.__init__` asserts on an initialized backend, the test
builds the instance with `__new__` and gives it only the attribute the
accessor reads. It sits with the other 1-bit Lamb tests rather than in a
new file, and it runs in `cpu-torch-latest` since that job runs all of
`unit/` and this module has no accelerator-level skip.

Fails before the change (`step() emptied the list returned to the
caller`) and passes after.

## Checks

`pre-commit run --files deepspeed/runtime/fp16/onebit/lamb.py
tests/unit/runtime/half_precision/onebit/test_onebit.py` is clean,
including yapf, flake8, check-torchdist, check-license and codespell.
Commit is signed off.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
 2026-08-11 00:33:49 +00:00
Thong Nguyen b39e07a717 Fix checkpoint rank selection for Ulysses sequence parallelism (#8226) 
## Summary

- use checkpoint model-parallel rank 0 for Ulysses sequence parallelism
- centralize checkpoint rank selection across model, expert, optimizer,
and ZeRO checkpoint paths
- preserve the real model-parallel rank for tensor-parallel MPUs

## Problem

Ulysses sequence parallelism does not shard model weights, so DeepSpeed
records `mp_world_size == 1`. However, the Ulysses MPU aliases
`get_model_parallel_rank()` to the sequence-parallel rank. Checkpoint
loading combined those values when selecting a model-state file, causing
SP ranks greater than zero to index past a single non-SP checkpoint
shard.

The same rank interpretation was also used for checkpoint filenames, so
this change keeps save and load behavior consistent for replicated SP
weights and optimizer shards.

## Tests

- `pytest tests/unit/sequence_parallelism/test_ulysses.py -k
'CheckpointRank or load_non_sequence_parallel_checkpoint' -q` (3 passed)
- `pytest tests/unit/checkpoint/test_latest_checkpoint.py -q` (2 passed)
- `pre-commit run --files deepspeed/runtime/engine.py
tests/unit/sequence_parallelism/test_ulysses.py`

The distributed regression saves without SP, loads module-only with
SP=2, then saves and fully resumes the SP checkpoint.

Signed-off-by: Thong Nguyen <thong.nguyen@snowflake.com>
 2026-08-10 23:23:14 +00:00
Olatunji Ruwase cf44300453 Unmanaged gradient accumulation: ZeRO offload support (#8225) 
## Summary
* Extends unmanaged gradient accumulation
(`managed_gradient_accumulation=false`) to **ZeRO optimizer-state and
parameter offload** (CPU/NVMe). Follow-up to #8217 (ZeRO stage 3, now
merged).
* Stage 2/3: grads still reduce/partition every `backward()`; `step()`
finalizes deferred offload boundary work (grad norms + FP32/NVMe copy)
via `finalize_gradient_accumulation_boundary()`.
* Stage 1: continues to reduce at `step()` via `allreduce_gradients()`,
which already performs offload boundary finalization when the boundary
flag is true.
* Pipeline parallelism, DeepCompile, Apex AMP, and stage-0/1
`overlap_comm` remain unsupported.

## Test plan
Validated on a 2-GPU node:
* [x] Full `-k Unmanaged` suite (**31 passed**), including:
  * `test_unmanaged_matches_managed_optimizer_offload[1|2|3]`
  * `test_unmanaged_matches_managed_param_offload` (stage 3)
* existing non-offload unmanaged equivalence / varying-GAS / rejection
tests
* [x] Docs updated (`config-json.md`, `training.rst`); 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: Ma, Guokai <guokai.ma@gmail.com>
v0.19.5 		2026-08-09 18:28:04 +00:00
Olatunji Ruwase 5ad9a97804 Split native host pinning into standalone pin_memory op (#8236) 
## 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>
 2026-08-09 13:15:28 +00:00
Hongwei Chen a2f8d9ea09 Fix compilation error for torch 2.12+ (#8238) 
Use `-std=c++20` for torch 2.12+. 
Fix #8235

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
 2026-08-09 08:00:53 +00:00
Masahiro Tanaka da06640728 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>
 2026-08-07 22:56:36 +00:00
Ma, Guokai 7dfa67aeaf fix(zero3): async grad offload + pinned offload buffers by default (#8207) 
## 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>
 2026-08-07 15:46:49 +00:00
Deng Zengyong 2d4488ef9d Fix ZeRO++ secondary shard copy for small params (#8210) 
## Summary

This PR fixes a ZeRO++ edge case in _partition_param_sec() when a small
parameter does not overlap with its computed secondary shard.

In the reported case from deepspeedai/DeepSpeed#6659, a parameter of
shape [32] is DP-aligned to 2048, and with zero_hpz_partition_size=16
the secondary shard size becomes 128. For some secondary-group ranks,
secondary_start is already beyond param.ds_numel, so sec_numel becomes
0. The old code still executes:

  one_dim_param.narrow(0, secondary_start, sec_numel)

PyTorch raises IndexError for this case even when sec_numel == 0 if the
start index is out of range.

  This change fixes the issue by:

  - skipping the secondary copy when sec_numel == 0
- zero-filling the secondary shard buffer first so uncovered padding
remains deterministic and does not leak uninitialized values into later
coalesced quantization

  This PR also adds focused regression tests covering:

  - the small-parameter, no-overlap secondary shard case
  - zeroed padding for partially covered secondary shards

  Fixes deepspeedai/DeepSpeed#6659

  ## Testing

python3 -m pytest -q tests/unit/runtime/zero/test_zeropp.py -k
'small_param_secondary_shard_without_overlap or
secondary_shard_padding_is_zeroed'

  Observed locally:

- current master reproduces the IndexError on the
deepspeedai/DeepSpeed#6659 geometry
  - this PR avoids the out-of-range narrow() call
  - 2 passed

---------

Signed-off-by: zengyong <2595650269@qq.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
 2026-08-07 10:37:47 +00:00
Aditya Singh 646e491c95 Do not let one op builder's compatibility probe break importing deepspeed (#8216) 
Fixes #7452

Importing deepspeed probes every op for build compatibility in
git_version_info.py, whether or not the caller will ever build that op.
Nothing catches a probe that fails, so on a machine with a visible GPU
but no CUDA toolkit the CUDA op builders raise MissingCUDAException out
of is_compatible() and import deepspeed fails outright, even for a
sharding only workload that needs no custom ops. ds_report fails the
same way. That is the nvcc assumption reported in the issue.

This adds probe_is_compatible, which reports an op whose probe fails as
not compatible and prints why, and uses it in the two places that scan
every op: the import time scan and the ds_report table.

It also makes installed_cuda_version raise MissingCUDAException when
nvcc cannot be run. CUDA_HOME regularly points at a runtime only install
with no nvcc under bin, and the raw FileNotFoundError from that case
slipped past the two except MissingCUDAException handlers already in
builder.py that fall back to a CPU only build.

Verification: added three tests to tests/unit/ops/test_op_builder.py
covering a runtime only CUDA_HOME, a probe that raises, and a probe that
answers. All 16 tests in that file pass, with the pre-existing CUDA fork
test skipped for lack of a GPU. The new tests fail against the
unmodified code, and I confirmed separately that installed_cuda_version
raised FileNotFoundError rather than MissingCUDAException before the
change. yapf and flake8 are clean on the changed files.

Signed-off-by: Aditya Singh <adisin650@gmail.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
 2026-08-07 10:35:48 +00:00
Logan Adams c6e727946d Update actions versions for workflows. (#8228) 
Signed-off-by: Logan Adams <loadams@microsoft.com>
 2026-08-07 01:24:28 +00:00
Masahiro Tanaka 20328cf551 Restore the custom DCO workflow (#8220) 
Restore the custom DCO workflow.

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
 2026-08-06 04:57:17 +00:00
Logan Adams 14364cd19e Update version.txt after 0.19.4 release (#8219) 
Signed-off-by: Logan Adams <loadams@microsoft.com>
 2026-08-06 04:43:10 +00:00
Logan Adams c455031422 Update version.txt pre-0.19.4 release (#8218) 
Signed-off-by: Logan Adams <loadams@microsoft.com>
v0.19.4 		2026-08-06 00:48:13 +00:00
Olatunji Ruwase e75d1150e5 Unmanaged gradient accumulation: ZeRO stage 3 support (#8217) 
## 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>
 2026-08-05 22:04:02 +00:00
Ehsan 4e3acb1e33 Fix ZeRO-3 all_reduce param fetch stride for padded parameters (#8158) 
## Problem

In `Init._all_gather_coalesced`, the branch taken when
`stage3_use_all_reduce_for_fetch_params` is enabled sizes the flat
buffer with the
*aligned* element count but walks it with the *unaligned* one:

```python
flat_buffer_size = sum(p.ds_numel_aligned for p in params)   # aligned
...
    start_param += param.ds_numel                            # unaligned
```

`ds_numel_aligned` is `ds_numel` rounded up to a multiple of the
partition world size, so
the two are equal only when `ds_numel % world_size == 0`. When a
parameter is padded, every
following parameter's base offset lands `padding_i` slots inside its
predecessor's aligned
span.

Those overlapping slots are the predecessor's partition padding. They
are not harmless:
`_partition_param` allocates the partition with `torch.empty` and, on
the one rank whose
partition extends past the end of the parameter, copies only
`elems_to_copy` elements, so
the tail is never written and holds whatever the allocator returned.
Because the fetch is
an `all_reduce` with SUM, that uninitialized tail is added into the next
parameter's
leading elements. The sibling `all_gather` path writes the same padding
but each parameter
narrows only its own `ds_numel` out of the result, so it is never read;
the corruption is
specific to the SUM reconstruction.

Total stride is smaller than the allocated buffer, so there is no
out-of-bounds access and
nothing crashes. The parameter simply comes back with wrong leading
values.

## Invariant

Each parameter owns a disjoint contiguous span of `ds_numel_aligned`
slots in `flat_tensor`,
which is exactly what `flat_buffer_size` already reserves for it.

`ds_numel_aligned` has a single writer in the package,
`param.ds_numel_aligned = tensor_size`
in `_partition_param`, set in the same block that creates `ds_tensor`.
It has two readers:
the buffer sizing above and the line changed here. Since the buffer
sizing already reads the
attribute for every param in this same loop, the change introduces no
new requirement on
when the attribute must exist.

## Fix

Advance `start_param` by `ds_numel_aligned`, matching the buffer sizing.

## Verification

Added `tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py`,
parametrized over a
padded case (numels 5 and 7 under `world_size=2`) and an aligned case
(numels 4 and 6):

| case | master | this branch |
|---|---|---|
| padded | fails | passes |
| aligned | passes | passes |

On master the padded case reports:

```
param p1 was not reconstructed exactly:
  expected [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
  got      [7778.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
```

7778 is 7777 (the sentinel filling `p0`'s uninitialized padding) plus
1.0 (`p1`'s real
first element), which is the SUM overlap rather than a generically wrong
value.

About that sentinel: the padding tail is uninitialized by construction,
and a fresh
allocation usually reads back as zero, in which case the SUM adds zero
and the bug is
invisible. A first attempt at this reproduction showed no corruption for
exactly that
reason. The test therefore fills new allocations with a sentinel so the
tail's contents are
deterministic rather than dependent on whether the allocator hands back
a recycled block.
It does not inject anything into the buffer under test.

Run in a CPU-only container, world_size 2 over gloo, `LOCAL_SIZE=2`:

```
python -m pytest unit/runtime/zero/test_zero_allreduce_fetch_params.py
```

`pre-commit run --files` passes on both changed files.

## What I did not verify

I have no multi-GPU machine here, so this was exercised only on CPU with
gloo, not with
NCCL on real devices. I also did not measure the effect on a full
training run, so I cannot
say how the corrupted leading elements affect convergence in practice;
the claim here is
limited to the reconstructed parameter values being wrong.

---------

Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Co-authored-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>
 2026-08-05 20:27:51 +00:00
Olatunji Ruwase 3c737a1c2f Unmanaged gradient accumulation: ZeRO stage 2 support (#8203) 
## 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>
 2026-08-05 16:12:40 +00:00
Vineeth Sai Varikuntla b8b4480965 Stop trim_mean sorting the caller's list in place (#8199) 
## The bug

`CommsLogger.comms_dict` stores parallel lists per message size,
`[count, latencies, algbws, busbws]`, where index `i` is the i-th
recorded op. `get_operation_summary()` and `log_all()` hand each of
those lists to `trim_mean`, which sorted **in place**:

```python
data.sort()
k = int(round(n * (trim_percent)))
return mean(data[k:n - k])
```

So summarising sorts the three lists independently and destroys the
correspondence between them. Driving the real `CommsLogger` with
latencies `[3.0, 1.0, 2.0]`:

```
BEFORE  latency: [3.0, 1.0, 2.0]   algbw: [0.0055, 0.0164, 0.0082]
AFTER   latency: [1.0, 2.0, 3.0]   algbw: [0.0055, 0.0082, 0.0164]

stored data mutated by a read-only summary call: True
latency[i] * algbw[i], constant by construction: [0.0055, 0.0164, 0.0492]
```

`algbw` is computed from `latency` in `calc_bw_log`, so `latency[i] *
algbw[i]` is constant for a fixed message size. After summarising it is
not, because row 0 now pairs the fastest op (1.0 ms) with the lowest
bandwidth. `get_raw_data()` hands that out to anyone consuming the log.

Two things make this look unintended rather than a quirk:

- `get_operation_summary()` already does `op_data =
self.comms_dict[operation_name].copy()` with the comment "Create a
snapshot to avoid concurrent modification issues". The intent not to
disturb the stored records is explicit; the copy is just shallow, so the
inner lists are shared and the sort reaches them anyway.
- A summary is a read. Calling it should not rewrite what was recorded,
and nothing documents that it does.

## It also breaks the straggler breakdown, which is the bigger effect

Raised in review and worth stating here, because it is worse than the
reordering above.

`log_all()` reads `vals[1]` twice. The summary loop calls
`trim_mean(vals[1], 0.1)`, which sorts it, and the `show_straggler`
block afterwards builds both `lats` and `min_lats` from that same,
now-sorted list:

```python
lats = torch.tensor(vals[1], device=device)
min_lats = torch.tensor(vals[1], device=device)
dist.all_reduce(min_lats, op=ReduceOp.MIN)
total_straggler = (lats - min_lats).sum().item()
```

`all_reduce(..., MIN)` is elementwise, so it assumes index `i` is the
same collective on every rank. Each rank has sorted its own list
independently by then, so index `i` is "the i-th fastest op on *that*
rank" and the reduction takes the minimum across unrelated operations.
`comms_dict_snapshot = self.comms_dict.copy()` at the top of `log_all()`
does not help, for the same reason the copy in `get_operation_summary()`
does not: it is shallow.

That is not a cosmetic reordering, it changes the number. Two ranks,
four index-aligned ops:

```
rank0 = [2.30, 1.60, 3.60, 1.29]
rank1 = [3.14, 2.46, 1.23, 3.03]
  true total_straggler per rank   [2.37, 3.44]
  with the in-place sort          [0.52, 1.59]   <- what gets reported
```

Over 20000 random four-op cases, 16246 report a different total, and the
error has a direction: sorting both ranks aligns them as closely as
their values allow, so the differences shrink and the straggler effect
is systematically **under**-reported. Small cases can coincide (a
three-op example I tried happened to agree), which is part of why this
is easy to miss.

## Fix

`data = sorted(data)`. That fixes every caller at the shared function
rather than patching `log_all` and `get_operation_summary` separately,
and the trimmed mean it returns is unchanged.

## Tests

Added to `tests/unit/comm/test_comms_logger.py`:

- `test_get_operation_summary_does_not_reorder_the_stored_records`
populates `comms_dict` directly and asserts the three stored lists
survive a summary call unchanged, that `latency[i] * algbw[i]` is still
constant, and that `avg_latency_ms` is still the correct trimmed mean.
- `test_trim_mean_does_not_mutate_its_argument` pins the contract at the
function itself.

Both are dist-free, like the existing test in that file. `comms_dict` is
populated directly rather than through `append()` because `append()`
calls `calc_bw_log`, which needs a live process group.

Fail-before / pass-after against the unmodified `timer.py`:

```
upstream timer.py, new tests present
  PASS   test_stop_profiling_comms_disables_prof_all
  FAIL   test_get_operation_summary_does_not_reorder_the_stored_records
  FAIL   test_trim_mean_does_not_mutate_its_argument

with this fix
  PASS   all three
```

The pre-existing test passing either way is deliberate: this is a
distinct failure from the one it covers.

How these were run, since I would rather say than imply a normal pytest
run: I do not have a GPU or a built DeepSpeed here, so I executed
`deepspeed/utils/timer.py` and `deepspeed/utils/comms_logging.py` from
source against stub `deepspeed.comm` / `deepspeed.accelerator` modules
and ast-extracted the tests. That exercises the real `CommsLogger` and
the real `trim_mean`; CI is the runner for the suite proper.

No existing issue or PR covers this; searching `trim_mean` across open
and closed returns only the merged PR that introduced the current form.

---------

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
 2026-08-05 15:11:27 +00:00
Olatunji Ruwase e1d6b4fe43 Share DeepNVMe pinned-tensor manager and route swap buffers through I/O handles (#8212) 
## Summary
DeepNVMe skips the bounce-buffer copy only when a buffer is torch-pinned
or managed by the handle's pinned-tensor manager. That manager was
**per-handle** and recognized only **exact base pointers**, so buffers
allocated by one handle — or narrows/views of a shared pool submitted
through a different read/write handle — were not recognized and always
bounced.

This PR makes the pinned-tensor manager a **process-wide shared instance
with range-based recognition**, and switches the swap subsystem to
obtain pinned memory and query pinned status through its I/O handles.

## Changes
- **C++**
- `deepspeed_pin_tensor_t` is now a process-wide `shared()` singleton
guarded by a `std::mutex`; `is_managed` is **range-based** so
slices/views of a locked buffer are recognized.
- `deepspeed_io_handle_t` and `cpu_op_desc_t` hold the manager via
`std::shared_ptr`.
- Added `handle.is_pinned(buffer)` with bindings in `py_ds_aio.cpp` and
`py_ds_gds.cpp` (GDS inherits the base).
- **Python (swap subsystem only)**
- `SwapBufferManager`/`SwapBufferPool` and the optimizer swappers take
an `aio_handle`; buffers are allocated via `new_cpu_locked_tensor` and
pinned status is queried via `handle.is_pinned`.
- Optimizer-swapper subclasses create their handle before
`super().__init__` so it can be threaded through.
- **Test**: `tests/unit/v1/nvme/test_pinned_manager.py` covers
narrow/view recognition and cross-handle sharing.

## Test plan
- [x] `pre-commit` (yapf/flake8/clang-format/check-license) on all
touched files.
- [x] `tests/unit/v1/nvme/test_pinned_manager.py` — 3/3 pass (narrow
recognition, cross-handle sharing, unmanaged buffer).
- [x] `tests/unit/v1/nvme/` + `tests/unit/utils/test_pin_memory.py` +
`tests/unit/v1/accelerator/test_accelerator.py` — 146 pass.
- [x] Swap smoke test
(`tests/unit/runtime/zero/test_nvme_checkpointing.py`): reproduces the
pre-existing baseline exactly (no regression; the failing
optimizer-on-NVMe configs fail identically on `master`).

Made with [Cursor](https://cursor.com)

---------

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
 2026-08-05 14:38:49 +00:00
Vineeth Sai Varikuntla 87e4c8c8c6 Make PipelineModule.set_checkpoint_interval actually change the interval (#8178) 
## Problem

`PipelineModule.set_checkpoint_interval()` does not change the
activation checkpoint interval.

```python
def set_checkpoint_interval(self, interval):
    assert interval >= 0
    self.checkpoint_interval = interval
```

Every reader uses `self.activation_checkpoint_interval`: `forward()`
branches on it and steps the layer loop by it (`module.py` L370-379),
`_precompute_checkpointable_values()` keys its cache on it (L223-230),
and `PipelineEngine` assigns it directly (`pipe/engine.py` L212).
`self.checkpoint_interval` is read nowhere in the repo, so the call
assigns a dead attribute and the schedule silently stays as it was.

Assigning the right attribute is not sufficient on its own, because the
recompute path it feeds is also broken:

```python
def _precompute_checkpointable_values(self):
    if self.activation_checkpoint_interval > 0 and self.is_checkpointable_results_interval != self.activation_checkpoint_interval:
        num_layers = len(self.forward_funcs)
        self.interval_was_zero = False
        for start_idx in range(0, num_layers, self.activation_checkpoint_interval):
            ...
            self.is_checkpointable_results.append(self._is_checkpointable(funcs))
        self.is_checkpointable_results_interval = self.activation_checkpoint_interval
```

The `is_checkpointable_results_interval !=
activation_checkpoint_interval` guard exists precisely so the values are
recomputed when the interval changes, but the loop appends to
`self.is_checkpointable_results` without clearing it, so results
computed for the previous interval stay at the front of the list.

`forward()` then pairs the layer blocks with that list positionally:

```python
for start_idx, is_checkpointable_result in \
    zip(range(0, num_layers, self.activation_checkpoint_interval), self.is_checkpointable_results):
```

so each block is checkpointed according to a decision made for a
different partitioning of the layers.

## Reproduction

8 layers, the first 4 without parameters and the last 4 with, so
`_is_checkpointable` genuinely differs per block. Going from interval 4
to interval 1:

```
interval=4  results: [False, True]

set_checkpoint_interval(1) -> activation_checkpoint_interval = 4     # unchanged
  results: [False, True]                                            # never recomputed

# assigning activation_checkpoint_interval directly, the way PipelineEngine does:
  results: [False, True, False, False, False, False, True, True, True, True]   # 10 entries for 8 blocks
  expected:               [False, False, False, False, True, True, True, True]

forward blocks : [(0, False), (1, True), (2, False), (3, False), (4, False), (5, False), (6, True), (7, True)]
expected       : [(0, False), (1, False), (2, False), (3, False), (4, True), (5, True), (6, True), (7, True)]
```

Blocks 1, 4 and 5 are checkpointed against the wrong decision: block 1
holds no parameters and is checkpointed anyway, blocks 4 and 5 hold
parameters and are not.

## Fix

Clear the cached results before recomputing them, and have the setter
assign `activation_checkpoint_interval` and rebuild the cache:

```python
self.is_checkpointable_results = []
```

```python
def set_checkpoint_interval(self, interval):
    assert interval >= 0
    self.activation_checkpoint_interval = interval
    self._precompute_checkpointable_values()
```

The setter has to do both. Assigning the interval alone would leave
`forward()` zipping the new, longer block range against a list still
sized for the old interval, and `zip` stops at the shorter one, so
trailing layer blocks would be dropped from the forward pass entirely.

Nothing changes for the normal path: `PipelineEngine` assigns the
interval once and calls `_precompute_checkpointable_values()` while the
cache is still empty, so clearing an empty list is a no-op and the guard
still skips the recompute when the interval is unchanged.

## Testing

`TestPipeModuleCheckpointInterval` in
`tests/unit/pipe/test_pipe_module.py` builds a `PipelineModule` at
interval 4, calls `set_checkpoint_interval(1)`, and asserts the interval
is updated and the results match a module constructed at interval 1
directly. It fails on master on both assertions (the interval stays 4,
and the results stay `[False, True]`) and passes with the fix. The mixed
`ReLU`/`Linear` model is deliberate: with a uniformly parameterised
model only the length of the list is wrong, and the misalignment would
not show.

`yapf` and `flake8` are clean on both changed files.

---------

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
 2026-08-04 21:40:08 +00:00
Hongwei Chen 58aab4200c Reduce redundant work in AutoEP token routing (#8209) 
**Motivations**
``count_tokens_per_expert`` function was called three times in each
forward pass, and the ``torch.bincount`` inside it will introduce
cpu-gpu sync. But the results of the first call could be reused.

**Changes**
• Reuse the router's histogram. The router already
computes num_tokens_per_expert; reuse it through compute_split_plan  and
the  ep_size == 1 path instead of recomputing it in
AutoEPMoELayer.forward .
• Faster count_tokens_per_expert. Replace  torch.bincount  with a
pre-sized  zeros(num_experts, int32)  +  scatter_add_ , avoiding the
device-to-host sync that  bincount needs . The helper now always returns
an int32 histogram; the unused  out_dtype / deterministic_safe  params
and padding logic are removed.
• Remove deterministic_safe path in ``count_tokens_per_expert``. The
histogram of integers is inherently deterministic. The op just sums 1
per bucket. Integer addition is associative and commutative, so the
atomic accumulation order has zero effect on the result — every run
produces identical counts
• Remove the  TokenReorderer  module. Its logic (argsort by expert +
score gather) is a two-liner, now inlined directly in the layer
forward. 

**Performance**
The time below is measured from the moe gate kernel to the last kernel
before first all-to-all communication.
A100: 2.3ms -> 1.7ms. 
H200: 0.89ms -> 0.53ms.

---------

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
 2026-08-04 07:35:55 +00:00