primitive-gradient-accumulation

name	primitive-gradient-accumulation
description	Compose multiple microbatch forward/backward passes into a single optimizer step, enabling effective batch sizes larger than memory permits. Activate when the user asks "gradient accumulation", "accumulate gradients", "effective batch size", "OOM at larger batch", or asks how to set the GA factor.

Primitive · GradientAccumulation

Microbatch schedule for one optimizer step. The only fully-implemented primitive in V1; serves as the model for the others.

What it does

Splits a target effective batch size into K microbatches. Per microbatch:

• zero_grad only on microbatch 0.
• Forward + backward, scaling loss by 1/K to keep gradient magnitude consistent.
• step_optimizer only on microbatch K-1.
• Optional gradient sync (DDP) only on microbatch K-1.

Interface

from curry_train.primitives import GradientAccumulation

ga = GradientAccumulation(steps=K)
ga.steps                           # K
ga.enabled                         # K > 1
ga.is_first(i), ga.is_last(i)      # microstep i is first/last
ga.train_step_kwargs(i)            # dict of {zero_grad, sync_gradients,
                                   #          step_optimizer, loss_divisor,
                                   #          collect_metrics}
ga.from_config(cfg, key="gradient_accumulation_steps")  # build from Hydra

When to use

• Effective batch size larger than fits in memory: K = effective_batch / microbatch.
• Training where you want a stable effective batch size while changing physical batch (e.g., for muP sweeps).

When NOT to use

• DDP across many ranks already provides effective-batch growth. Don't double-count.
• For very small models, the overhead of K extra forwards is real; prefer larger physical batch.

Interaction with parallelism primitives

• DDP: GA does not need any DDP awareness; the per-step sync_gradients flag is what train_step_kwargs(i) produces.
• FSDP: same as DDP; sync only on the last microbatch.
• TP: GA is orthogonal; each microbatch does its own TP comms.
• PP: GA microbatch count is usually replaced by PP schedule's microbatch count. Don't stack.

Boundaries

• GA is a scheduler, not an algorithm. The actual gradient computation happens in the model's forward/backward; GA just coordinates when the optimizer fires.
• loss_divisor must be applied consistently. Forgetting it produces gradients K× too large.

Implementation status

V1: fully implemented at template/curry_train/primitives/grad_accum.py.

Related

• template/curry_train/benchmark.py:run_accumulated_step — uses this primitive.
• template/curry_train/loop.py:run_training_steps — uses this primitive.
• skills/bench — bench command exercises GA.