DeepSpeed
eebe24a6 - Route ZeRO/SuperOffload pin sites through accelerator pin_memory (#8256)

Commit
1 day ago
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>
Parents
Loading