Olatunji Ruwase

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>

Z3: Support for activation ckpt with frozen params (#8148) Build on #8130 --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Route ZeRO/SuperOffload pin sites through accelerator pin_memory (#8256) ## Summary - Route ZeRO-1/2 `offload_states` hp/lp pins, `offload_optimizer_states`, ZeRO++ secondary shards, and SuperOffload grad buffers through `get_accelerator().pin_memory()` so `DS_PIN_MEMORY_BACKEND=native` applies. - Use `make_copy=False` for empty scratch destinations to avoid an extra host alloc/copy under the native backend. - Clarify in `memory.rst` that ZeRO `offload_*.pin_memory` (whether) is orthogonal to `DS_PIN_MEMORY_BACKEND` (how); env-only, no ds_config backend field. - Add `tests/unit/v1/pin_memory/test_offload_route.py` covering `offload_optimizer_states` under native. ## Test plan - [x] `pre-commit run --files` on touched paths (already run locally) - [x] `DS_PIN_MEMORY_BACKEND=native` + `pytest tests/unit/v1/pin_memory/test_offload_route.py` (+ pin_memory/accelerator UTs) on H200 host with pin_memory op — **23 passed** (incl. both offload_route tests) - [x] Smoke ZeRO-1/2 CPU offload with default torch backend — `test_offload_states.py` + `test_destroy_unpin.py` — **133 passed** - [x] SuperOffload worker pin site: source uses `get_accelerator().pin_memory(..., make_copy=False)`; native empty-buffer pin/unpin smoke OK Evidence: autorun `job-20260818T143819Z` on `e6dc5f9d` / `tjruwase/pin-memory-route-zero` (EXIT 0). GitHub CI also green on the tip commit. --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

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>

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>

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>

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>

Register native pinned host memory with CUDA for GPU DMA (#8283) ## Summary - Native pin (`posix_memalign` + `mlock`) is device-independent and sufficient for DeepNVMe, but the GPU DMA engine does not treat `mlock` pages as pinned. After allocation, optionally call `cudaHostRegister` (and `cudaHostUnregister` before free) via accelerator hooks so `copy_(..., non_blocking=True)` can DMA. - Default-on via `DS_PIN_MEMORY_REGISTER_DEVICE` (`1`/`true`/`yes`/`on`; `0` keeps mlock-only). Registration failure logs once and continues with mlock. CPU accelerators no-op. - Adds `benchmarks/pin_memory/h2d_d2h_bench.py` (torch vs native-unregistered vs native-registered). ## H200 H2D/D2H (1 GPU) Host `tunji-h200-n1g2-ds-pin-0`, NVIDIA H200, `--sizes-mib 4 64 256 --warmup 10 --iters 50`, autorun `job-20260820T155549Z`. | Arm | Size (MiB) | H2D (GB/s) | D2H (GB/s) | `torch.is_pinned` | |-----|------------|------------|------------|-------------------| | torch | 4 | 53.20 | 52.96 | True | | torch | 64 | 55.32 | 52.49 | True | | torch | 256 | 55.41 | 52.30 | True | | native-unregistered | 4 | 16.95 | 14.05 | False | | native-unregistered | 64 | 16.45 | 15.45 | False | | native-unregistered | 256 | 9.98 | 16.42 | False | | native-registered | 4 | 53.16 | 52.78 | True | | native-registered | 64 | 55.29 | 54.79 | True | | native-registered | 256 | 55.40 | 54.94 | True | Registered native matches torch bandwidth (~53–55 GB/s). `mlock` without `cudaHostRegister` stays in the ~10–17 GB/s pageable range and reports `torch.is_pinned=False`. ## Test plan - [x] `pre-commit run --files` on touched paths - [x] `pytest tests/unit/v1/pin_memory/test_pin_memory.py` **16 passed** (register on/off, CPU no-op, CUDA cudart mock, register failure keeps mlock, GC unregister, invalid env) - [x] GPU autorun `job-20260820T155549Z`: pin_memory + accelerator tests **25 passed**, H2D/D2H bench exit 0 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>

Mixed-precision: per-policy param/buffer dtype cast (preserve fp32 buffers) (#8066) ## Summary - Add `data_types.param_dtype` and `data_types.buffer_dtype` (both default `None`), mirroring FSDP `MixedPrecisionPolicy`. - Replace the blanket `module.half()` / `module.bfloat16()` in `_configure_distributed_model` with a targeted cast: parameters go to `param_dtype`; floating buffers keep their loaded dtype unless `buffer_dtype` is explicitly set. ## Motivation The blanket cast downcasts every floating buffer, including the rotary `inv_freq` buffer that HF/FSDP2 keep in fp32. On long contexts the bf16 `inv_freq` loses precision, RoPE angles drift, and logits/grads diverge from the FSDP2 reference. Preserving fp32 buffers by default fixes this; `buffer_dtype` is the escape hatch to reproduce the legacy behavior. ## Behavior - `param_dtype` unset -> derived from the fp16/bf16 enabled flag (legacy param behavior). - `buffer_dtype` unset -> buffers keep their loaded dtype (e.g. fp32 `inv_freq`). - `buffer_dtype` set -> buffers force-cast (legacy blanket-cast parity). ## Test plan - [ ] `param_dtype=bf16`, `buffer_dtype` unset -> params bf16, `inv_freq` stays fp32. - [ ] `buffer_dtype=bf16` -> buffers downcast (legacy parity). - [ ] bf16/fp16 run with neither key set behaves as before except fp32 buffers preserved. - [ ] 8B / 32B ZeRO-3 long-context run -> grad_norm tracks the FSDP2 reference. Made with [Cursor](https://cursor.com) --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Signed-off-by: Stas Bekman <stas@stason.org> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com> Co-authored-by: Stas Bekman <stas00@users.noreply.github.com> Co-authored-by: Stas Bekman <stas@stason.org>

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>