master
1768 Commits
| 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> |
||
|
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> |
||
|
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>
|
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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. |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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>
|
||
|
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>
|
||
|
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>
|
||
|
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>
|
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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>
|
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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>
|
||
|
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> |
||
|
Masahiro Tanaka
|
d8dd9e43c7 |
Fix DeepCompile ZeRO-3 gathered parameter ownership (#8157)
## Problem DeepCompile inserts ZeRO-3 parameter all-gather and release operations into compiled graphs. When Dynamo skips a frame because of a graph break, however, that frame executes eagerly and does not run those graph operations. The eager fallback introduced in #8059 handles this case by all-gathering a partitioned parameter when the skipped frame accesses it through `ZeROOrderedDict`. The fallback is enabled around `DeepSpeedEngine.forward()`. Dynamo guard evaluation occurs inside that outer forward context and also resolves parameters through `ZeROOrderedDict`, while `torch.compiler.is_compiling()` is false. The fallback could therefore mistake a guard lookup for actual eager execution and unnecessarily all-gather the parameter. Parameters gathered by the fallback are normally partitioned after backward, but that cleanup does not run when backward is skipped. A fallback-gathered parameter may also be passed to an explicit `GatheredParameters` context, which must keep the full tensor available until the context exits. ## Why it matters These cases require different behavior: - Dynamo guard evaluation should not trigger an all-gather. - A parameter gathered for an eagerly executed frame must remain available through backward and then be partitioned. - If backward does not run, a leftover full parameter must be partitioned before the next outermost forward. - A parameter covered by `GatheredParameters` must remain fully gathered until that context exits. Without distinguishing these cases, a full parameter can remain allocated into a later forward, or fallback cleanup can partition it while a `GatheredParameters` block is still using it. ## Solution This PR: - detects parameter access during Dynamo guard evaluation and skips the eager fallback all-gather; - partitions leftover nonpersistent full parameters before the next outermost forward when the normal post-backward cleanup did not run; - removes a parameter from fallback cleanup when it is passed to `GatheredParameters`, so that context alone partitions it on exit; - restores the `GatheredParameters` state even when context exit raises; and - rejects nested `GatheredParameters` contexts that overlap on the same parameter, while continuing to allow nesting over disjoint parameter sets. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> |
||
|
pengdurice
|
76928ff9bb |
Wiring up offload_opt_states (#8186)
### Make DeepCompile's optimizer-state offloading work under inductor and reachable from config #### Summary of Changes - **Registered ops instead of closures** — four ops in the existing `dc` namespace. The graph carries a tensor anchor and an integer index; live tensors stay in module state. - **ORDERED effect registration** so stock inductor keeps the ops *and* preserves their order. Reload-before-sync is correctness-critical. A test fails without it. - **Config wiring, capacity-first schedule** `[(0,[z3]), (1,[for_init, z3, move_opt_states])]`: states are emptied to host before profiling, so the plan is made against the floor and a job that only fits with offloading never runs a step with everything resident. (The pass author's own ordering from their test harness; the budget formula is unchanged.) - **Three stream races fixed**, all silent if wrong: copies wait for the optimizer's writes; `record_stream` protects reload buffers from early reuse; and the copy stream waits for the compute stream before writing a reload buffer — without it a mid-backward reload overwrites a live activation, which showed up as NaN losses. All are stream dependencies, no host waits. - **Reload buffers allocate from the pool with room** (compute stream, just freed by backward). The wrong pool costs an allocator retry plus a device-wide sync per step. - **Multi-graph gating** so graph breaks do not duplicate copies; `empty_cache` once per compile phase rather than per step (per-step measured +28%); mutual exclusion with `offload_parameters`. #### Results Qwen3-14B, 8×H200 (141 GB), ZeRO-3, micro-batch 4, fp32 states (22.2 GB/rank), `expandable_segments:True`. Medians over each run's final phase, single campaign. | configuration | seq 2048 | seq 2400 | |---|---|---| | plain eager ZeRO-3, no compile | OOM | OOM | | ZeRO-Offload (CPU Adam) | 10.11 | OOM | | DeepCompile, no offload (z3 only) | 1.89 | dies at the recompile | | DeepCompile, no offload (default schedule) | 1.82 | dies at the recompile | | this pass, blocking variant (eager-only) | 4.06 | 4.42 | | **this pass (async), eager** | 2.43 | 4.32 | | **this pass (async), inductor** | **2.01** | **3.70** | | this pass at micro-batch 1 | 0.76 — planner offloads nothing | 0.85 | #### Limitations - Runs **in place of** prefetch and selective gather (worth 3–27%); combining them is future work. - Designed for `gradient_accumulation_steps=1`: the graph runs per micro-batch, so accumulation repeats the whole cycle. Documented in the config docstring. - Mutually exclusive with `offload_parameters`. #### Tests `tests/unit/v1/compile/test_offload_opt_states.py` — op and ORDERED-effect registration; a mechanism test compiling side-effect ops through stock inductor and asserting program order; budget planning and node placement; re-run and multi-graph gating; once-per-phase `empty_cache`; and a 2-GPU end-to-end loss-parity test whose op counters prove the ops ran in the compiled graph (reloads are skipped while profiling, so a nonzero reload count is the proof). --------- Signed-off-by: pengdurice <pengduhit@gmail.com> |
||
|
Hongwei Chen
|
df84f6d842 |
Add Triton grouped-GEMM for MoE experts on Ampere/Ada (#8180)
**Motivation** AutoEP expert computation uses torch._grouped_mm , which only has a fused grouped-GEMM kernel on Hopper (sm90) and newer. On Ampere/Ada GPUs (sm80/sm86 — A100/A6000/A40/RTX 30xx) it silently falls back to a per-group Python loop ( _grouped_mm_fallback ), issuing one at::mm per expert plus a device→host sync on offs . This is slow for MoE, which invokes grouped GEMM 3× per layer in the forward and again in the backward. **What this PR does** Adds deepspeed/moe/group_gemm_triton.py : an autograd-aware Triton grouped GEMM that is a drop-in for torch._grouped_mm (2D×3D + offs ), and wires it into GroupedExperts so it is auto-selected on sm < 9.0 (native path is kept on sm90+). **Key points:** • Single fused kernel per grouped GEMM, no per-group launches and no device→host sync (per-tile → group mapping is computed on device from offs ). • Support forward + backward for both operands, supporting fp16/bf16/fp32. • trans_b option: expert weights are passed in their native [E, N, K] layout and transposed via strides inside the autograd Function, keeping .transpose off the tape. The weight gradient is produced directly in the native layout, avoiding a contiguous-materialization copy in backward. • Device-side group metadata via a small Triton kernel ( _group_meta_kernel ). • Auto-selection is overridable via DS_DISABLE_TRITON_GROUPED_MM=1 . • Auto-tunning won't be called twice for same experts shape. **Performance** Configuration: `dim=2048`, `hidden=768`, `avg_tokens/expert=64`, `dtype=bf16`, distribution: **unbalanced** | Experts | Tokens | Method | Forward (ms) | Fwd + Bwd (ms) | Backward (ms) | |---:|---:|:---|---:|---:|---:| | 16 | 760 | Triton | 0.107 | 0.589 | 0.482 | | 16 | 760 | Torch | 0.256 | 1.093 | 0.837 | | 16 | 760 | For-loop | 0.776 | 6.337 | 5.561 | | 32 | 1,848 | Triton | 0.213 | 0.622 | 0.410 | | 32 | 1,848 | Torch | 0.755 | 2.222 | 1.466 | | 32 | 1,848 | For-loop | 1.603 | 22.732 | 21.129 | --------- Signed-off-by: Hongwei <hongweichen@microsoft.com> Signed-off-by: Hongwei Chen <hongweichen@microsoft.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> |
||
|
Masahiro Tanaka
|
f6a386a320 |
Preserve tensor learning rates in OneCycle (#8205)
Follow up #8202: When an optimizer starts with a tensor learning rate, OneCycle initialization replaces it with a Python scalar while applying `cycle_min_lr`. This loses the caller's tensor identity, shape, and dtype before later scheduler updates can preserve them. This PR initializes OneCycle learning rates through the existing tensor-aware update helper, matching the path used by subsequent scheduler steps. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> |
||
|
Zupeng Wang
|
ebf1531ebc |
Add keyword argument support to activation checkpointing (#8182)
## Summary DeepSpeed's activation checkpointing wrapper currently accepts positional arguments only, while `torch.utils.checkpoint.checkpoint` also supports keyword arguments. This change allows callers to pass keyword arguments through `deepspeed.checkpointing.checkpoint`. Keyword names and non-Tensor values are retained for reconstruction during the forward and recompute passes. Tensor keyword values are flattened into the inputs passed to `CheckpointFunction`, so autograd tracks them and returns their gradients correctly. The activation checkpointing documentation now describes keyword argument support, and the regression test covers both Tensor and non-Tensor keyword arguments as well as gradient propagation. ## Validation - Activation checkpointing unit tests: 27 passed - Pre-commit checks for all changed files: passed - 1-GPU CUDA correctness smoke: direct and checkpointed execution matched, with zero maximum gradient error - 2-GPU distributed CUDA correctness smoke: each rank matched direct execution, with zero maximum gradient error - DCO sign-off is included in the commit Fixes #7038 Signed-off-by: Wang Zupeng <zupenwang@gmail.com> |
||
|
iLeGend
|
619e3199df |
MoE dispatch: fix silent Tutel + TP corruption, speed up native and Tutel paths (#8195)
Follow-up to #8174. Three independent changes to `deepspeed/moe/sharded_moe.py`. ### 1. Speed up Tutel route extraction for k > 2 `topkgating` rebuilt a dense `[s, e]` one-hot mask for every route, so gating cost grew linearly with `k`. On launch-bound configurations Tutel therefore lost to the native path once `k > 2`. Since the top-k columns already name the selected experts, the per-route index/location/gate values are now a `gather`. Routes are handed to Tutel as contiguous int32 so its own casts become no-ops, and the capacity tensor is resolved once instead of forcing a device-to-host sync per use. Routing output is bit-identical. ### 2. Shard the Tutel dispatch buffer along capacity under tensor parallelism Tutel's `encode()` returns a flat `[e * c, m]` buffer, whereas the dense path produces `[e, c, m]`. The `drop_tokens(dispatched_input, dim=1)` that follows therefore sliced the **model** dim instead of the **capacity** dim. The result still reshaped cleanly into the expert input (`e * c * m / tp` divided by `e * m` happens to yield `c / tp`), so there was no error — experts silently ran on a hidden-state slice reinterpreted as capacity, producing wrong results. Reshaping the buffer to `[e, c, m]` before sharding fixes it. This is a pre-existing bug on `master`, not introduced by #8174, but #8174 widens its reach by lifting the `k == 1` restriction on `use_tutel`. ### 3. Dispatch MoE tokens by index instead of a dense one-hot einsum The native path materialised a dense `[s, e, c]` one-hot and ran `einsum("sec,sm->ecm")`, costing `O(s * e * c * m)` to move only `s * k * m` elements. Tokens are now routed through their capacity slots directly, so dispatch and combine each become a single gather. The sparse routing metadata already existed for Tutel, so the gate always returns it and `MOELayer` selects a backend. The dense return path is kept because `deepspeed/ops/transformer/inference/moe_inference.py` still consumes it. Combine accumulates in fp32, matching the einsum it replaces — the einsum accumulated in fp32 inside the tensor-core matmul, so a naive low-precision gather would have been *less* accurate than the code it replaced. ## Results Measured on 2 GPUs, bf16, forward+backward. Tutel/native step-time ratio (change 1): | k | before | after | |---|--------|-------| | 2 | 0.947 | 0.937 | | 3 | 1.078 | 1.01 | | 4 | 1.158 | 1.04 | Native path (change 3), `s=8192 m=2048 e=16 k=3`: | | step time | peak memory | |---|-----------|-------------| | before | 58.33 ms | 4353 MB | | after | 37.25 ms | 2918 MB | The native path now matches Tutel (37.25 vs 37.49 ms). ## Testing `tests/unit/moe/`: 50 passed with Tutel installed, 48 passed / 2 skipped without. - `test_sparse_dispatch_matches_dense` (k=1/2/3) checks the index path against the dense einsum for both dispatch and combine. - `TestMOETensorParallelTutel` checks Tutel against the dense path under tensor parallelism; it fails without change 2 and passes with it. Cross-checked end-to-end against `master` (fp32 max diff 1.2e-07) and verified that bf16/fp16 accuracy versus an fp32 golden reference does not regress. > `TestMOETensorParallelTutel` requires Tutel and is skipped in CI, which does not > install it. Change 2 was validated locally on 2 GPUs: max diff `7.27e-01` before the > fix, `0.00e+00` after. --------- Signed-off-by: iLeGend <824040212@qq.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> |
||
|
Olatunji Ruwase
|
7a31fe66b1 |
Add managed_gradient_accumulation for ZeRO stage 0/1 (#8184)
When managed_gradient_accumulation=false, disable micro-step tracking and treat each engine.step() as the accumulation boundary: reduce locally accumulated grads then apply the optimizer update. Stage 2/3 and pipeline remain unsupported in this change. Part fix for #8183 --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com> |
||
|
Ma, Guokai
|
eec237ee89 |
[AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP (#8168)
AutoTP + ZeRO-3 silently produced incomplete checkpoints: both export paths handled only the ZeRO data-parallel dimension and dropped the tensor-parallel shards. - ds_to_universal.py: stage3 conversion recovers the (tp,dp) grid from checkpoint file names, extracts shards under the real tp_index, and reuses the stage<=2 TP-aware merge when tp_degree>1 (DP-only path preserved for tp_degree==1 -> no regression for plain ZeRO-3). - engine.py: _zero3_consolidated_16bit_state_dict nests GatherReplacedLayerParams inside GatheredParameters so save_16bit_model gathers both DP and TP; remove the blanket autotp+zero3 training block now that checkpoint consolidation is implemented. - stage3.py: load_hp_checkpoint_state resolves the TP shard before the ZeRO-DP partition, so universal checkpoint restore round-trips. Add end-to-end universal conversion tests and update existing tests for the refactored merge_tp_slices / extract_zero_shards_stage3 signatures. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> |