Autograd, custom pieces and performance
By the end: You can extend PyTorch and make a slow training run fast.
- What autograd records
- Writing a custom module and dataset
- Mixed precision and compile
- Profiling
New to pytorch? PyTorch is the Python library almost everyone uses to build and train neural networks. It hands you two things you would otherwise have to write: arrays of numbers that can live on a GPU instead of the main processor, and automatic differentiation — it watches the arithmetic you perform and works out on its own how every weight should move to reduce the error. Layers, optimisers and data loading come with it. Why it matters
You can write a training loop. This section goes underneath it, to what PyTorch is recording when you call an operation on a tensor, which explains most of the confusing errors, and to writing your own dataset and module rather than bending an existing one to fit. The second half is speed: two changes that usually cost nothing, and the habit of profiling first, since the time in a slow training run is rarely where anyone guesses it is.
What autograd does#
Every operation on a tensor with requires_grad records itself in a graph. Calling backward walks that graph in reverse, applying the chain rule. Understanding this explains most confusing errors: a gradient of None means the operation that produced the tensor was not recorded.
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 3
print(type(y.grad_fn).__name__) # PowBackward0 -> the node that knows how to undo **3
y.backward()
print(x.grad.item()) # 12.0 = 3 * x**2
# Common trap: detaching or converting breaks the graph
z = torch.tensor(y.item()) # a new tensor, no history
z.requires_grad # FalseThe graph is rebuilt from scratch on every forward pass, which is what define-by-run means and why a PyTorch model can contain ordinary Python control flow. It is also why the graph is freed as soon as backward has walked it: calling backward twice on the same loss raises rather than silently doubling your gradients.
There are exactly four ways to lose the history, and every gradient-is-None question is one of them.
a = torch.tensor(2.0, requires_grad=True)
print((a * 3).requires_grad) # True normal
print((a * 3).detach().requires_grad) # False detach
print(torch.tensor((a * 3).item()).requires_grad) # False through a Python float
with torch.no_grad():
print((a * 3).requires_grad) # False inside no_grad
print(torch.from_numpy((a * 3).detach().numpy()).requires_grad) # False via numpyThe related confusion is that only leaves get a .grad. A leaf is a tensor you created rather than computed, which in practice means the parameters. Intermediate values have gradients computed and immediately discarded, because keeping every one of them is what the memory is being saved from.
x = torch.tensor(2.0, requires_grad=True)
h = x * 3
(h ** 2).backward()
print(x.grad.item()) # 36.0 d/dx of (3x)^2 is 18x
print(h.grad) # None plus a UserWarning telling you exactly this
x = torch.tensor(2.0, requires_grad=True)
h = x * 3
h.retain_grad() # ask for it explicitly
(h ** 2).backward()
print(h.grad.item()) # 12.0 = 2hThe memory leak that looks like a leak in PyTorch
Accumulating a running loss with total = total + loss keeps every graph alive for the whole epoch, because the sum still references them, and the process grows until it is killed. Adding .item() or .detach() fixes it. This is the single most common cause of a training run whose memory use climbs steadily with no change in batch size, and it is one character of difference.
total = 0.0
for batch in loader:
total = total + criterion(model(batch.x), batch.y) # BAD: graph retained per batch
print(type(total.grad_fn).__name__) # AddBackward0 -> every graph is still reachable
total = 0.0
for batch in loader:
total += criterion(model(batch.x), batch.y).item() # a plain floatA custom dataset#
from torch.utils.data import Dataset
class JsonlDataset(Dataset):
def __init__(self, path, tokenizer):
self.rows = [json.loads(l) for l in open(path)]
self.tok = tokenizer
def __len__(self):
return len(self.rows)
def __getitem__(self, i):
row = self.rows[i]
return self.tok(row["text"]), row["label"]Cheap means what it says. If __getitem__ opens a file, decodes an image and resizes it, that work happens once per example per epoch, and on a hundred epochs you have decoded the same JPEG a hundred times. Anything that does not change between epochs belongs in __init__ or in a preprocessing pass written to disk. Anything random, such as augmentation, has to stay in __getitem__, because that is the point of it.
Opening a file handle in __init__ breaks with workers
An open HDF5 or LMDB handle created in __init__ is inherited by every forked worker process, and they will corrupt each other's reads in ways that look like random data corruption rather than a bug. Open it lazily on first use inside __getitem__ instead, so each worker gets its own. This is the standard way a dataset works perfectly at num_workers=0 and produces nonsense at 4.
A custom module#
A Module is a container that knows which tensors are parameters, so that .parameters(), .to(device) and .state_dict() find them. The rule is that anything assigned to self as a Parameter or a Module is tracked automatically, and anything else is not.
class Residual(nn.Module):
def __init__(self, dim):
super().__init__() # never omit this
self.norm = nn.LayerNorm(dim) # tracked: a Module
self.fc = nn.Linear(dim, dim) # tracked: a Module
self.scale = nn.Parameter(torch.ones(1)) # tracked: a Parameter
self.register_buffer("calls", torch.zeros(1)) # saved, but not trained
self.dropout_p = 0.1 # a plain float: not tracked, fine
def forward(self, x):
return x + self.scale * self.fc(self.norm(x))
m = Residual(64)
print(sum(p.numel() for p in m.parameters())) # 4289 = 64*64 + 64 + 64 + 64 + 1
# linear weight, linear bias, norm weight,
# norm bias, scale. The buffer is not counted.The mistake is assigning a plain tensor, self.mask = torch.ones(n), which is neither a parameter nor a buffer. It stays on the CPU when the model moves to the GPU, and it is missing from the checkpoint, so the model loads without complaint and behaves differently. Buffers exist precisely for this case.
Two free speedups#
- Mixed precision with autocast: roughly half the memory and often nearly twice the throughput, at no quality cost for most models.
- torch.compile: traces the model and fuses operations. A one-line change that frequently gives a substantial speedup, and it can make errors harder to read, so add it last.
scaler = torch.amp.GradScaler("cuda")
for xb, yb in loader:
opt.zero_grad(set_to_none=True)
with torch.autocast("cuda", dtype=torch.bfloat16):
loss = criterion(model(xb), yb) # forward in 16-bit
scaler.scale(loss).backward() # scale up, so small gradients survive
scaler.unscale_(opt) # back to true scale before clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt) # skips the step if the gradients overflowed
scaler.update()The master weights stay in float32 throughout; only the forward pass and the gradients are 16-bit. With bfloat16 the GradScaler is often unnecessary, because bfloat16 keeps float32's exponent range and simply cannot underflow where float32 would not, which is the whole reason it is preferred over float16 for training.
Profile before optimising#
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
for _ in range(10):
train_step(batch)
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))Name Self CPU % Self CPU CPU total # of Calls
aten::mm 20.21% 3.432ms 3.874ms 30
aten::addmm 18.16% 3.084ms 3.495ms 20
Optimizer.step#AdamW.step 7.98% 1.355ms 7.345ms 10
aten::mul_ 7.58% 1.287ms 1.402ms 80
aten::copy_ 5.36% 909.658us 909.658us 250
aten::sum 5.07% 861.917us 921.620us 30
Self CPU time total: 16.984msRead the two time columns differently. Self time is spent in that operation alone; total time includes everything it called. An entry with small self time and large total time is a wrapper, and optimising it means optimising its children. The aten::copy_ row with 250 calls is the usual signature of data being moved around more than necessary.
GPU timings need a synchronise
CUDA operations are queued asynchronously, so a Python timer around a forward pass measures how long it took to enqueue the work, not to do it. Call torch.cuda.synchronize() before reading the clock, or use the profiler, which handles this. Timings that look impossibly fast are almost always this.
The bottleneck is rarely where you guess. Data loading, host-to-device transfers and synchronisation points dominate far more often than the model's arithmetic. A profile takes a minute and saves days of optimising the wrong thing.
Where people go wrong
Optimising the model code before profiling. The time is usually in data loading or synchronisation, and the model was fine.
If this comes up in an interview
Why might a tensor's gradient be None after backward?
The operation that produced it was not recorded in the autograd graph, usually because it was detached, created from a Python number, or computed inside no_grad. Autograd can only backpropagate through what it recorded.
How would you speed up a training run?
Profile first, because the bottleneck is usually data loading or synchronisation rather than arithmetic. Then mixed precision, more data loader workers, removing per-step .item() calls, and torch.compile.
Build this: Profile a training loop, identify the top three costs, fix them one at a time, and measure throughput after each change.
Watch it explained
This is one section of PyTorch, 30 hours across 6 sections. Read it in full to get the whole topic in order.