Unmanaged gradient accumulation: ZeRO offload support - #8225
Open
sfc-gh-truwase wants to merge 5 commits into
Open
Unmanaged gradient accumulation: ZeRO offload support#8225sfc-gh-truwase wants to merge 5 commits into
sfc-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>
sfc-gh-truwase
requested review from
loadams,
tjruwase and
tohtana
as code owners
August 6, 2026 15:22
![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>
delock
reviewed
Aug 8, 2026
delock
reviewed
Aug 8, 2026
Collaborator
|
Hi @sfc-gh-truwase , I have left my comments. One impression is |
delock
approved these changes
Aug 8, 2026
delock
reviewed
Aug 8, 2026
delock
requested changes
Aug 8, 2026
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>
Collaborator
Author
|
@delock, thanks for catching the issues with |
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
2 participants
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.
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