Making it fast
Lesson 7 of 8
7Making it fast
The big pictureReal speed comes from three places in order: vectorised code instead of Python loops, lower numeric precision where it does no harm, and compiling the model. Every speed-up must be measured on your own hardware, because what speeds up a GPU can slow down a CPU.
What the picture shows: Sixteen bits mean half the memory: 14 GB instead of 28. Both 16-bit types are the same size, but bfloat16 keeps float32's range so it does not overflow easily, which is why it became the preferred choice for training where the GPU supports it.
Map of this part
The concepts in order, click any one to jump to it:
Two ways into the same idea
You can make training faster in two ways: compute with shorter numbers (like rounding prices to the nearest riyal instead of halalas), or have the compiler prepare the computation ahead. But speed costs complexity, so pay only if slowness is really holding you back.
Mixed precision with torch.autocast computes in float16 or bfloat16 while keeping weights in float32, with GradScaler to stop gradients vanishing. torch.compile fuses operations into faster kernels. You measure epoch time and memory before and after, and check accuracy did not drop.
Before You Start
🧊Tensors
You know dtype sets the size of each number in memory, and that a tensor lives on a device.
🚚Dataset and DataLoader
You know num_workers, since the first speed-up is keeping the GPU busy.
The Core Idea
Real speed comes from three places, in order. In this lesson you remove Python loops first, then learn the limits of 16-bit numbers, then wrap the step in mixed precision, then measure on your machine before you believe anything.
1Vectorise first
Before any advanced trick, make sure the heavy computation is not running in a Python loop. That alone is often the biggest speed-up you will get, and the cheapest.
import time
import numpy as np
x = np.random.default_rng(0).normal(size=1_000_000)
start = time.perf_counter()
total = 0.0
for v in x:
total += v * v
loop_s = time.perf_counter() - start
start = time.perf_counter()
total_np = float((x * x).sum())
vec_s = time.perf_counter() - start
print("same result:", abs(total - total_np) < 1e-6)
print("vectorised at least 3x faster:", vec_s * 3 < loop_s)
The same result, with the vectorised version many times faster (tens or hundreds of times depending on the machine). The reason you saw in the arrays lesson: the vectorised operation runs one loop written in C, while the Python loop interprets each element separately.
2The limits of 16-bit numbers
A 16-bit number takes half the memory of a 32-bit one, but it pays in range or in precision. Look at the picture: float16 gives the exponent only 5 bits, which is why its largest value is 65504.
import numpy as np
x = np.float32(1.0)
print("float32 1 + 1e-8:", x + np.float32(1e-8) == x)
h = np.float16(1000.0)
print("float16 1000 + 0.3:", h + np.float16(0.3))
print("float16 max:", np.finfo(np.float16).max)
print("float16 70000:", np.float16(70000.0))
print("bytes per value: float32 =", np.dtype(np.float32).itemsize, "float16 =", np.dtype(np.float16).itemsize)
float16 cannot represent 1000.3, so it rounds to 1000.5, its maximum is 65504, and 70000 becomes infinity. That is why float16 needs a GradScaler, which scales the loss up before backward so small gradients do not become zero, then scales back. bfloat16 avoids the range problem and pays with less precision.
float16 is more precise but overflows easily, while bfloat16 is less precise but has the same range as float32, so it needs no loss-scaling tricks.
3Mixed precision: autocast and GradScaler
Mixed precision runs most operations in 16 bits and keeps the sensitive ones in 32. Look at the picture: the master weights stay float32, and only the computation drops to 16 bits.
This is the full shape of a float16 training step. It needs a GPU, so it is here for reading.
scaler = torch.amp.GradScaler("cuda")
model = torch.compile(model)
for xb, yb in loader:
xb, yb = xb.to("cuda"), yb.to("cuda")
with torch.autocast("cuda", dtype=torch.float16):
loss = loss_fn(model(xb), yb)
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
The five core lines did not change; they were only wrapped. autocast chooses the precision per operation, and scaler protects small gradients. Note torch.compile at the top: one line, and the first batch will take a minute or more to compile.
4Measure on your machine before you believe
Run the suspect part more than once with time.perf_counter(), because the first run includes loading and warm-up. Look at the picture: the general advice says lower precision is faster, and this CPU says the opposite.
This code needs PyTorch on your machine, so the output shown is recorded from a real run. Install it with pip install torch and run it yourself.
import time
import torch
from torch import nn
torch.manual_seed(0)
model = nn.Sequential(nn.Linear(512, 2048), nn.GELU(), nn.Linear(2048, 512))
x = torch.randn(256, 512)
def timed(fn, repeats=20):
fn()
start = time.perf_counter()
for _ in range(repeats):
fn()
return (time.perf_counter() - start) / repeats * 1000
with torch.no_grad():
fp32 = timed(lambda: model(x))
with torch.autocast("cpu", dtype=torch.bfloat16):
out = model(x)
bf16 = timed(lambda: model(x))
print("output dtype under autocast:", out.dtype)
print(f"fp32 {fp32:.1f} ms, bf16 {bf16:.1f} ms per batch")
On this CPU bfloat16 is more than twice as slow. Ordinary CPUs lack fast 16-bit arithmetic units, so the conversion costs without paying back. On a modern GPU the result flips. This is exactly why the rule is "measure before you accept": the general advice "lower precision is faster" is wrong here.
On a GPU, call torch.cuda.synchronize() before reading the clock, because operations there are asynchronous. torch.compile comes last, and only for long training.
How It Works, Step by Step: how to speed up training, in order
Measure first with time.perf_counter(), with warm-up and repeats, and call torch.cuda.synchronize() on a GPU before reading the clock.
Move heavy computation out of Python loops into vectorised operations, often the biggest speed-up of all.
Wrap the forward pass and loss in torch.autocast, and add a GradScaler if you use float16.
Add torch.compile(model) for long training, and ignore the first batches when timing because they include compile time.
Worked example
Before any advanced trick, make sure the heavy computation is not running in a Python loop. That alone is often the biggest speed-up you will get.
2. Measure before you optimise
Wrap the suspect part in time.perf_counter() and run it more than once, since the first run includes loading and warm-up. On a GPU, call torch.cuda.synchronize() before reading the clock, because operations there are asynchronous.
torch.cuda.synchronize(). Operations are queued on the card and return immediately before finishing, so you time the queuing, not the computation, and think your change made the code a hundred times faster.
torch.compile. Log time per epoch before and after each change.
Try it without code
Myth vs Reality
| Myth | Reality |
|---|---|
| Lower precision is always faster. | On the lesson's CPU bfloat16 was more than twice as slow, because ordinary CPUs lack fast 16-bit arithmetic units. |
| In mixed precision the weights are stored in 16 bits. | The master weights stay float32 and only the computation runs at lower precision. A small update like 0.0001 is lost when added to a float16 weight. |
| torch.compile speeds things up from the first batch. | The first run is slower because it analyses the model and generates code. The speed-up shows in long training, and may not be worth it for runs of a few minutes. |
Try it yourself
You added torch.compile and training got slower in the first minute. The explanation?
torch.compile analyses the model and generates code on the first call, which takes time. Correct timing ignores the first batches. For training that lasts a few minutes it may not be worth it at all.
Which type suits mixed-precision training better if your GPU supports both, and why?
The problem with 16-bit training is range, not precision: gradients and activations can exceed 65504 or underflow to zero. bfloat16 fixes that directly, which is why it became the default for training large models on GPUs that support it.
Quick review cards
torch.compile doesCheck Your Understanding
Why are the master weights not stored in float16 in mixed precision, but kept in float32?
Because a weight update is usually tiny compared with the weight itself, like adding 0.0001 to 1.0. In float16 that addition rounds to nothing, as you saw with 1000 + 0.3, and learning stops silently. Keeping weights in float32 lets small updates accumulate, while only the heavy computation runs in 16 bits.
Why did 70000 become infinity in float16 while 1000 stayed a normal number?
Because the largest value float16 can represent is 65504, so anything bigger overflows to infinity. That is why float16 needs a GradScaler and bfloat16 does not.
Resources for this lesson
Grouped by type. Start with the official docs for precision, or an article or video for the explanation.
Official docs3
Done with this lesson?
Sign in to save your progress.