Gradient accumulation

The implementation treats batch_size as the effective optimizer batch and micro_batch_size as the per-forward/backward slice. The invariant is that optimizer state, scheduler state, global_step, checkpoint cadence, and logging cadence all advance once per effective batch, never once per micro-batch.

For a loss that is already a mean over some units, each micro-batch must contribute in proportion to the number of units it was averaged over:

This is why the shared helper scales micro-batch losses as:

scaled_loss = mean_loss * (micro_units / total_units)
scaled_loss.backward()

The subtle part is choosing u_m correctly.

• Entangled training uses token-level next-token CE, so u_m must be the number of supervised shifted tokens in the micro-batch. Weighting by examples would change the objective when supervision lengths differ.
• Disentangled training currently preserves the existing separable objective, whose reconstruction-style terms are means over examples of per-example reduced quantities, so the separable accumulation path uses u_m = len(micro_batch_records).

router_marginal_kl_to_prior_loss is the known exception. It is computed from batch-level summary statistics, so one-pass micro-batch accumulation is only an approximation:

That term currently warns under micro-batching rather than pretending to be exact.

Implementation notes:

• micro_batch_size="auto" is resolved once at startup by trying divisors of batch_size from largest to smallest.
• The trial step runs one real forward/backward pass and then restores RNG / gradient state so resolution does not perturb the real run.
• Declared config keeps the requested value, while runtime metadata records resolved_micro_batch_size separately because the numeric result is hardware-dependent.

For shutdown/requeue behavior and faithful within-epoch restart, see Preemption recovery and resume.