fix(zero-offload): mitigate CPU RSS waste on ZeRO-2 optimizer offload (#7693)#8132
fix(zero-offload): mitigate CPU RSS waste on ZeRO-2 optimizer offload (#7693)#8132liuyun7345 wants to merge 1 commit into
Conversation
![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: 9ef04cf853
ℹ️ 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".
| # result buffer even for 0-dim tensors, pinning ~4 bytes/param in the | ||
| # CPU allocator cache that is never reused by later allocations. | ||
| # mul_(float) is guaranteed in-place. See issue #7693. | ||
| combined_scale = float(combined_scale) |
There was a problem hiding this comment.
When ZeRO-1/2 runs without CPU offload and clip_grad > 0, scaled_global_norm() produces total_norm on the accelerator and this unconditional float() reads the CUDA scalar back to Python, forcing a device synchronization for every optimizer step/param group before unscaling gradients. The allocator leak described here is CPU-specific, so limiting the conversion to CPU gradients/scales avoids regressing the normal GPU ZeRO gradient-clipping path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Thanks — agreed. Updated so float(combined_scale) only runs when self.cpu_offload is true, keeping the GPU ZeRO clipping path on-device and avoiding the sync regression.
| if not getattr(self, '_first_step_done', False): | ||
| self._first_step_done = True | ||
| _malloc_trim() |
There was a problem hiding this comment.
In ZeRO-2 CPU offload, stage_1_and_2._optimizer_step() narrows self.optimizer.param_groups to a single group and calls DeepSpeedCPUAdam.step() once per bit16_groups entry. With multiple optimizer param groups, this optimizer-wide flag is set after group 0, so later groups still lazily allocate their optimizer state and temporary buffers but skip _malloc_trim(), leaving the RSS waste this patch is meant to reclaim for all parameters outside the first group.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Thanks — agreed. Replaced the optimizer-wide _first_step_done flag with per-group / per-subgroup _trimmed_groups tracking so each ZeRO param group still gets malloc_trim after its first lazy state init.
liuyun7345
force-pushed
the
fix/cpu-offload-memory-leak-7693
branch
from
9ef04cf to
8e4b0d3
Compare
liuyun7345
changed the title
| auto params_c = params.is_contiguous() ? params : params.contiguous(); | ||
| auto grads_c = grads.is_contiguous() ? grads : grads.contiguous(); | ||
| auto exp_avg_c = exp_avg.is_contiguous() ? exp_avg : exp_avg.contiguous(); | ||
| auto exp_avg_sq_c = exp_avg_sq.is_contiguous() ? exp_avg_sq : exp_avg_sq.contiguous(); |
There was a problem hiding this comment.
From this document, this change should not be necessary, PyTorch ensures return the same tensor if the tensor is already contiguous.
https://docs.pytorch.org/docs/2.13/generated/torch.Tensor.contiguous.html
There was a problem hiding this comment.
Thanks @delock — you're right. Tensor.contiguous() already returns self when the tensor is contiguous, so the explicit is_contiguous() check is a no-op. I'll remove this change from the PR.
Addresses deepspeedai#7693 with targeted mitigations for CPU-offload paths: 1. unscale_and_clip_grads: convert combined_scale to float only when cpu_offload is enabled, avoiding a GPU device-sync regression on the normal ZeRO-1/2 clipping path (Codex review). 2. DeepSpeedCPUAdam: call malloc_trim(0) once per param group/subgroup on its first step so later ZeRO groups also release RSS after lazy state init (Codex review). Also adds regression tests (LF line endings). Dropped the ds_adam_step .contiguous() guard: Tensor.contiguous() already returns self when contiguous (review feedback from @delock). This is a best-effort mitigation, not a complete resolution of all CPU memory overhead tracked in deepspeedai#7693 (e.g. H2D tiling remains open). Signed-off-by: liuyun7345 <liuyun7345@sina.com>
liuyun7345
force-pushed
the
fix/cpu-offload-memory-leak-7693
branch
from
8e4b0d3 to
b3ed969
Compare
| # float() would force a device sync every step. The allocator concern in | ||
| # #7693 is CPU-specific, so keep the GPU path on-device. See #7693 / #8132. | ||
| if self.cpu_offload: | ||
| combined_scale = float(combined_scale) |
There was a problem hiding this comment.
I didn't observed the 4byte/element memory usage from combined_scale. The pytorch version I use is torch 2.11. Did you observe this memory bloat on your environment? Thanks!
|
Hi @liuyun7345 thanks for tracking this issue and do the best effort fix. Yet I found some claims in the original issue may not be the real reason. More over, on a single 5090 system with CPU offload, we didn't observe significant RSS change before and after the PR. Did you observe RSS change after applying the patch? I would suggest to digg deeper to find the real reason for RSS usage. If the cause claimed in #7693 is not true, then we should avoid adding these complexity for no return. |
2 participants
Summary
Best-effort mitigations for CPU RSS waste in ZeRO-2 + CPU optimizer offload tracked by #7693.
This PR does not claim a complete resolution of all overhead in that issue (H2D cast/copy buffer tiling and other unaccounted bytes remain open). Another commenter also reported being unable to reproduce the two ~4 bytes/param leaks on a newer PyTorch build.
Changes
1.
unscale_and_clip_grads— float conversion only on CPU offload (stage_1_and_2.py)Convert
combined_scaleto a Pythonfloatonly whencpu_offloadis enabled, somul_receives a scalar on the CPU path.Per Codex review: an unconditional
float()would read a CUDA scalar back to Python when ZeRO-1/2 runs without CPU offload andclip_grad > 0, forcing a device sync every step. The GPU clipping path stays on-device.2.
DeepSpeedCPUAdam— per-groupmalloc_trim(0)(cpu_adam.py)Call
malloc_trim(0)once per param group / subgroup on that group's first step.Per Codex review: ZeRO-2 CPU offload invokes
step()once perbit16_groupsentry. An optimizer-wide “first step” flag would skip trim for later groups after group 0. Per-group tracking fixes that.This is a best-effort glibc RSS mitigation (Linux-only no-op elsewhere), not a correctness fix.
3. Tests (
test_cpu_offload_memory.py)Removed after review
ds_adam_stepis_contiguous()guard incpu_adam_impl.cpp. As @delock pointed out,Tensor.contiguous()already returnsselfwhen the tensor is contiguous, so the explicit check was a no-op.Test plan
test_malloc_trim_helper_existstest_malloc_trim_called_per_param_grouptest_cpu_adam_first_step_initializes_optimizer_statestest_unscale_and_clip_grads_cpu_offload_uses_float_scaletest_unscale_and_clip_grads_no_cliptest_unscale_and_clip_grads_gpu_path_skips_float_conversionNotes
.contiguous()change.Signed-off-by: liuyun7345 <liuyun7345@sina.com>.Addresses #7693