primitive-gradient-accumulation

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.