Module 7 — GPU training and mixed precision
The training loop from module 5 runs on a laptop CPU and reaches around 91 % accuracy on Fashion-MNIST after a couple of minutes per epoch. Real projects run on GPUs, sometimes for days, and squeezing time and cost out of them changes what is feasible in a research or product cycle. This module lays out the two orthogonal levers PyTorch gives you: moving the computation to a GPU, and switching part of it to lower-precision arithmetic. The code additions are small, the wall-clock impact is large, and both come with specific traps worth learning before you deploy them at scale.
The Fashion-MNIST classifier does not need a GPU to train, but it is small enough that every diagnostic here can be reproduced on Colab's free tier in minutes. We keep the same MLP and the AdamW plus cosine schedule from module 6.
Devices in PyTorch: a single mental model
PyTorch tensors live on a device, and every operation requires all its inputs to share it. There are three you might meet: cpu, cuda:0 (an NVIDIA GPU with a driver and CUDA toolkit) and mps (Apple Silicon's Metal Performance Shaders). The rest of this module uses cuda:0 since it covers the overwhelming majority of practical deep learning.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = FashionMLP().to(device)
.to(device) moves the model in place and returns the same object; assigning back is idiomatic but redundant. The move copies every parameter and buffer once, at model-construction time. From then on, only batches travel between host and device.
The training loop gains two .to(device) calls, one per batch:
for x, y in train_loader:
x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
optimizer.zero_grad()
logits = model(x)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
non_blocking=True allows the transfer to overlap with computation, provided the source memory is pinned. That is exactly what pin_memory=True on the DataLoader (module 4) enables. Without pinned memory, non_blocking is silently ignored.
Transfers are the bottleneck you did not know you had
A common surprise on the first GPU run is that training is barely faster than on the CPU, sometimes even slower. The culprit is almost never the compute — it is the transfers.
Every batch travels from host RAM to GPU VRAM over PCIe. On typical hardware that link peaks around 15 GB/s in one direction. A batch of 256 Fashion-MNIST images at 28-by-28 float32 is 800 KB, which is negligible. But a batch of 256 ImageNet-style images at 224-by-224 float32 is 150 MB, and the transfer alone costs several milliseconds — the same order as the actual training step.
Three cheap fixes cover most cases. First, pin_memory=True plus non_blocking=True, together, so transfer overlaps compute. Second, num_workers set high enough that the next batch is ready before the GPU asks for it. Third, transfer the raw uint8 image and cast on the GPU: dividing the transferred volume by four is worth the small kernel launch on the target device.
.to(device) is asynchronous — measuring is not straightforwardCUDA operations queue asynchronously. time.time() around a training step measures the queue-append time, not the compute time. Correct benchmarking requires torch.cuda.synchronize() before reading the clock, or torch.cuda.Event measured on the stream. Skipping this yields numbers that look great and are entirely fictional.
Mixed precision: half the bits, most of the accuracy
The default dtype is float32, 32 bits per number. Modern GPUs have specialised units — Tensor Cores on NVIDIA — that operate on float16 (or bfloat16) at two to eight times the throughput. Mixed precision casts most of the forward and backward pass to 16 bits while keeping the master copy of the weights in 32 bits, capturing most of the speed-up with negligible accuracy loss.
PyTorch exposes this as one context manager and one gradient scaler.
scaler = torch.cuda.amp.GradScaler()
for x, y in train_loader:
x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
optimizer.zero_grad()
with torch.cuda.amp.autocast(dtype=torch.float16):
logits = model(x)
loss = criterion(logits, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
Inside the autocast block, PyTorch chooses per operation whether float16 is safe (matrix multiplications, convolutions) or whether float32 must be kept (reductions, softmax, loss). Everything outside the block stays in float32. Weights, optimiser state and gradients accumulate in float32: that is the "mixed" part.
GradScaler: what it prevents
float16 has a much smaller range than float32. Small gradient values — common toward the end of training or on deep networks — underflow to zero, silently disabling learning on those parameters. GradScaler multiplies the loss by a large factor before backward, which shifts every gradient into the representable range, then divides them back before the optimiser applies them.
scaler.scale(loss).backward() scales, scaler.step(optimizer) unscales and steps (or skips the step if any gradient is inf or nan), scaler.update() adapts the scale factor for the next iteration. All three calls are mandatory. Removing scaler.update() freezes the scale and eventually reproduces the underflow it was meant to prevent.
bfloat16 versus float16
On A100, H100 and newer, bfloat16 is available. It has the same range as float32 and less precision than float16. Range is the trickier property to trade off, so bfloat16 usually trains stably without a GradScaler:
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
logits = model(x)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
If the hardware supports it, bfloat16 is simpler than float16 for the same speed-up. On older or consumer GPUs where only float16 is available, keep GradScaler.
torch.compile: a preview
Introduced with PyTorch 2, torch.compile(model) traces the forward pass, fuses operations and generates optimised kernels. For many models the speed-up is significant — 20 to 60 % — for one line of code:
model = torch.compile(model)
The first iteration is slower than usual because compilation happens then; subsequent iterations are faster. Compilation is invalidated when the input shapes change, so a stable batch_size matters more than usual. torch.compile composes well with autocast and GradScaler, and there is no reason not to try it on any model that will train for more than a few minutes.
Measuring throughput properly
The right unit to track is examples per second, not seconds per epoch. It normalises across batch sizes, models and hardware, and it makes the impact of each change immediately readable.
import time
torch.cuda.synchronize()
t0 = time.time()
for x, y in train_loader:
# ... one training step ...
torch.cuda.synchronize()
dt = time.time() - t0
print(f"{len(train_set) / dt:.0f} examples per second")
Baseline on CPU, on GPU, on GPU with mixed precision, on GPU with torch.compile: four numbers, always in the same units, and the decision to buy hardware or optimise code becomes evidence-based.
On a laptop-scale model, float16 on a GPU is often no faster than float32 because the model is too small to saturate the Tensor Cores. Measure first: if the GPU utilisation is already below 80 % on float32, the bottleneck is elsewhere, and mixed precision buys nothing.
In summary
.to(device, non_blocking=True)combined withpin_memory=Trueoverlaps host-to-GPU transfers with compute; without both, the accelerator waits.- Mixed precision through
autocastusesfloat16(orbfloat16) for compute andfloat32for weights and optimiser state, and typically doubles throughput on modern GPUs. GradScaleris required withfloat16to prevent gradient underflow, optional withbfloat16; forgettingscaler.update()reproduces the very bug the scaler prevents.torch.compile(model)fuses operations for a 20-60 % speed-up at the cost of a slow first iteration; it composes cleanly with mixed precision.
Next module: checkpoints and resume, because a run this long deserves the ability to survive a crash.