12.5.7 — Training at Scale

Cartwise's photo team decided to go beyond the frozen network of 12.5.5 and fine-tune a larger image model on two million seller photos. They rented a cloud machine with a powerful GPU at a discount, using a spot instance: spare capacity the cloud provider sells cheaply on the condition that it can take the machine back at short notice. The job was expected to take three days. At hour thirty, the provider reclaimed the machine. The script saved the model only at the very end, so thirty hours of training, and the money spent on them, were simply gone.

The second attempt failed differently. They moved to a bigger model, and the job crashed in the first second with CUDA out of memory, although the model's weights were only a fraction of the GPU's memory.

Both failures come from the same gap: knowing how training works, which Chapters 12.5.1 to 12.5.4 covered, is not the same as knowing what it costs in memory, time and failure risk once models and datasets get large. This page covers that: where the memory goes, how to make a model fit, how training is spread across many GPUs, how to search for good settings, and how to keep a long run safe and every experiment reproducible.

1. The hardware, and the real limit

12.2.2 explained why GPUs run neural networks well: a layer is thousands of independent dot products, and a GPU has thousands of simple cores that do the same operation on many numbers at once. Two more facts about modern GPUs matter in practice.

They are often limited by how fast they can read memory, not by how fast they can multiply. A data-centre GPU can move data from its memory to its cores at over 3,000 gigabytes per second, many times faster than a typical server processor, because neural networks constantly stream huge weight matrices through the cores. Modern GPUs also contain tensor cores, circuits built specifically to multiply small blocks of a matrix in a single step.

Their memory is small and fixed. A top data-centre GPU has around 80 gigabytes of its own memory, called VRAM. Everything the training step needs must fit there at once, or it does not run at all. That is the limit you hit first, and it is the reason for most of the techniques on this page.

2. Where the memory goes during training

The out of memory crash is easy to understand once you count what training actually stores. For every parameter, training with the Adam optimiser from 12.5.3 keeps:

• The weight itself. 4 bytes in standard 32-bit floating point.
• Its gradient, filled in by backpropagation. Another 4 bytes.
• Adam's two running averages, m and s. Another 8 bytes.

That is about 16 bytes per parameter, four times the size of the weights alone. On top of that come the activations: every intermediate value from the forward pass, kept for the backward pass, as 12.5.2 explained. Activations grow with the batch size and with the size of each input, and for large inputs they are often the biggest item of all.

Model size	Weights alone (32-bit)	Training state (≈16 bytes/param)	Fits one 80 GB GPU?
25 million (a ResNet-50)	0.1 GB	0.4 GB	easily, plus activations
1 billion	4 GB	16 GB	yes, with room for activations
7 billion	28 GB	112 GB	no, before any activations

So a 7-billion-parameter model, whose 32-bit weights alone are 28 gigabytes and look comfortable on an 80-gigabyte GPU, needs about 112 gigabytes just for its training state. That is the crash. Running the same model, without training it, needs only the weights, and in 16-bit format only 14 gigabytes. Training needs several times the memory of inference, and that single fact explains why fine-tuning large models needed special tricks, the subject of Chapter 12.11.2.

3. Making it fit on one GPU

Smaller numbers: mixed precision

Chapter 1.4 explained how a floating-point number stores a sign, an exponent (which sets how large or small the number can be) and a mantissa (which sets how many digits of precision it has). Neural networks turn out not to need 32 bits for most of their arithmetic. Three formats matter:

Format	Bits: sign / exponent / mantissa	Largest value	Notes
float32	1 / 8 / 23	about 3 \times 10^{38}	the safe default
float16	1 / 5 / 10	65,504	precise but narrow range
bfloat16	1 / 8 / 7	about 3 \times 10^{38}	same range as float32, less precision

The important comparison is float16 against bfloat16. float16 keeps more digits but only 5 exponent bits, so its largest value is 65,504, and its smallest normal values are about 6 \times 10^{-5}. During training, gradients can be far smaller than that and round to zero, or activations can exceed 65,504 and overflow to infinity. Early 16-bit training needed a workaround called loss scaling to dodge this. bfloat16, the "brain floating point" format designed at Google Brain, keeps float32's full 8 exponent bits and gives up mantissa instead. It has the same range as float32 and so almost never overflows or underflows. For training, range matters more than the last few digits of precision, so bfloat16 is the format of choice wherever the hardware supports it.

In mixed precision training, most of the forward and backward arithmetic runs in 16 bits, roughly halving activation memory and running much faster on tensor cores, while a few sensitive operations, and a master copy of the weights, stay in 32 bits. PyTorch makes it a one-line change around the forward pass:

for xb, yb in train_dl:
    opt.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):   # (1)
        loss = loss_fn(model(xb), yb)
    loss.backward()
    opt.step()

Line (1) tells PyTorch to run the operations inside the block in bfloat16 where that is safe, and in 32 bits where it is not, choosing automatically per operation.

Bigger batches without the memory: gradient accumulation

A batch of 256 photos may not fit, while a batch of 32 does. Because PyTorch adds new gradients to the stored ones, the behaviour that made zero_grad necessary in 12.5.3, you can run eight batches of 32, let their gradients add up, and only then take one step. Mathematically this is the same as one batch of 256, at the memory cost of 32:

ACCUM = 8
opt.zero_grad()
for i, (xb, yb) in enumerate(train_dl):
    loss = loss_fn(model(xb), yb) / ACCUM    # (1) average over the 8 mini-batches
    loss.backward()                          # (2) gradients add up in .grad
    if (i + 1) % ACCUM == 0:
        opt.step()                           # (3) one update per 8 mini-batches
        opt.zero_grad()

Line (1) divides each loss by 8, so the summed gradient equals the average over 256 examples rather than eight times it. Line (2) backpropagates without clearing, so the gradients accumulate. Line (3) takes one step every eighth batch, then clears. It is called gradient accumulation.

Store less, compute more: gradient checkpointing

The activations kept for backpropagation can be traded for time. Gradient checkpointing keeps only the activations at a few chosen points in the network and throws the rest away; during the backward pass, it recomputes each missing stretch from the nearest kept point when it is needed. That costs roughly one extra forward pass, often around a third more training time, and can cut activation memory several times over. It is the standard move when a model almost fits.

4. Spreading training across many GPUs

When one GPU is too slow, or the model will not fit on one at all, training is spread across many. There are a few basic ways to divide the work, and large training runs combine them.

The two basic ways to spread training. Data parallelism (left) copies the whole model onto every GPU and splits the batch. Model parallelism (right) splits the model itself: by groups of layers (pipeline) or by cutting individual matrices (tensor).

Data parallelism is the simplest and most common. Every GPU holds a complete copy of the model. Each batch is split, and each GPU runs the forward and backward pass on its share. Then the GPUs average their gradients, so every copy takes exactly the same step and the copies stay identical. The averaging is called an all-reduce: every GPU ends up holding the sum of everyone's values, computed with a communication pattern that keeps every link busy instead of funnelling everything through one machine. With four GPUs, you process four times as many examples per step, minus the time spent communicating. PyTorch's version is called DistributedDataParallel. Its limit is in the figure: the whole model, with its training state, must still fit on each GPU.

Sharding the training state. Notice how wasteful plain data parallelism is: every GPU holds a full copy of the weights, the gradients and Adam's averages, so with eight GPUs, the 112 gigabytes of training state for a 7-billion-parameter model are stored eight times. The fix, introduced by Microsoft's DeepSpeed team under the name ZeRO (zero redundancy optimiser) and built into PyTorch as FSDP (fully sharded data parallel), gives each GPU only one eighth of the weights, gradients and optimiser state, and passes the pieces around just in time for each layer's computation. Eight GPUs then hold about 112 / 8 = 14 gigabytes of state each, and the model fits.

Pipeline parallelism puts different groups of layers on different GPUs, and the data flows from one to the next like an assembly line. The problem is idle time: GPU 2 cannot start until GPU 1 has finished, and on the backward pass the order reverses. That idle time is called the bubble. Splitting each batch into several small micro-batches, so that GPU 1 can start the second micro-batch while GPU 2 works on the first, keeps everyone busier.

Tensor parallelism splits individual matrix multiplications across GPUs, for example the left and right halves of a weight matrix, as in the figure. The GPUs must exchange partial results inside every layer, so this only works over very fast connections, typically between GPUs in the same server.

The largest language models are trained with all of these at once across thousands of GPUs, which is sometimes called 3D parallelism. As an application engineer you will rarely set this up yourself, but knowing the shapes tells you why a fine-tuning job needs the hardware it asks for, and what "this model needs eight GPUs to train" actually means.

5. Transfer learning and fine-tuning

Almost nobody trains a large model from scratch, and the reason is economic. A model pre-trained on a huge general dataset has already learned the low-level structure of its kind of data: edges and textures for images, grammar and word meanings for text. Adapting it to a specific task needs a tiny fraction of the data and compute. 12.5.5 did exactly this for photos. There are three levels, from cheapest to most thorough:

1. Feature extraction. Freeze the whole pre-trained network and train only a new final layer. It works with a few hundred labelled examples and trains in minutes, but the network's internal features cannot adapt.
2. Full fine-tuning. Unfreeze some or all layers and keep training on your data with a much lower learning rate, typically 10 to 100 times lower than for training from scratch. A high learning rate makes large changes to the pre-trained weights and wipes out what they knew, a failure called catastrophic forgetting. A common refinement is to unfreeze gradually, from the last layers backwards.
3. Parameter-efficient fine-tuning. Keep the original weights frozen and train a small number of added parameters alongside them. Because the frozen weights need no gradients or optimiser state, the 16-bytes-per-parameter bill applies only to the small added part. This is how large language models are adapted today, and Chapter 12.11.2 covers its most popular form, LoRA.

6. Searching for good settings

The learning rate, batch size, weight decay, dropout rate and number of layers are hyperparameters from 12.3.3, chosen on validation data. With several of them, trying combinations by hand is slow and biased towards whatever you tried first. Three systematic approaches:

Grid search tries every combination from a list of values for each setting. With three values for each of two settings, that is 9 runs.

Random search tries the same number of runs with random combinations from the ranges. It sounds less careful, and it usually works better. The reason, set out by James Bergstra and Yoshua Bengio in 2012, is that usually only one or two settings matter much, and you do not know in advance which ones.

Why random search usually beats grid search. If only the learning rate really matters (the purple curve below each square), the grid's nine runs test just three distinct learning rates, while random search's nine runs test nine, and so are more likely to land near the best one.

With a 3 by 3 grid, if dropout barely matters, three of your nine runs are wasted repeats of each learning rate, and you have effectively tested only three learning rates. Nine random runs test nine different learning rates. Searching on a logarithmic scale is also important for learning rates: sample the exponent, so that 0.0001, 0.001 and 0.01 are equally likely, rather than sampling between 0.0001 and 0.01 evenly, which would almost never try anything below 0.001.

Smarter search builds on this. Bayesian optimisation fits a small model of "which settings give which validation score" from the runs so far and chooses the next run where it expects the biggest improvement. Successive halving (and its extension, Hyperband) starts many runs, stops the worst half after a short time, gives the survivors longer, and repeats, so most of the compute goes to promising settings. Libraries such as Optuna implement both.

7. Keeping every experiment: tracking and checkpoints

Experiment tracking

A month into the project, the team had run 140 training jobs, and nobody could say for certain which settings produced the model in production. A spreadsheet had been started and abandoned. This is normal, and the fix is a tool that records every run automatically.

Experiment tracking tools log, for each run, its settings, its code version, its data version, the hardware it ran on, and every metric over time, and show them on a web dashboard where runs can be compared. Weights & Biases, usually called W&B (its website is wandb.ai), is the most widely used hosted service; MLflow is the common open-source one that you can run on your own servers. Adding W&B to the training loop from 12.5.3 takes a few lines:

import wandb

run = wandb.init(project="cartwise-photo-checker",                # (1)
                 config={"lr": 1e-3, "batch_size": 64, "arch": "resnet18",
                         "data_snapshot": "photos-2026-10-01"})
for epoch in range(10):
    train_one_epoch(model)
    val_loss, val_acc = evaluate(model, val_dl)
    wandb.log({"epoch": epoch, "val_loss": val_loss, "val_acc": val_acc})   # (2)
run.finish()

Line (1) starts a run and records its settings, including the data snapshot identifier from 12.4.4; the library also records the Git commit and machine details by itself. Line (2) sends metrics after each epoch, which appear as live charts. When a model reaches production, its run page answers "what exactly made this?" in one click, and comparing two runs side by side shows which change helped.

Checkpoints

The thirty lost hours had one cause: nothing was saved until the end. A checkpoint is a saved snapshot of the training state, written regularly, from which training can resume exactly as if it had never stopped. Resuming exactly needs more than the weights:

def save_checkpoint(path, model, opt, scheduler, epoch, step):
    torch.save({
        "model": model.state_dict(),          # (1) the weights
        "opt": opt.state_dict(),              # (2) Adam's running averages
        "sched": scheduler.state_dict(),      # (3) where the learning rate schedule is
        "epoch": epoch, "step": step,         # (4) where we are in the data
        "rng": torch.get_rng_state(),         # (5) the random number generator
    }, path)

Line (1) saves the weights. Line (2) saves the optimiser's state: without Adam's averages, resuming would behave like a fresh start with a cold optimiser and could jolt the weights. Line (3) saves the schedule's position, so the learning rate continues from where it was rather than restarting the warmup. Line (4) records how far through the data training had got, and line (5) the random state, so that shuffling and dropout continue the same sequence. Saving every half hour turns a reclaimed spot instance from a disaster into a ten-minute delay, which is why spot instances, often far cheaper than the normal price, are the standard way to pay for long training runs. Checkpoints also give you a history: if the model gets worse late in training, you can go back to an earlier one, as early stopping in 12.5.4 did automatically.

Recall

• GPU memory is the first limit. Training with Adam stores about 16 bytes per parameter (weights, gradients, two averages) plus activations, so a 7-billion-parameter model needs about 112 GB of training state, while inference in 16 bits needs about 14 GB.
• bfloat16 keeps float32's 8 exponent bits, so it rarely overflows; mixed precision runs most arithmetic in 16 bits with torch.autocast. Gradient accumulation sums gradients over small batches to act like a big one; gradient checkpointing recomputes activations to save memory.
• Data parallelism copies the model and averages gradients with an all-reduce; sharding (ZeRO, FSDP) splits the training state across GPUs; pipeline parallelism splits by layer (the bubble); tensor parallelism splits matrices and needs fast links.
• Fine-tune pre-trained models at a 10 to 100 times lower learning rate to avoid catastrophic forgetting; parameter-efficient methods avoid the training-state bill for frozen weights.
• Random search beats grid search when few settings matter; sample learning rates on a log scale. Experiment tracking (W&B, MLflow) records settings, code, data and metrics for every run. A checkpoint saves weights, optimiser and scheduler state, data position and random state, making cheap spot instances safe.

Self-test: Why does a 7-billion-parameter model whose weights take 28 GB not train on an 80 GB GPU? With eight GPUs and full sharding, about how much training state does each hold? What goes wrong if a checkpoint saves the weights but not the optimiser state?

Next: 12.6.1 What an Embedding Is — the 512 numbers a photo network produces just before its classifier are a compressed description of the photo; that idea, turning anything into a short vector that captures its meaning, powers search, recommendations and every LLM, and it is where Chapter 12.6 begins.