GPU & Efficient Training

GPU training moves both model parameters and each batch's tensors to the same accelerator device. Mixed precision uses lower-precision arithmetic where it is safe and maintains critical values or scaling in higher precision, often reducing memory use and increasing throughput on modern hardware. Gradient accumulation sums appropriately scaled gradients across several micro-batches before one optimiser step, emulating a larger effective batch when activations would not fit together in memory. Data loading, host-to-device transfer, sequence padding, and activation checkpointing can dominate performance as much as the model's arithmetic.

Performance is often limited by input loading, memory bandwidth, or activation storage rather than raw FLOPs. Measuring those bottlenecks is more useful than guessing: profile step time, accelerator utilisation, memory peaks, and data-loader wait time before changing model design. Efficiency work must preserve numerical and training semantics; a faster run with a changed effective batch or unstable reduced precision is not a fair comparison.

A CPU has a few dozen powerful cores optimised for low-latency execution of branchy, dependent instructions. A GPU has thousands of simple cores and very wide memory bandwidth, optimised for throughput on the same operation applied to many independent elements. Deep learning is almost entirely the second shape: a matrix multiply is millions of independent multiply-accumulates with no branches, and an activation is the same scalar function applied element by element.

This is the NumPy vectorisation lesson taken to hardware. There, the win came from replacing a Python loop with one compiled operation over a whole array; here, the same restructuring lets thousands of cores work at once. It also explains the shape of every optimisation in this module. A GPU is fast at *bulk uniform arithmetic*, so anything that breaks uniformity — a Python loop over batch elements, per-example branching, tiny kernels, a synchronising .item() inside the inner loop — destroys the advantage. And because bandwidth is finite, many deep-learning operations are memory-bound rather than compute-bound: they are waiting for numbers to arrive, not for arithmetic to finish, which is why lower precision speeds things up even when it does not reduce the operation count.

Mixed precision runs most of the forward and backward pass in 16-bit floating point while keeping a 32-bit master copy of the weights that the optimiser actually updates. With float16 the exponent range is narrow enough that small gradients underflow to zero, so a loss scaler multiplies the loss by a large factor before backward() and divides the gradients afterwards. bfloat16 has the same exponent range as float32 and fewer mantissa bits, so it needs no loss scaling — which is why it has become the default on hardware that supports it.

It buys two things at once. Memory: activations and gradients halve, which is usually the binding constraint. Speed: modern accelerators have dedicated tensor cores that do 16-bit matrix multiplies several times faster than 32-bit ones, and halving the bytes moved also relieves the memory-bandwidth limit. Accuracy holds up because the parts of training that are numerically delicate — accumulating a weight update, reducing a sum, normalising — are kept in float32 by the autocast rules, while the parts that are robust to a few lost mantissa bits, chiefly the matrix multiplies, are not. The win is close to free, which is why mixed precision is a default rather than a tuning option.

Gradient accumulation decouples the batch size the optimiser sees from the batch size that fits in memory. Run N micro-batches, dividing each loss by N and calling backward() on each so gradients sum in .grad, then call optimizer.step() once. The resulting update is mathematically the same as one step on a batch N times larger, because the mean of means over equal-sized micro-batches is the mean over the whole batch.

It matters because batch size is not a free parameter. Large-batch training needs a matched learning rate, interacts with the warmup schedule, and changes the gradient noise that some models depend on; a result reported at batch 512 cannot be reproduced at batch 32 by changing nothing else. Accumulation is what lets you keep the *training recipe* fixed while adapting to whatever memory you actually have — the same run on one 24 GB GPU and on eight 80 GB ones, if the effective batch matches. The caveats are that it costs proportional time (you do all the same work, just serially) and that the division must be exact: unequal final micro-batches weight their examples differently unless you normalise by token or sample count instead of by N.

Data parallelism replicates the whole model on every GPU, gives each a different slice of the batch, and averages the gradients across devices before the update — so every replica stays identical. Model parallelism instead splits the *model*: tensor parallelism shards individual weight matrices across devices, and pipeline parallelism assigns different layers to different devices and streams micro-batches through them. Fully-sharded approaches combine the ideas by sharding parameters, gradients and optimiser state across the data-parallel group and gathering each layer's weights only while it is in use.

The choice is forced by a single question: does one copy of the model, its gradients and its optimiser state fit on one device? If yes, use data parallelism — it is simpler, it scales nearly linearly, and the only communication is one gradient all-reduce per step. If no, you have no option but to split the model, and you accept the cost: tensor parallelism communicates *within* every layer, so it needs very fast interconnect, and pipeline parallelism leaves devices idle at the start and end of each batch ('bubbles') unless micro-batching keeps the pipe full. Reaching for model parallelism before you need it buys complexity and a slower run.

Training memory has four components and they scale differently. Parameters are fixed by the model. Gradients are one value per trainable parameter. Optimiser state is the largest surprise: AdamW keeps two float32 moment estimates per parameter, so it costs roughly twice the model in fp32. Activations are everything the forward pass stored for the backward pass, and unlike the first three they scale with *batch size and sequence length* — which is why the same model trains fine at batch 8 and dies at batch 16.

This breakdown is the most practically valuable table in the module, because it tells you which lever to pull. If activations dominate, reduce batch size, shorten sequences, or use gradient checkpointing — changing the optimiser will not help. If optimiser state dominates, switch to 8-bit Adam or shard it with ZeRO/FSDP — reducing batch size will not help. Without the breakdown, out-of-memory debugging is guesswork; with it, the fix is usually obvious from which row is biggest.

Normally the forward pass stores every intermediate activation because the backward pass needs them to compute local derivatives. Gradient checkpointing stores only a few — the boundaries between segments of the network — and discards the rest. During the backward pass, each segment's activations are recomputed from its stored input just before they are needed, then thrown away again. It is an explicit trade of compute for memory.

The trade is unusually favourable, which is why it is a standard tool rather than a last resort. Storing every activation costs memory proportional to depth; storing only segment boundaries costs roughly the square root of that, while the extra compute is one additional forward pass, about 30% more time. For a deep network that is a 60 to 70% cut in activation memory for a third more time — and if the alternative is not being able to train at all, or training at a batch size so small the gradient is too noisy, that is an easy trade. It is also composable: checkpointing plus mixed precision plus accumulation are three independent levers on three different terms of the budget.

torch.cuda.OutOfMemoryError reports how much was requested, how much is free, and how much PyTorch has reserved. Read all three: a request for 2 MB failing with 8 GB reserved and 50 MB free means fragmentation, not a genuine shortage. The error also names the moment of failure, which is informative — failing in the backward pass points at activations, failing at optimizer.step() on the first update points at optimiser state being allocated lazily.

The reason to have a checklist is that the instinctive fix — halve the batch size — often does not work, because the batch was never the biggest term. Work down the list instead: is something accumulating across steps (a loss tensor in a list, a metric holding a graph, a missing zero_grad)? Is the validation loop missing no_grad, so it retains activations at inference batch sizes? Did memory grow monotonically over epochs, which means a leak rather than a budget? Only then reduce batch size, and if that is not enough turn to the structural levers: checkpointing, an 8-bit optimiser, sharding, or LoRA.