insideout Get Pro

11 · Training across many GPUs

Pro

Training and Tuning, lesson 11

Training across many GPUs

How one training run spreads across many GPUs.

Big training runs spread over dozens or thousands of GPUs, for speed and because the training state won't fit on one. This lesson works out how the batch, the learning rate and the memory get split, and what a checkpoint needs so a run survives a crash. You need the training loop from "Pretraining" and the memory count from "Adapters and LoRA".

By the end, you can:

  1. Work out a run's effective batch from its GPUs, per-GPU batch and accumulation steps.
  2. Scale the learning rate when the batch grows, and warm it up.
  3. Count training memory per GPU, with and without sharding.
  4. List what a checkpoint needs for a run to carry on exactly, and size it.

Part 1. Split the batch, average the gradients

Work out a run's effective batch from its GPUs, per-GPU batch and accumulation steps.

A training step averages the loss over a batch, then takes one step downhill. A big batch won't fit on one GPU, and one GPU would be slow anyway. So give every GPU a full copy of the model and a different slice of the batch.

Example 1.1

Eight GPUs, one update

Each of 8 GPUs has room for 4 sequences of 4,096 tokens at a time. Each GPU runs 8 of these micro-batches and adds up their gradients before the weights change. How many sequences and tokens go into one update, and how often do the GPUs swap gradients?

  1. One GPU, one pass4 sequences

    As many as fit in its memory at once, with their activations.

  2. Accumulate4 × 8 = 32 sequences per GPU

    The gradients from each micro-batch are added up, and nothing is updated in between.

  3. Add the GPUs32 × 8 = 256 sequences

    Every GPU saw a different 32.

  4. Count the tokens256 × 4,096 = 1,048,576

    About a million tokens per update.

  5. Swap once1 all-reduce per update

    The GPUs average gradients only after the 8th micro-batch.

Answer256 sequences, about a million tokens, and one swap per update.

The batch loss is an average, and the gradient of an average is the average of the gradients. With equal slices, averaging 8 GPUs' gradients gives exactly the gradient of the whole 256-sequence batch. Every copy then makes the same update, so the copies stay identical.

one batch, split four waysslice 1GPU 1full modelg₁slice 2GPU 2full modelg₂slice 3GPU 3full modelg₃slice 4GPU 4full modelg₄all-reduce: (g₁ + g₂ + g₃ + g₄) ÷ 4same updatesame updatesame updatesame update
one batch, split four waysslice 1GPU 1full modelg₁slice 2GPU 2full modelg₂slice 3GPU 3full modelg₃slice 4GPU 4full modelg₄all-reduce: the average(g₁ + g₂ + g₃ + g₄) ÷ 4every GPU makes the same update
Figure 1Four GPUs, for short. Each holds the whole model and computes a gradient on its own slice. The all-reduce leaves every GPU holding the same average, so every copy takes the same step.

Data parallel, all-reduce and gradient accumulation

Data parallel training gives every GPU a full copy of the model and a different slice of each batch. An all-reduce sums a value across GPUs and hands every GPU the result. Gradient accumulation adds up gradients over several micro-batches before one update.

effective batch = per-GPU batch × GPUs × accumulation steps

The average must weigh every token equally. Averaging micro-batches' mean losses doesn't: 100 tokens at loss 2.0 and 300 at 1.0 give 1.5, but the per-token loss is 1.25. Hugging Face fixed this in its Trainer in October 2024. Sum the losses and divide by the total token count.

Your turn

A run uses 64 GPUs. Each takes 2 sequences of 4,096 tokens per micro-batch and accumulates 16 micro-batches per update. How many tokens go into one update?

64 × 2 × 16 = 2,048 sequences, and 2,048 × 4,096 = 8,388,608 tokens. 524,288 forgets the accumulation steps, 131,072 is one GPU's share, and 2,048 counts sequences, not tokens.

Part 2. A bigger batch takes bigger steps

Scale the learning rate when the batch grows, and warm it up.

Keep going with Pro

You've done the first part. Pro opens the rest of this lesson, and every other one.

Still to come
  1. A bigger batch takes bigger steps
  2. When the model doesn't fit
  3. Checkpoints that resume exactly
See Pro plansI have Pro: sign in
Next lessonDistillationTrain a small model on a big one's outputs. Start lesson 12

Training and Tuning overviewSpotted a mistake? Tell us.