Multi-GPU calculation of gradients for accelerated training

A half-million-token optimizer step

In Karpathy’s GPT-2 walkthrough, he takes the roughly half-million-token batch target from the 125M-parameter row of the GPT-3 training table because that model is close in size to GPT-2 124M. In code he chooses the nearby power-of-two value total_batch_size = 524288, or 2**19, tokens per optimizer update. This is a chosen optimizer batch size, not a limit imposed by the GPT-2 architecture.

A micro-batch is the data handled by one GPU in one forward and backward call. That call is one micro-step. The global batch is all data that contributes to one optimizer update, across every micro-step and GPU.

Each sequence contains T = 1024 tokens, so 524,288 tokens means 512 sequences. Sending all 512 through one GPU at once would require it to retain the intermediate tensor elements, called activations, for all 512 sequences until backward finishes. That does not fit the GPU used in the walkthrough.

Karpathy instead uses B = 16 sequences per micro-batch:

One GPU therefore needs 32 serial micro-steps to cover the global batch. Each backward pass adds to the parameter gradients, and the optimizer waits until all 32 have finished. Gradient accumulation makes the batch fit by doing the work in smaller pieces. On one GPU, those 32 micro-steps still run sequentially.

Eight GPUs shorten the serial schedule

The walkthrough starts one process per GPU with torchrun, then wraps each process’s model in PyTorch DistributedDataParallel, or DDP. Each process has a rank. Every rank holds a model replica, and DataLoaderLite gives it a different 16-sequence micro-batch.

During one parallel wave, eight ranks process

at the same time. A wave is one micro-step running concurrently on all eight ranks. Four waves cover the global batch:

The one-GPU schedule has 32 serial micro-steps. The eight-GPU schedule has four serial waves. If the computation scaled perfectly, this part would take one eighth as long. In practice, exchanging gradients and other overhead reduce the speedup.

Each rank accumulates gradients from its four micro-batches. On the final backward pass, DDP averages the gradients across all eight ranks. Every model replica then applies the same optimizer update.

Figure 1. More GPUs reduce the serial work needed for the same global batch. D2 source · SVG

In the eight-GPU run, grad_accum_steps = 4: each rank executes four micro-steps.

grad_accum_steps = total_batch_size // (B * T * ddp_world_size)
# 524288 // (16 * 1024 * 8) = 4

Each rank therefore divides its losses by four:

loss = loss / grad_accum_steps

The factor of eight is absent because DDP supplies the separate average across ranks.

Cross-entropy averages within one micro-batch

PyTorch receives B * T next-token predictions:

loss = F.cross_entropy(
    logits.view(-1, logits.size(-1)),
    targets.view(-1),
)

The code uses every target token and gives them equal weight. With the default reduction="mean", PyTorch divides the sum of the token losses by B * T. Targets from later micro-steps are not arguments to this call, so they are absent from its denominator.

Gradient accumulation averages across micro-steps

In this GPT-2 loop, every micro-step on every GPU has the same B * T valid target tokens. Under that condition, averaging the micro-batch means gives the same global mean as processing all targets together.

Here is a tiny version with two micro-steps and two token predictions per micro-batch. Suppose the individual losses are 2, 4, 6, 8. Processing all four together would give

Split them into two micro-batches and cross-entropy returns:

micro-step 1: mean(2, 4) = 3
micro-step 2: mean(6, 8) = 7

Calling .backward() twice adds into each parameter’s .grad. Without another division, this gives the gradient of $L_1 + L_2$. We instead want the gradient of their mean:

For these values:

The identity holds for the loss functions themselves, so their gradients match.

Figure 2. Cross-entropy supplies the mean inside each micro-batch. Dividing by two supplies the missing mean across two micro-steps. D2 source · SVG

DDP averages gradients across processes

Now use two GPUs, two micro-steps per rank, and two token predictions per micro-batch.

• GPU 0 sees means 3 and 7, so its scaled local objective is 3/2 + 7/2 = 5.
• GPU 1 sees losses 10, 12, 14, 16. Its micro-batch means are 11 and 15, so its scaled local objective is 11/2 + 15/2 = 13.

Call the two local parameter gradients $g_0$ and $g_1$. DDP gives every rank their average, $(g_0 + g_1) / 2$. Using the gradient identity above, that is the gradient of

This is exactly the whole eight-token mean:

Figure 3. Each rank averages across its own micro-steps. DDP performs the final gradient average across ranks. D2 source · SVG

Do not divide the loss used for .backward() by the number of GPUs. DDP already divides the synchronized parameter gradients by the number of ranks. Dividing by two again in this example would halve the intended gradient.

DDP can skip the gradient all-reduce for the first grad_accum_steps - 1 micro-steps. The supported no_sync() context must wrap both forward and backward. The video instead assigned model.require_backward_grad_sync after forward; a later fix moved that assignment before forward because DDP also consults the flag during forward.

References

• Andrej Karpathy, gradient accumulation and DDP gradient averaging in Let’s reproduce GPT-2 (124M).
• Tom B. Brown et al., Language Models are Few-Shot Learners, Table 2 (the 125M-parameter model uses a 0.5M-token batch).
• Karpathy, video-era cross-entropy call and gradient-accumulation/DDP loop.
• PyTorch, torch.nn.functional.cross_entropy, whose default reduction is mean.
• PyTorch, DistributedDataParallel and the no_sync() gradient-accumulation guidance.