fix(zero3): async grad offload + pinned offload buffers by default - #8207
Draft
delock wants to merge 4 commits into
Draft
fix(zero3): async grad offload + pinned offload buffers by default#8207delock wants to merge 4 commits into
delock wants to merge 4 commits into
Conversation
The ZeRO-3 grad GPU->CPU offload copy in partition_grads() was a blocking copy_() without non_blocking, and its destination buffer defaulted to pageable host memory. Together these forced the grad 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. - offload_config.py: default offload_optimizer/offload_param pin_memory to True. Pinned (page-locked) host memory is required for asynchronous, full-bandwidth DMA; the prior False default silently selected the slow staged pageable copy. Disable only on hosts with tight memlock limits. - stage3.py: issue the grad offload copy with non_blocking=True. The copy stays on reduce_and_partition_stream, and correctness is preserved by the existing reduce_and_partition_stream.synchronize() barriers ahead of step(). Validated on 4x RTX 4080-SUPER (autotp=2, offload_optimizer, cpu_adam): BWD -18% to -28% (1.5B/3B), no memory regression, FWD/STEP unchanged. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
`async_inplace_copy_grad_to_fp32_buffer_from_gpu` in stage3.py has been dead code since deepspeedai#1453 (2022-01, "Various ZeRO Stage3 Optimizations + Improvements"). That rewrite moved grad handling from the per-param path to the batched partition_grads() path and deleted both the call site and the `self.copy_grad_stream` initialization, but left this function definition behind. It now (a) has zero callers in stage3 and (b) references `self.copy_grad_stream`, an attribute that no longer exists in stage3, so it would raise AttributeError if ever invoked. Keeping it is a maintenance hazard: it reads like working async-offload infrastructure and obscures the fact that stage3's grad offload was actually synchronous (fixed in the previous commit). The live, called version in stage_1_and_2.py (still using non_blocking, called at ~line 1590) is unaffected and intentionally left in place. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
delock
marked this pull request as ready for review
August 4, 2026 10:40
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d537ae3764
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| """ | ||
|
|
||
| pin_memory: bool = False | ||
| pin_memory: bool = True |
There was a problem hiding this comment.
When a config still uses the supported deprecated form cpu_offload: true together with cpu_offload_use_pin_memory: false, DeepSpeedZeroConfig maps cpu_offload to DeepSpeedZeroOffloadOptimizerConfig(device=cpu) and the deprecated pin-memory field has set_new_param=False, so this new default makes that explicit opt-out pin memory anyway. On hosts with low memlock limits this can turn previously working legacy configs into initialization failures; wire the deprecated flag into the new offload config or keep the default opt-out-compatible for the migration path.
Useful? React with 👍 / 👎.
delock
marked this pull request as draft
August 4, 2026 13:54
Expand the pin_memory descriptions for offload_param and offload_optimizer in config-json.md to cover the mechanism (async full-bandwidth DMA + compute overlap), the precondition (overlap_comm for the grad offload path), and the failure mode (pinned memory is non-swappable, counts against ulimit -l; may cause init failures or indirect OOM on hosts with tight memlock limits). Update the default to true and add a release notice. Fix memory.rst Pinned Memory section: replace the deprecated cpu_offload_use_pin_memory flag with the new offload_optimizer / offload_param pin_memory fields, correct the stale 'ZeRO-2 can't be controlled' claim, and note the new default. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
ZeRO-3's gradient GPU→CPU offload in
partition_grads()used a blockingcopy_()withoutnon_blocking, and its destination buffer defaulted topageable 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: defaultoffload_optimizer/offload_param.pin_memoryto
True. Pinned (page-locked) host memory is required for async,full-bandwidth DMA; the prior
Falsesilently selected the slow stagedpageable copy. Disable only on hosts with tight
ulimit -lmemlock.stage3.py: issue the grad offload copy withnon_blocking=True. Stayson
reduce_and_partition_stream.stage3.py(2nd commit): remove an orphaned helper(
async_inplace_copy_grad_to_fp32_buffer_from_gpu) that referenced anuninitialized 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-rankCPU affinity:
Memory footprint unchanged; FWD/STEP unchanged.