Unmanaged gradient accumulation: ZeRO stage 2 support - #8203
Open
sfc-gh-truwase wants to merge 17 commits into
Open
Unmanaged gradient accumulation: ZeRO stage 2 support#8203sfc-gh-truwase wants to merge 17 commits into
sfc-gh-truwase wants to merge 17 commits into
Conversation
When managed_gradient_accumulation=false, disable micro-step tracking and treat each engine.step() as the accumulation boundary: reduce locally accumulated grads then apply the optimizer update. Stage 2/3 and pipeline remain unsupported in this change. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Reject managed_gradient_accumulation=False with DeepCompile and Apex AMP, since both perform boundary-gated work during backward() that the unmanaged path (boundary only at step()) breaks. Move all unmanaged-mode guards into _do_sanity_check() and tersify related comments. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Address review feedback (codex r3668352534): the training.rst and config-json entries described ZeRO stage 2/3 and optimizer-offload behavior for unmanaged gradient accumulation, but initialization rejects those combinations on this branch. Reframe them as planned/not-yet-available so users are not misled into an init-time AssertionError. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…False Address review feedback from @stas00: the unmanaged-mode example called model_engine.backward(loss) while the accompanying note said to use scale_wrt_gas=False and average the loss manually. Since the caller owns the number of backward() calls per step (which may differ from the configured gradient_accumulation_steps), the default 1/gas scaling would be incorrect. Update the example to disable DeepSpeed's scaling and average over the actual micro-batch count, and clarify the note accordingly. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Address review feedback from @stas00 (r3670124969): pull the loss division out of the backward() call into an explicit averaged_loss statement so the manual averaging is more prominent and teachable. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…load in unmanaged mode, fix test micro-batch accounting - Reject set_gradient_accumulation_boundary() when managed_gradient_accumulation=False, since setting _is_gradient_accumulation_boundary breaks the step-owned boundary contract. - Reject ZeRO optimizer/param offload in unmanaged mode, closing the ZeRO-1 offload gap so behavior matches the docs (offload not yet supported). - Fix unmanaged GAS tests: build_managed_gas_config now uses micro-batch size 1 so random_dataloader's total_samples equals the micro-batch count, fixing the off-by-2x step-count assertions. - Add tests for the set_gradient_accumulation_boundary and offload rejections. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
With managed_gradient_accumulation=False, backward() accumulates locally (boundary is False) so the reduce hooks are skipped, and DeepSpeedZeroOptimizer.reduce_gradients() skips its param walk when overlap_comm=True -- leaving gradients unreduced at step(). Reject overlap_comm at initialization (defaults to False for the supported ZeRO stage 0/1) and document the incompatibility. Adds a validation test. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
In unmanaged gradient accumulation, the caller owns the boundary and the number of backward() calls per step is variable, so advancing global_samples by the fixed train_batch_size() was incorrect. Count backward() calls since the last step() and advance global_samples by the actual micro-batch count. Add a test exercising a varying number of backward() calls per step, matching a managed manual-boundary reference and validating the global_samples accounting. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Finalize averaged_gradients at step() after per-backward reduce/partition. Still rejects ZeRO stage 3, ZeRO offload, and pipeline parallelism. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Stage 2 reduces/partitions gradients on every backward, and the epilogue synchronizes the reduction stream before finalizing accumulated partitions, so overlap_comm is safe. Relax the guard to allow overlap_comm with stage 2 while still rejecting it for stage 0/1 (where reduction is deferred to step()). Add a stage-2 equivalence test (unmanaged vs managed with overlap_comm) and update docs to describe the stage-2 support. 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 1, 2026 23:44
![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: c317cbe819
ℹ️ 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".
coalesce_grad_reduction() coalesces reduction to its own with-block boundary and finalizes averaged_gradients on exit (clearing all_grad_tensors). In unmanaged mode the caller already owns the boundary via step(), so a following step() would re-finalize against a cleared all_grad_tensors and crash. Reject the combination up front at context entry (addresses Codex review on #8203). Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
stas00
approved these changes
Aug 4, 2026
|
|
||
| def finalize_gradient_accumulation_boundary(self): | ||
| # 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 offload" |
Collaborator
There was a problem hiding this comment.
would it be more user friendly to be more specific here, and say optim states offload?
param offload doens't need to be mentioned since Zero 3 isn't supported, right?
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
Extends unmanaged gradient accumulation (
managed_gradient_accumulation=false) to ZeRO stage 2. Stacked on top of the stage 0/1 foundation in #8184 (base branchsfc-gh-truwase/gas_mgmt_zero01); review that PR first.Unlike stage 0/1 (where
backward()accumulates locally andstep()performs the reduction), ZeRO stage 2 must reduce/partition gradients on everybackward()to preserve its memory characteristics. This is compatible with unmanaged mode because reduce-scatter is linear: accumulating the reduced partitions across N caller-controlled backwards is equivalent to reducing once at the boundary. Micro-step tracking stays disabled and the caller still owns the boundary; only theaveraged_gradientsfinalization is deferred tostep().ZeROOptimizer.finalize_gradient_accumulation_boundary()(stage 1/2) buildsaveraged_gradientsfrom the accumulatedall_grad_tensorsatstep().not partition_weights), still rejecting stage 3 and ZeRO offload (follow-up PRs).overlap_commis now supported for stage 2: its async reduction is confined to the per-backward path and the epilogue synchronizes before finalizing, so it behaves exactly as in managed mode. It remains rejected for stage 0/1 (where reduction is deferred tostep()).Test plan
Validated on a 2-GPU node (full
-k Unmanagedsuite, 22 passed):test_unmanaged_matches_managed[2]— unmanaged stage-2 matches managed stage-2test_unmanaged_varying_backward_count[2]— variable backward count per step, stage 2test_unmanaged_matches_managed_overlap_comm— unmanaged stage-2 withoverlap_comm=Truematches managed referencetest_unmanaged_rejects_stage3,test_unmanaged_rejects_zero_offload— stage 3 / offload rejectedtest_unmanaged_rejects_overlap_comm[0,1]— overlap_comm still rejected for stage 0/1pre-commit(yapf/flake8/codespell) cleanDocs (
config-json.md,training.rst) updated to describe stage-2 behavior and the overlap_comm support; previewable onrtd-staging.Made with Cursor