Unmanaged gradient accumulation: ZeRO offload support - #8225
Unmanaged gradient accumulation: ZeRO offload support#8225sfc-gh-truwase wants to merge 5 commits into
Conversation
Allow managed_gradient_accumulation=False with ZeRO optimizer-state and parameter offload (CPU/NVMe). Stage 2/3 defer boundary norm and FP32/NVMe gradient copy into finalize_gradient_accumulation_boundary(); stage 1 continues to finalize via allreduce_gradients() at step(). Docs and equivalence tests cover stages 1-3 optimizer offload and stage-3 param offload. Validated on a 2-GPU node (full -k Unmanaged suite, 31 passed) incl. optimizer offload [1,2,3] and param+optimizer offload stage 3. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
![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: d3dcf9ad4e
ℹ️ 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".
Track params reduced in the current accumulation window and only finalize those at step(), matching managed offload (which only copies grads that flow through the boundary backward). Prevents inactive ZeRO-2/3 params from being updated from leftover CPU/partition buffers. Adds a regression test that alternates active heads across windows. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Cover optimizer NVMe and param+optimizer NVMe paths against managed mode, skipping when async-io is unavailable. Shared helper drives CPU and NVMe offload cases; NVMe uses a large enough model to exercise swap_out_gradients. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Resolve stage3.py conflict with async grad offload (#8207): keep unmanaged boundary helper and retain non_blocking FP32 grad copy. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
| assert not self.offload_optimizer and not self.offload_param, \ | ||
| "unmanaged gradient accumulation does not support ZeRO offload" | ||
| # Unmanaged mode: partitions already accumulate in __param_id_to_grad_partition; offload still needs deferred boundary copy. | ||
| self.is_gradient_accumulation_boundary = True |
There was a problem hiding this comment.
It looks like self.is_gradient_accumulation_boundary is not consumed in this path, is it necessary to set it to true? Maybe add a comments in initilization of this flag is for managed code only will help understanding.
| continue | ||
| grad_buffer = self.__param_id_to_grad_partition[param.ds_id] | ||
| if not get_accelerator().on_accelerator(grad_buffer): | ||
| grad_buffer = grad_buffer.to(get_accelerator().current_device_name(), non_blocking=True) |
There was a problem hiding this comment.
I think whether moving grads from host to device for normalization and fp32 conversion depends on host capability. Do we consider make it an option in the furture?
|
Hi @sfc-gh-truwase , I have left my comments. One impression is |
| # Unmanaged mode: grads were reduced/accumulated into all_grad_tensors each backward; finalize averaged_gradients for step(). | ||
| assert not self.cpu_offload, "unmanaged gradient accumulation does not support ZeRO optimizer state offload" | ||
| # Unmanaged mode: grads accumulated each backward; finalize for step() (averaged_gradients or offload fp32 copy). | ||
| self.is_gradient_accumulation_boundary = True |
There was a problem hiding this comment.
Same here, this variable seems not used in following path.
Move optimizer is_gradient_accumulation_boundary behind ZeROOptimizer get/set methods so the engine mirrors boundary state through an explicit API, with comments clarifying managed vs unmanaged use. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
2 participants
Summary
managed_gradient_accumulation=false) to ZeRO optimizer-state and parameter offload (CPU/NVMe). Follow-up to Unmanaged gradient accumulation: ZeRO stage 3 support #8217 (ZeRO stage 3, now merged).backward();step()finalizes deferred offload boundary work (grad norms + FP32/NVMe copy) viafinalize_gradient_accumulation_boundary().step()viaallreduce_gradients(), which already performs offload boundary finalization when the boundary flag is true.overlap_commremain unsupported.Test plan
Validated on a 2-GPU node:
-k Unmanagedsuite (31 passed), including:test_unmanaged_matches_managed_optimizer_offload[1|2|3]test_unmanaged_matches_managed_param_offload(stage 3)config-json.md,training.rst); previewable onrtd-stagingMade with Cursor