Gradient accumulation — how a huge training batch fits on finite GPUs

Yesterday we saw what happens after backpropagation: AdamW turns gradients into parameter updates, while maintaining substantial optimizer state.

Today we will focus on a deceptively simple question:

If a training run wants a batch containing millions of tokens, how can that batch possibly fit in GPU memory?

Often, it doesn't.

Instead, training systems use gradient accumulation: process a large logical batch as several smaller microbatches, accumulate their gradients, and update the weights only after all of them have been processed.

The central idea is:

\[ \boxed{\text{the batch used for an optimizer update does not have to exist in memory all at once}} \]

This is one of the basic tricks that makes large-scale LLM training practical.

Start with ordinary minibatch training¶

Suppose our batch contains \(B\) training examples:

\[ x_1,x_2,\ldots,x_B. \]

For each example, the model produces a loss:

\[ \ell_i(\theta). \]

The batch loss is usually the mean:

\[ L(\theta) = \frac{1}{B} \sum_{i=1}^{B} \ell_i(\theta). \]

Backpropagation gives:

\[ \nabla_\theta L = \frac{1}{B} \sum_{i=1}^{B} \nabla_\theta \ell_i. \]

Notice something important.

The gradient of the average is the average of the gradients.

That linearity is what makes gradient accumulation work.

A tiny example¶

Suppose a batch has four examples.

Their gradients for one parameter are:

example 1:  +2
example 2:  +4
example 3:  -1
example 4:  +3

The batch-average gradient is:

\[ g = \frac{2+4-1+3}{4} = 2. \]

If all four examples fit in memory, we can compute that in one forward/backward pass.

But suppose only two examples fit at once.

Process the first microbatch:

example 1
example 2

Its gradient sum is:

\[ 2+4=6. \]

Keep that gradient.

Then process:

example 3
example 4

whose gradient sum is:

\[ -1+3=2. \]

Add it to the existing gradient:

\[ 6+2=8. \]

Finally divide by the total number of examples:

\[ \frac{8}{4}=2. \]

Same result.

We never needed all four examples' activations in GPU memory simultaneously.

Microbatch versus global batch¶

This introduces two different notions of batch size.

Suppose:

\[ B_{\text{micro}}=8 \]

examples fit on each GPU.

We have:

\[ N=16 \]

GPUs.

And we accumulate gradients for:

\[ A=32 \]

microbatch steps before updating the weights.

Then the effective global batch size is:

\[ B_{\text{global}} = B_{\text{micro}} \times N \times A. \]

So:

\[ B_{\text{global}} = 8\times16\times32 = 4096. \]

Each GPU only processes 8 examples at a time.

But AdamW sees an update corresponding to 4,096 examples.

For LLMs, it is often more useful to think in tokens rather than examples, because sequence lengths vary.

If each microbatch contributes 16,384 tokens per GPU, then:

\[ 16{,}384 \times16 \times32 = 8{,}388{,}608 \]

tokens contribute to one optimizer step.

That is an effective batch of about:

\[ 8.4\text{ million tokens}. \]

The warehouse mental model¶

Imagine you need to weigh all the grain delivered by 100 trucks.

Your scale can only hold one truck at a time.

You do not need to build a scale large enough for all 100 trucks simultaneously.

Instead:

weigh truck 1 → add result to running total
weigh truck 2 → add result to running total
...
weigh truck 100 → add result to running total

Then compute the average.

GPU memory is the scale.

Microbatches are the trucks.

The accumulated gradient is the running total.

The optimizer update happens only after the entire convoy has been measured.

Why this saves so much memory¶

During training, forward activations are expensive.

For a microbatch, the model must retain enough intermediate information to run backpropagation.

Very roughly, activation memory grows with quantities such as:

\[ B_{\text{micro}} \times T \times d_{\text{model}} \times L, \]

where:

\(B_{\text{micro}}\) is microbatch size,
\(T\) is sequence length,
\(d_{\text{model}}\) is hidden width,
\(L\) is layer count.

If we tried to make the entire global batch resident simultaneously, activation memory could explode.

Gradient accumulation lets us keep:

\[ B_{\text{micro}} \]

small while making:

\[ B_{\text{global}} \]

large.

The key is that after one microbatch's backward pass finishes, most of its activations can be discarded.

Only the accumulated parameter gradients need to survive.

Conceptually:

microbatch 1:
forward → activations → backward → gradients
                         ↓
                  discard activations

microbatch 2:
forward → activations → backward → add gradients
                         ↓
                  discard activations

...

optimizer step

But the weights must not change between microbatches¶

This detail is essential.

Suppose we process microbatch 1 and immediately update:

\[ \theta. \]

Then microbatch 2 computes gradients using different parameters.

We would no longer be reproducing the gradient of one larger batch.

So during accumulation:

microbatch 1 → backward
microbatch 2 → backward
microbatch 3 → backward
microbatch 4 → backward

the model parameters remain fixed.

Only after all accumulation steps do we run:

AdamW step

Then:

\[ \theta_t\rightarrow\theta_{t+1}. \]

This means there is an important distinction between a backward pass and an optimizer step.

They are not necessarily one-to-one.

One optimizer step may contain many backward passes¶

Suppose:

\[ A=16. \]

Then training looks like:

forward/backward 1
forward/backward 2
forward/backward 3
...
forward/backward 16
        ↓
optimizer step
        ↓
forward/backward 17
...

So if someone says:

"The model trained for 100,000 steps"

you should ask:

"Optimizer steps, or microbatch iterations?"

Those can represent very different amounts of training compute and data.

How distributed training fits in¶

Now add multiple GPUs.

Suppose four GPUs each process one microbatch:

GPU 0 → microbatch A
GPU 1 → microbatch B
GPU 2 → microbatch C
GPU 3 → microbatch D

Each GPU computes local gradients:

\[ g_0,g_1,g_2,g_3. \]

Data-parallel training then combines them, commonly using an all-reduce operation:

\[ g = \frac{g_0+g_1+g_2+g_3}{4}. \]

With gradient accumulation, we have another dimension.

Each GPU may process several microbatches before the optimizer update.

Conceptually:

             accumulation step 1
GPU 0 → batch A1
GPU 1 → batch B1
GPU 2 → batch C1
GPU 3 → batch D1

             accumulation step 2
GPU 0 → batch A2
GPU 1 → batch B2
GPU 2 → batch C2
GPU 3 → batch D2

             accumulation step 3
...

             ↓

        optimizer update

The effective batch includes data processed across both:

\[ \text{devices} \]

and:

\[ \text{accumulation steps}. \]

A communication optimization¶

There is a subtle systems opportunity here.

If gradients are synchronized across GPUs after every microbatch, then accumulation may trigger repeated expensive network communication.

But the weights are not changing yet.

So some distributed systems can delay gradient synchronization during intermediate accumulation steps.

Conceptually, instead of:

microbatch 1
→ backward
→ synchronize gradients

microbatch 2
→ backward
→ synchronize gradients

microbatch 3
→ backward
→ synchronize gradients

we would prefer, when the parallelism strategy permits it:

microbatch 1
→ backward
→ accumulate locally

microbatch 2
→ backward
→ accumulate locally

microbatch 3
→ backward
→ accumulate locally

final microbatch
→ synchronize
→ optimizer step

This reduces communication frequency.

Framework details differ, especially when gradients or parameters themselves are sharded, but the general principle is important:

\[ \boxed{\text{do not communicate more often than the algorithm requires}} \]

At large scale, network traffic can be almost as important as GPU arithmetic.

Why use a large batch at all?¶

If smaller batches fit more easily, why not simply update the model after every tiny microbatch?

Because gradients are noisy estimates of the direction we want.

Imagine the true dataset gradient is:

\[ g^*. \]

A small random batch gives an estimate:

\[ \hat g = g^* + \epsilon, \]

where \(\epsilon\) represents sampling noise.

A larger batch averages over more examples, reducing some of that noise.

Very loosely:

\[ \operatorname{Var}(\hat g) \propto \frac{1}{B}. \]

So increasing batch size can make gradient estimates more stable.

That can allow efficient use of huge parallel machines.

But there is a catch.

Bigger batches are not automatically better¶

Suppose you have a fixed dataset of:

\[ 1\text{ billion tokens}. \]

If each optimizer step uses:

\[ 1\text{ million tokens}, \]

you get:

\[ 1000 \]

parameter updates per pass through that data.

If each optimizer step instead uses:

\[ 10\text{ million tokens}, \]

you get only:

\[ 100 \]

updates.

So increasing batch size gives you more accurate gradient estimates but fewer opportunities to change the parameters for a fixed amount of data.

This creates a tradeoff:

small batch
→ noisier gradients
→ more optimizer updates

large batch
→ cleaner gradients
→ fewer optimizer updates

There is usually a useful batch-size regime rather than "as large as possible."

Batch size and learning rate interact¶

Suppose we increase global batch size dramatically.

The gradient estimate becomes less noisy, and the optimization dynamics change.

Training recipes therefore often adjust the learning rate with batch size.

A common heuristic is some form of learning-rate scaling, though the exact rule depends on the optimizer, model, training phase, and empirical recipe.

The important point is:

\[ \boxed{\text{batch size is an optimization hyperparameter, not merely a memory setting}} \]

Gradient accumulation lets us choose a large effective batch independently of what physically fits at once.

But that does not mean the chosen effective batch is automatically good.

Sequence length complicates the picture¶

For language models, "batch size = 8" is ambiguous.

Consider:

batch A:
8 sequences × 512 tokens
= 4,096 tokens

batch B:
8 sequences × 8,192 tokens
= 65,536 tokens

Both contain 8 sequences.

But batch B contains:

\[ 16\times \]

as many tokens.

Since next-token prediction creates a training target at almost every token position, token count is usually a more meaningful measure of training batch scale.

So large LLM training discussions often describe global batches in terms of:

\[ \text{tokens per optimizer step}. \]

Packing helps use those tokens efficiently¶

Suppose the training sequence length is:

\[ T=4096. \]

But the dataset contains documents of lengths:

700 tokens
1,200 tokens
350 tokens
900 tokens
...

If every document occupies its own 4,096-token training slot, enormous amounts of padding may be wasted.

Instead, training systems can pack multiple documents into a sequence.

Conceptually:

4096-token training sequence:

[document A][document B][document C][part of D]

rather than:

[document A][padding padding padding ...]

This increases the fraction of GPU work spent on real training tokens.

Packing and gradient accumulation solve different problems:

packing
→ reduce wasted token slots inside a microbatch

gradient accumulation
→ combine multiple microbatches into one optimizer batch

Both improve training efficiency.

Accumulation does not eliminate activation memory¶

Suppose even one microbatch is too large.

Gradient accumulation cannot help further if:

\[ B_{\text{micro}}=1 \]

and a single long sequence still does not fit.

Then we need other techniques.

One is activation checkpointing, which we will examine separately.

Instead of storing every forward activation, checkpointing stores only selected activations and recomputes others during backward.

So:

gradient accumulation
→ reduce how many examples are resident together

activation checkpointing
→ reduce how much intermediate state is stored per example

These techniques are complementary.

Accumulation also does not reduce optimizer-state size¶

Recall yesterday's AdamW memory:

parameters
gradients
first moments
second moments
possibly master weights

Gradient accumulation mainly attacks activation memory associated with the global batch.

It does not magically make Adam's:

\[ m_t \]

or:

\[ v_t \]

smaller.

Those are still parameter-sized tensors.

That is why large training systems combine multiple techniques:

gradient accumulation
activation checkpointing
optimizer sharding
parameter sharding
tensor parallelism
pipeline parallelism
mixed precision

Each attacks a different resource constraint.

The deeper mental model: separate physical batch from statistical batch¶

This is the most useful idea to remember.

There are really two questions:

How much data can the hardware process at once?

and:

How much data should contribute to one optimization decision?

Those do not have to have the same answer.

The first is the physical microbatch.

The second is the effective global batch.

Gradient accumulation creates the bridge:

\[ \boxed{ B_{\text{global}} = B_{\text{micro}} \times N_{\text{devices}} \times N_{\text{accumulation}} } \]

or, more usefully for LLMs:

\[ \boxed{ T_{\text{global}} = T_{\text{micro per device}} \times N_{\text{devices}} \times N_{\text{accumulation}} } \]

where \(T\) counts training tokens.

This is a recurring pattern in large-model engineering.

The mathematical object we want may be enormous:

a multi-million-token gradient estimate

while the hardware can only hold a small piece at any instant.

So we decompose the mathematical operation into smaller physical operations while preserving the intended result.

That same philosophy appears throughout LLM systems:

huge model
→ shard parameters

huge optimizer
→ shard optimizer state

huge batch
→ accumulate microbatches

huge activation footprint
→ recompute activations

huge inference workload
→ schedule token batches

The model may look like one giant computation on paper.

In practice, large-scale AI is largely the art of deciding which pieces must exist simultaneously and which pieces can be spread across time or machines.

Gradient accumulation is one of the cleanest examples: the optimizer can learn from millions of tokens at once, even though the GPU never sees those millions of tokens at once.