Skip to content

fix(zero3): async grad offload + pinned offload buffers by default #8207

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sfc-gh-truwase merged 7 commits into from
Aug 7, 2026
Merged

fix(zero3): async grad offload + pinned offload buffers by default #8207

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d6f43c5
fix(zero3): async grad offload + pinned offload buffers by default
delock Aug 3, 2026
d537ae3
refactor(zero3): remove orphaned async grad copy stub
delock Aug 3, 2026
d4bc2fc
Merge branch 'master' into gma/async_zero_offload
delock Aug 4, 2026
12bdbc5
docs(zero): document pin_memory default, mechanism, and memlock caveat
delock Aug 4, 2026
fbb0277
feat(accel): track pinned host memory for OOM diagnostics
delock Aug 5, 2026
0b8426f
Merge branch 'master' into gma/async_zero_offload
delock Aug 6, 2026
ef6a75d
Merge branch 'master' into gma/async_zero_offload
sfc-gh-truwase Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions deepspeed/runtime/zero/offload_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,13 @@ class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel):
NVMe is enabled.
"""

pin_memory: bool = False
pin_memory: bool = True
"""
Offload to page-locked CPU memory. This could boost throughput at the cost
of extra memory overhead.
Offload to page-locked (pinned) CPU memory. Required for asynchronous,
full-bandwidth GPU<->CPU transfers and for overlap of grad/param offload
with compute. Defaults to True. Disable only on hosts with tight memlock
limits (ulimit -l) or very limited resident RAM, since pinned memory
cannot be paged out.
"""


Expand All @@ -69,10 +72,13 @@ class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel):
gradient, momentum, and variance).
"""

pin_memory: bool = False
pin_memory: bool = True

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Honor legacy pin-memory opt-out

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 👍 / 👎.

"""
Offload to page-locked CPU memory. This could boost throughput at the cost
of extra memory overhead.
Offload to page-locked (pinned) CPU memory. Required for asynchronous,
full-bandwidth GPU<->CPU transfers and for overlap of grad/param offload
with compute. Defaults to True. Disable only on hosts with tight memlock
limits (ulimit -l) or very limited resident RAM, since pinned memory
cannot be paged out.
"""

pipeline_read: bool = False
Expand Down
11 changes: 2 additions & 9 deletions deepspeed/runtime/zero/stage3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1772,14 +1772,6 @@ def set_norm_for_param_grad_in_gpu(self, param):
#Using a more memory efficient version
self.norm_for_param_grads[param_id] = self._constant_buffered_norm2(param.grad)

def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param, fp32_grad_tensor):
with get_accelerator().stream(self.copy_grad_stream):
param_id = self.get_param_id(param)
src_tensor = param.grad.view(-1).to(dtype=self.master_weights_and_grads_dtype)
#print(f"src_tensor {src_tensor.size()} and fp32 grad {fp32_grad_tensor.size()}")
fp32_grad_tensor.copy_(src_tensor, non_blocking=True)
param.grad = None

def complete_grad_norm_calculation_for_cpu_offload(self, params):
self._assert_same_partition_group(params)
process_group = self._get_param_partition_group(params[0])
Expand Down Expand Up @@ -1858,7 +1850,8 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
else:
fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(
0, dest_offset, grad_buffer.numel())
fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype))
fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype),
non_blocking=True)

# free the gradient
if not get_accelerator().is_synchronized_device():
Expand Down
Loading