Motivation
DeepSpeed owns the gradient-accumulation boundary via an internal micro-step counter:
backward() bumps micro_steps, and step() only reduces gradients + runs the optimizer when
(self.micro_steps + 1) % self.gradient_accumulation_steps() == 0
This enables symmetric forward/backward/step on every micro-batch, which is ideal for the
standard training loop.
It breaks down for callers that want to own the boundary themselves — notably HTTP/RPC-driven RL
backends (e.g. Tinker-protocol servers) where one logical optimizer step arrives as N backward()
calls followed by exactly one step(), with N client-controlled and not known at config time.
FSDP-based backends handle this natively (they call optimizer.step() when they choose); DeepSpeed's
only lever today is set_gradient_accumulation_boundary(bool), which must be toggled around every
backward, is easy to get wrong, and couples the caller to engine internals.
Proposal
Add a boolean ds_config flag managed_gradient_accumulation (default True, i.e. current behavior).
When False ("unmanaged" mode):
- micro-step tracking is disabled; there is no counter-driven boundary,
is_gradient_accumulation_boundary() is True only while step() executes,backward() accumulates; step() finalizes accumulated grads and applies exactly one optimizer
update — unconditionally, every call,- the caller is responsible for calling
step() at the true boundary.
This gives a first-class, config-level API for "backward and step are separate primitives, the client
owns the boundary" without per-backward toggling.
Per-stage semantics
- ZeRO 0/1 / DDP:
backward() accumulates gradients locally (no reduction); step() performs the
all-reduce then the optimizer update.
- ZeRO 2/3: keep reducing/partitioning on every
backward() (their natural behavior), but defer
the averaged_gradients finalization + optimizer update to step() via a
finalize_gradient_accumulation_boundary() hook.
- ZeRO offload (CPU/NVMe optimizer): supported; the boundary-time grad-norm computation and
fp32/optimizer-buffer copies move into that finalize hook invoked from step().
Gradient scaling (open question)
DeepSpeed scales loss/grads by the configured gradient_accumulation_steps:
gas_scaled_loss = loss / self.gradient_accumulation_steps() if scale_wrt_gas else loss
In unmanaged mode the number of backward() calls per step may differ from the configured GAS (and
may vary per step). Callers who accumulate a variable count should either set
gradient_accumulation_steps accordingly or use backward(loss, scale_wrt_gas=False) and average
client-side. Open question: should unmanaged mode default to scale_wrt_gas=False, or expose a
per-step "effective accumulation count"?
Non-goals / open questions
- Pipeline parallelism (
PipelineModule.train_batch() owns the whole schedule) is out of scope
for v1; managed_gradient_accumulation=False would be rejected there.
- Naming:
managed_gradient_accumulation vs external_grad_accum_boundary / manual_optimizer_step.
- Scaling knob: should there be a first-class control for variable accumulation counts (see above)?
- Interactions with ZenFlow, eigenvalue computation, and gradient clipping edge cases.
Reference implementation
A working prototype exists covering ZeRO 0-3 including CPU optimizer offload, with unit tests asserting
numerical parity against managed mode (N backwards + 1 step == managed GAS=N). Happy to contribute as
stacked PRs (ZeRO 0/1 -> 2 -> 3 -> offload) if there's interest.
Motivation
DeepSpeed owns the gradient-accumulation boundary via an internal micro-step counter:
backward()bumpsmicro_steps, andstep()only reduces gradients + runs the optimizer whenThis enables symmetric
forward/backward/stepon every micro-batch, which is ideal for thestandard training loop.
It breaks down for callers that want to own the boundary themselves — notably HTTP/RPC-driven RL
backends (e.g. Tinker-protocol servers) where one logical optimizer step arrives as N
backward()calls followed by exactly one
step(), with N client-controlled and not known at config time.FSDP-based backends handle this natively (they call
optimizer.step()when they choose); DeepSpeed'sonly lever today is
set_gradient_accumulation_boundary(bool), which must be toggled around everybackward, is easy to get wrong, and couples the caller to engine internals.
Proposal
Add a boolean ds_config flag
managed_gradient_accumulation(defaultTrue, i.e. current behavior).When
False("unmanaged" mode):is_gradient_accumulation_boundary()isTrueonly whilestep()executes,backward()accumulates;step()finalizes accumulated grads and applies exactly one optimizerupdate — unconditionally, every call,
step()at the true boundary.This gives a first-class, config-level API for "backward and step are separate primitives, the client
owns the boundary" without per-backward toggling.
Per-stage semantics
backward()accumulates gradients locally (no reduction);step()performs theall-reduce then the optimizer update.
backward()(their natural behavior), but deferthe
averaged_gradientsfinalization + optimizer update tostep()via afinalize_gradient_accumulation_boundary()hook.fp32/optimizer-buffer copies move into that finalize hook invoked from
step().Gradient scaling (open question)
DeepSpeed scales loss/grads by the configured
gradient_accumulation_steps:In unmanaged mode the number of
backward()calls per step may differ from the configured GAS (andmay vary per step). Callers who accumulate a variable count should either set
gradient_accumulation_stepsaccordingly or usebackward(loss, scale_wrt_gas=False)and averageclient-side. Open question: should unmanaged mode default to
scale_wrt_gas=False, or expose aper-step "effective accumulation count"?
Non-goals / open questions
PipelineModule.train_batch()owns the whole schedule) is out of scopefor v1;
managed_gradient_accumulation=Falsewould be rejected there.managed_gradient_accumulationvsexternal_grad_accum_boundary/manual_optimizer_step.Reference implementation
A working prototype exists covering ZeRO 0-3 including CPU optimizer offload, with unit tests asserting
numerical parity against managed mode (N backwards + 1 step == managed GAS=N). Happy to contribute as
stacked PRs (ZeRO 0/1 -> 2 -> 3 -> offload) if there's interest.