Add managed_gradient_accumulation for ZeRO stage 0/1 - #8184
Add managed_gradient_accumulation for ZeRO stage 0/1#8184sfc-gh-truwase wants to merge 9 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>
![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: 00800f7a9c
ℹ️ 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".
| if not self.managed_gradient_accumulation(): | ||
| if self.enable_backward_allreduce and not self.inside_no_sync_ctxt: | ||
| self.allreduce_gradients() |
There was a problem hiding this comment.
When ZeRO stage 1 uses active DeepCompile with managed_gradient_accumulation=false, every backward reports a non-boundary, so deepspeed/compile/init_z1.py:147-153 clears optimizer.averaged_gradients; this added step-time call then reaches allreduce_gradients(), which returns immediately for DeepCompile at engine.py:2773-2774. The subsequent optimizer step indexes the missing averaged gradients and fails instead of updating, so this combination must either be rejected or given a DeepCompile-aware step-time reduction path.
Useful? React with 👍 / 👎.
| if not self.managed_gradient_accumulation(): | ||
| # Unmanaged mode: micro-step tracking is disabled and the caller's | ||
| # step() call is the accumulation boundary. This flag is only set | ||
| # while step() runs, so backward accumulates and step reduces/updates. | ||
| return self._running_engine_step |
There was a problem hiding this comment.
When ZeRO stage 0 uses Apex AMP with managed_gradient_accumulation=false, this branch returns False during every backward, making the existing delay_unscale = not self.is_gradient_accumulation_boundary() always true. The boundary only becomes true later inside step(), after the final amp.scale_loss context has already exited, so Apex never receives the required non-delayed unscale before the optimizer update and the accumulated gradients are not finalized correctly; reject this combination or perform the final AMP unscale before reducing and stepping.
Useful? React with 👍 / 👎.
stas00
left a comment
stas00
left a comment
There was a problem hiding this comment.
Looks good, but missing docs/examples of how one would do managed gas
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>
|
@stas00 thanks. I will work on docs: https://deepspeed.readthedocs.io/en/latest/training.html#gradient-accumulation |
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
|
@stas00, see the docs in the staging: https://deepspeed.readthedocs.io/en/rtd-staging/training.html#gradient-accumulation |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4c6ae671ec
ℹ️ 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".
| if not self.managed_gradient_accumulation(): | ||
| if self.enable_backward_allreduce and not self.inside_no_sync_ctxt: | ||
| self.allreduce_gradients() |
There was a problem hiding this comment.
When managed_gradient_accumulation=false is used with ZeRO stage 1 and overlap_comm=true, this step-time call does not actually enqueue gradients for reduction: DeepSpeedZeroOptimizer.reduce_gradients() skips its param walk when overlap_comm is true, while the backward hooks also skipped reduce_ready_partitions_and_remove_grads() because the boundary was false during every backward. In multi-rank runs each optimizer partition is then updated from local, unaveraged gradients, so this combination should either be rejected or force a reduction path that walks the params at step().
Useful? React with 👍 / 👎.
| * **ZeRO stage 2/3**: gradients are still reduced/partitioned on every ``backward()`` (unchanged | ||
| from managed mode); only the ``averaged_gradients`` finalization and the optimizer update are | ||
| deferred to ``step()``. | ||
| * **ZeRO optimizer offload (CPU/NVMe)**: the boundary-time gradient-norm computation and the | ||
| fp32/optimizer-buffer copies run during ``step()``. |
There was a problem hiding this comment.
These bullets describe ZeRO-2/3 and optimizer-offload behavior for unmanaged accumulation, but initialization currently rejects ZeRO-2/3 with managed_gradient_accumulation=False and the note below says that support is still future work. Users following this section will get an init-time AssertionError instead of the described behavior, so the stage-2/3/offload bullets should be removed or clearly marked as not available yet.
Useful? React with 👍 / 👎.
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>
|
@sfc-gh-truwase, looking clear and it's easy to understand, thank you. I just have a question about:
If this is the case, shouldn't the example use W/o it how would the example even work correctly? |
…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>
| model_engine.step() # reduce + optimizer update | ||
| # scale_wrt_gas=False disables DeepSpeed's 1/gradient_accumulation_steps scaling; | ||
| # average over the actual micro-batch count instead. | ||
| model_engine.backward(loss / num_micro_batches, scale_wrt_gas=False) # accumulate only |
There was a problem hiding this comment.
I first missed the division - might it be more prominent / obvious / teachable if you were to add an explicit code, something along the lines of:
scaled_loss = loss / num_micro_batches
model_engine.backward(scaled_loss, ...
I'm not sure if scaled_loss is the best temp var name
The reader can then "optimize" it away in their own code.
tohtana
left a comment
tohtana
left a comment
There was a problem hiding this comment.
Hi Tunji,
This is a very useful setting to make DeepSpeed more flexible for various use cases.
I left a few inline comments. We would need to fix the tests at least.
| bool: if the current step is a gradient accumulation boundary. | ||
|
|
||
| """ | ||
| if self._is_gradient_accumulation_boundary is None: |
There was a problem hiding this comment.
For unmanaged mode to work as intended, _is_gradient_accumulation_boundary must remain None so that is_gradient_accumulation_boundary() uses the step-owned _running_engine_step marker. However, set_gradient_accumulation_boundary() can still set it to True or False, and either value breaks this contract. With False, the step-time allreduce_gradients() call performs no reduction and the optimizer update is skipped. With True, the backward epilogue reduces the gradients, and step() then enters the reduction path again.
Can we reject set_gradient_accumulation_boundary() when managed_gradient_accumulation=False?
There was a problem hiding this comment.
@tohtana, yes that makes sense.
Seperately. do you think it user experience will be simplified by deprecating set_gradient_accumulation_boundary() in the long run?
| def build_managed_gas_config(zero_stage, gradient_accumulation_steps, managed_gradient_accumulation): | ||
| """fp32 config toggling managed_gradient_accumulation for exact managed-vs-unmanaged comparison.""" | ||
| return { | ||
| "train_micro_batch_size_per_gpu": 2, |
There was a problem hiding this comment.
total_samples in random_dataloader() is the number of individual samples, not the number of micro-batches. Here it is 8, and train_micro_batch_size_per_gpu is 2, so the loader yields only 4 micro-batches. However, the assertion expects 8 step results, corresponding to 8 micro-batches.
The same accounting issue occurs below: 12 samples produce only 6 micro-batches, but the test indexes 12.
| .. note:: | ||
| Unmanaged mode is being added incrementally. Only ZeRO stage 0/1 (and DDP) is supported today; | ||
| ZeRO stage 2/3 and ZeRO optimizer offload are planned but **not yet available** -- enabling them | ||
| with ``managed_gradient_accumulation=false`` raises an ``AssertionError`` at initialization. |
There was a problem hiding this comment.
The new check rejects ZeRO stages 2/3, but ZeRO-1 optimizer offload still initializes with unmanaged accumulation.
There was a problem hiding this comment.
@tohtana good catch. I have follow up PRs for stages 2/3.
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>
3 participants
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.
Part fix for #8183