L6.3 · LESSON 3 OF 6

Batches and accumulation

Describe a pretraining loop and reason about data, batch size, precision and a simplified compute budget.

THE IDEA

How many tokens per update

One optimizer update is based on a batch of sequences. Its size in token positions is batch × sequence length × accumulation steps × devices. When the batch you want does not fit in memory, run several smaller microbatches, add up their gradients, and update once. This gradient accumulation gives the same gradient (for equal-sized microbatches) as the big batch.

microbatch size4sequence length8accumulation steps2devices1

tokens per update = 4 × 8 × 2 = 64 memory needed at once: only 4 sequences of 8 tokens (32 positions) per device

Accumulation runs several microbatches one after another, adds their gradients and only then updates the weights. With equal-sized microbatches the averaged gradient equals the gradient of one big batch, without holding the big batch in memory.

Change the numbers. Memory depends on the microbatch only; the update size depends on all four.

Predict first. Batch 4, length 8 and accumulation 2 correspond to how many token positions per update?

Python

opt.zero_grad()
for micro in microbatches:        # accum steps
    loss = loss_fn(model(micro.x), micro.y) / accum
    loss.backward()               # gradients add up in .grad
opt.step()

Chapter context

Pretraining optimizes next-token predictions over a dataset. Its costs depend on the architecture, sequence lengths, implementation and hardware. A useful budget model states which work it approximates and which work it ignores.

TRY IT YOURSELF

Part 1. Do it hands-on

Compute update sizes by hand, then write the arithmetic and prove that accumulation equals one big batch.

microbatch size4sequence length8accumulation steps2devices1

tokens per update = 4 × 8 × 2 = 64 memory needed at once: only 4 sequences of 8 tokens (32 positions) per device

4 × 8 × 2 16 × 1024 × 8 updates to see 1,048,576 tokens

To do: Batch 4, length 8, accumulation 2: token positions per optimizer update?
To do: Batch 16, length 1024, accumulation 8?
To do: How many such updates does it take to see 1,048,576 tokens once?

Part 2. Do it in code

tokens_per_update(batch, seq, accum, devices=1): batch · seq · accum · devices. All arguments must be at least 1 (ValueError).
updates_needed(total_tokens, per_update): ceil(total_tokens / per_update).
linear_grad(w, X, y): the gradient of the mean squared error mean((X@w − y)²) with respect to w, i.e. (2/n)·Xᵀ(Xw − y).
accumulated_grad(w, X, y, micro): split the rows into `micro` equal parts (ValueError if the rows are not divisible), compute linear_grad on each part and average the gradients.

import numpy as np

def tokens_per_update(batch, seq, accum, devices=1):
    pass

def updates_needed(total_tokens, per_update):
    pass

def linear_grad(w, X, y):
    pass

def accumulated_grad(w, X, y, micro):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Batch 4, length 8 and accumulation 2 correspond to how many token positions per update?

A64 before any masking exclusions: 4 × 8 = 32 per microbatch, accumulated twice.B32, because accumulation does not add positions.C14, because the three numbers are added.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

A commonly used rough dense-transformer training estimate is

C\approx6ND,

with N model parameters, D training tokens and C floating-point operations. It omits some architecture-dependent work and is not a hardware throughput guarantee. Effective token batch is approximately $B\times T\times A$ for microbatch B, length T and A accumulation steps.

Worked example

For N=1 million parameters and D=100 million training tokens, the rough compute estimate is 6×10^6×10^8=6×10^14 FLOPs. A microbatch of 4 sequences of length 8 accumulated twice contributes 64 token positions per optimizer update. Neither calculation includes padding exclusions or communication time.

Reference code

The following example uses the Python standard library and does not call a remote service.

N,D=1_000_000,100_000_000
assert 6*N*D == 600_000_000_000_000
batch,length,accumulation=4,8,2
assert batch*length*accumulation == 64

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Compute-optimal ratios are model- and study-dependent. Train/test leakage can arise from shared documents and overlap. Tokens per update and optimizer steps are different units.