PyTorch
The framework research and production both use.
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.
Doing this by hand means deriving and coding the backwards pass for every layer you invent, then rewriting it whenever you change the model. That is days of work and a reliable source of silent bugs, repeated for each experiment. With a framework, changing an architecture is changing a few lines. It is also the shared language: papers ship PyTorch code, and so does most of what you will need to read.
It sits on top of neural networks and backpropagation — those ideas as working software — and beneath everything that trains: fine-tuning libraries, distributed training and most model code you will open are PyTorch underneath. It borrows its array conventions from numpy and pandas, and it pairs with GPU fundamentals, since most PyTorch frustration turns out to be a memory or device problem.
It keeps track of the arithmetic so it can compute the adjustments for you; what is left is the modelling.
In more detailTensors, autograd, modules, datasets and the training loop. Most of the day is spent on shape errors, device placement, and making data loading fast enough to keep the GPU busy. Writing the loop yourself once before adopting a trainer abstraction is worth the time.
Learn first
PyTorch in practice
By the end: You can write and debug a training pipeline.
- Tensors and devices
- Modules and the training loop
- Datasets and data loading
- Where the time actually goes
This is where the ideas become something that runs. PyTorch is what nearly all research and a good deal of production is written in, and the section covers the parts you touch daily: tensors and the hardware they sit on, modules, the training loop, and getting data in quickly enough that the GPU is not left idle. Expect the difficulties to be mechanical rather than conceptual, which is why shapes and devices get most of the attention.
The error you will actually see What it means
mat1 and mat2 shapes cannot be multiplied (32x768 and 512x256) inner dimensions disagree
Tensor for argument input is on cpu but expected on mps model moved, batch not (CUDA words it
as: all tensors on the same device)
expected target dtype to be Long or Byte, but got Float cross_entropy wants integer class labels
Trying to backward through the graph a second time backward called twice, or a tensor kept across steps
element 0 of tensors does not require grad you are inside no_grad, or the input was detached
Can't call numpy() on Tensor that requires grad add .detach()
CUDA out of memory batch too large, or you accumulated a tensor
that dragged its graph along with itTensors#
A tensor is an array that tracks operations so gradients can be computed, and that can live on a GPU. Most day-to-day friction is shapes and devices.
It carries five properties, and nearly every error is one of them being wrong. Print all five when something is confusing rather than guessing which one it is.
import torch
x = torch.randn(32, 768)
print(tuple(x.shape)) # (32, 768) what it holds
print(x.dtype) # torch.float32 how each number is stored
print(x.device) # cpu where the memory lives
print(x.requires_grad) # False whether autograd is watching
print(x.element_size() * x.nelement()) # 98304 bytes, = 32 * 768 * 4dtype bytes 1000x1000 Use it for
torch.float32 4 4.0 MB the default; everything, until memory hurts
torch.bfloat16 2 2.0 MB mixed-precision training; same exponent range as float32
torch.float16 2 2.0 MB inference and older GPUs; overflows above 65504
torch.int64 8 8.0 MB class labels and indices; PyTorch's default integer
torch.int32 4 4.0 MB indices, when int64 is wasting memory
torch.bool 1 1.0 MB masksThe conversion between numpy and torch has a trap in it that is worth meeting deliberately. torch.from_numpy shares the underlying memory, so editing the array changes the tensor; torch.tensor copies. One is fast and surprising, the other is safe and slower.
import numpy as np
a = np.array([1.0, 2.0, 3.0])
shared = torch.from_numpy(a) # shares memory with a
copied = torch.tensor(a) # copies
a[0] = 99.0
print(a) # [99. 2. 3.]
print(shared.numpy()) # [99. 2. 3.] followed the edit
print(copied.numpy()) # [ 1. 2. 3.] did not
g = torch.ones(3, requires_grad=True)
g.numpy() # RuntimeError: Can't call numpy() on Tensor that requires grad.
g.detach().numpy() # [1. 1. 1.]Devices#
Every tensor in an operation must live on the same device. Decide once at the top of the script and move everything to the same place, rather than sprinkling .cuda() through the code and discovering at step 300 that one buffer stayed behind.
device = ("cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available() # Apple silicon
else "cpu")
model = model.to(device) # in place for a Module
x = x.to(device) # NOT in place for a Tensor: returns a copy
for xb, yb in loader:
xb, yb = xb.to(device), yb.to(device) # the batch must be moved too, every stepx = torch.randn(4, 8).to("mps")
x @ torch.randn(8, 4)
# RuntimeError: Tensor for argument #2 'mat2' is on CPU, but expected it to be on GPU
# (while checking arguments for mm)Where the model and the data get separated
Almost always the model was moved in the setup code and the batch is moved nowhere, because the loader yields CPU tensors and nothing says otherwise. The second most common cause is a tensor you created inside forward, such as torch.arange for positions, which lands on the CPU unless you pass device=x.device. Writing that argument by reflex costs nothing and removes a whole category of bug.
A model#
import torch.nn as nn
class Classifier(nn.Module):
def __init__(self, in_dim, hidden, classes):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(hidden, classes),
)
def forward(self, x):
return self.net(x) # raw scores, loss applies softmaxThe training loop, line by line#
Six lines, and every one of them is load-bearing. Learn what each does and the loop stops being something you copy.
model.train() # 1
for xb, yb in train_loader: # 2
xb, yb = xb.to(device), yb.to(device) # 3
optimizer.zero_grad(set_to_none=True) # 4
logits = model(xb) # 5
loss = criterion(logits, yb) # 6
loss.backward() # 7
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 8
optimizer.step() # 91 model.train() switches dropout and batchnorm into training behaviour.
2 for xb, yb in ... the loader yields one batch; workers prepared it while the GPU worked.
3 .to(device) the loader always yields CPU tensors. This is the line people omit.
4 zero_grad gradients ACCUMULATE in PyTorch. Without this, batch 50 carries
the sum of the first 50 batches and the run diverges silently.
set_to_none=True frees the buffers rather than filling them with 0.
5 model(xb) forward pass. Every operation is recorded onto the autograd graph.
6 criterion(...) one scalar. Feed it raw logits: the loss applies softmax itself.
7 loss.backward() walks that graph in reverse, adding each parameter's gradient
into p.grad. Adds, not assigns, which is why line 4 exists.
8 clip_grad_norm_ rescales all gradients if their combined norm exceeds 1.0.
Optional for small models, standard for transformers.
9 optimizer.step() applies p -= lr * p.grad, or Adam's version of it.Why gradients accumulate at all
It looks like a design mistake until you need it. Accumulation is what lets you simulate a batch of 512 on hardware that fits 64: run eight forward-and-backward passes, then step once. Divide the loss by 8 so the gradient magnitude matches a real batch of 512. Gradient checkpointing, multi-task losses and most memory tricks lean on the same behaviour.
Train and eval mode#
Two calls matter here and both fail quietly. Dropout and batch normalisation behave differently during training and inference, so forgetting model.eval() before validation gives noisy, pessimistic numbers, and forgetting model.train() afterwards silently disables dropout for the rest of the run.
It is worth knowing exactly what changes, because the answer is narrower than most people assume. The flag does nothing to gradients, nothing to memory and nothing to speed. It sets a boolean on every submodule, and only two common layer types read it.
model = nn.Sequential(nn.Linear(10, 10), nn.Dropout(0.5), nn.BatchNorm1d(10))
x = torch.randn(8, 10)
model.train()
print([round(model(x)[0, 0].item(), 4) for _ in range(5)])
# [1.5063, -0.366, -0.4132, 1.4859, -0.4225] same input, five different answers
model.eval()
print([round(model(x)[0, 0].item(), 4) for _ in range(5)])
# [0.6156, 0.6156, 0.6156, 0.6156, 0.6156] deterministicDropout zeroes a random half of the activations during training, and to keep the average signal the same it scales the survivors up by two. In eval it does nothing at all, passing the input straight through. That scaling is why you must not compensate by hand.
d = nn.Dropout(0.5); t = torch.ones(100000)
d.train(); print(round(d(t).mean().item(), 4), round((d(t) == 0).float().mean().item(), 3))
d.eval(); print(round(d(t).mean().item(), 4), round((d(t) == 0).float().mean().item(), 3))
# train mean 0.9986 zeros 0.500 half deleted, the rest doubled
# eval mean 1.0000 zeros 0.000 untouchedBatch normalisation is the more damaging of the two to get wrong, because it does not merely add noise; in training mode it normalises using the statistics of whatever batch it is handed. Give it two examples and each is forced to the batch's own mean and range, which produces the same output regardless of the input.
torch.manual_seed(0)
bn = nn.BatchNorm1d(4)
big = torch.randn(256, 4) * 5 + 3 # mean about 3, variance about 25
bn.train()
for _ in range(50):
bn(big) # learns running statistics from real batches
print(bn.running_mean.round(decimals=2).tolist()) # [2.65, 3.80, 3.41, 2.72]
print(bn.running_var.round(decimals=1).tolist()) # [20.8, 29.3, 28.7, 25.7]
two = big[:2]
bn.train(); print(bn(two)[0].round(decimals=3).tolist())
bn.eval(); print(bn(two)[0].round(decimals=3).tolist())
# train mode, batch of 2: [-1.0, -1.0, 1.0, 1.0] every value is exactly +-1
# eval mode, batch of 2: [-1.08, -1.149, -0.291, -0.249] the real numbersbn.train(); bn(big[:1])
# ValueError: Expected more than 1 value per channel when training,
# got input size torch.Size([1, 4])Those running statistics are also why a batch-normalised model has state beyond its weights. They live in the state dict as buffers, they must be saved with the checkpoint, and a model evaluated before they have stabilised will score badly for reasons that have nothing to do with its weights.
What no_grad actually changes#
A separate mechanism from eval, controlling a separate thing, and the two are needed together rather than instead of one another. eval changes layer behaviour; no_grad stops the autograd graph being built. Inside it, operations run normally but nothing records what it did, so no intermediate activations are kept.
m = nn.Sequential(*[nn.Linear(512, 512) for _ in range(20)])
x = torch.randn(256, 512)
print(m(x).grad_fn is not None) # True the graph is being built
with torch.no_grad():
print(m(x).grad_fn is not None) # False nothing recorded
# measured on CPU, 20 layers, batch 256:
# forward with gradients tracked 4.4 ms live tensors 32.0 MB
# forward under no_grad 3.9 ms live tensors 22.1 MB
#
# difference 10.0 MB, against 20 x 256 x 512 x 4 bytes = 10.5 MB of
# stored activations. That is exactly what was not kept. model.eval() torch.no_grad()
changes dropout yes no
changes batchnorm yes no
saves memory no yes, a lot
makes it faster no slightly
prevents backward() no yes
scope until .train() the with-block onlymodel.train()
for batch in train_loader:
...
model.eval()
with torch.no_grad(): # no gradients needed, saves memory and time
for batch in val_loader:
...
model.train() # switch back, or dropout stays off for the rest of the runDataset and DataLoader#
A Dataset answers two questions: how many examples are there, and give me the i-th one. That is the entire interface, and writing one is three methods. The DataLoader wraps it and handles batching, shuffling, parallel loading and collation.
from torch.utils.data import Dataset, DataLoader
class Rows(Dataset):
def __init__(self, X, y):
self.X = torch.as_tensor(X, dtype=torch.float32)
self.y = torch.as_tensor(y, dtype=torch.long) # long: cross_entropy demands it
def __len__(self):
return len(self.y) # how many
def __getitem__(self, i):
return self.X[i], self.y[i] # the i-th, unbatched
ds = Rows(torch.randn(1000, 20), torch.randint(0, 3, (1000,)))
print(len(ds), tuple(ds[0][0].shape)) # 1000 (20,) one example, no batch dimensionloader = DataLoader(ds, batch_size=64, shuffle=True, num_workers=0)
xb, yb = next(iter(loader))
print(tuple(xb.shape), tuple(yb.shape), len(loader)) # (64, 20) (64,) 16
# 1000 / 64 = 15 full batches and a final one of 40.
print(len(DataLoader(ds, batch_size=64, drop_last=True))) # 15The default collate function stacks the examples, which requires them all to be the same shape. Variable-length data needs your own, and this is the one everybody eventually writes.
from torch.nn.utils.rnn import pad_sequence
def collate(batch):
lengths = torch.tensor([len(s) for s in batch])
return pad_sequence(batch, batch_first=True, padding_value=0), lengths
padded, lengths = collate([torch.arange(n) for n in (3, 7, 5)])
print(tuple(padded.shape), lengths.tolist()) # (3, 7) [3, 7, 5]
# [[0, 1, 2, 0, 0, 0, 0],
# [0, 1, 2, 3, 4, 5, 6],
# [0, 1, 2, 3, 4, 0, 0]]A DataLoader batches, shuffles and loads in parallel. It is frequently the bottleneck: a GPU sitting idle waiting for data is the most common performance problem in training, and it looks like the model being slow.
loader = DataLoader(
dataset,
batch_size=64,
shuffle=True, # training only. Never shuffle validation; it wastes time.
num_workers=4, # parallel loading, the usual fix for an idle GPU
pin_memory=True, # faster host to device transfer
persistent_workers=True, # do not respawn the workers every epoch
prefetch_factor=2, # batches queued per worker, ready before they are asked for
)num_workers is not free either. Each worker is a separate process holding its own copy of the dataset object, so a Dataset that loaded a large array in __init__ is now in memory several times over, and workers respawn every epoch unless persistent_workers is set, which on a small dataset can cost more than the loading saved. Raise it from zero while measuring throughput and stop where it stops helping, usually around the number of physical cores.
Almost everything here is PyTorch, which is the right default: it is what most research code and most production serving is written in. The exception worth knowing about is JAX, used heavily at some frontier labs and built around a different idea — you write pure functions and transform them, with jit for compilation, grad for derivatives, vmap for batching and pmap or shard_map for spreading work across devices. It is also the native stack for TPUs, Google's accelerators, where the unit of thinking is a mesh of chips and how your arrays are sharded across it. You do not need to learn it before an interview unless the job says so, but you should not be surprised by the word, and if you are applying somewhere whose published work is in JAX, read enough to follow their code.
Where the time goes#
- Data loading, usually. Check GPU utilisation: if it is low, the model is not the problem.
- Moving data between CPU and GPU. Do it once per batch, not per item.
- Small batches, which underuse the hardware.
- Full precision, where mixed precision would be roughly twice as fast at the same quality.
Where people go wrong
Calling .item() or .cpu() inside the training loop to log something. Each one forces synchronisation with the GPU and stalls the pipeline. Accumulate on the device and read once per epoch.
Build this: Write a full training loop by hand: dataset, loader, model, optimiser, train and eval modes, and per-epoch metrics. Then measure GPU utilisation and improve it.
Open this section on its own pageAutograd, 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
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.
Build this: Profile a training loop, identify the top three costs, fix them one at a time, and measure throughput after each change.
Open this section on its own pageShapes, broadcasting and the bugs they cause
By the end: You can predict the shape of any tensor expression, reshape safely, and spot the broadcasting bug that trains without an error.
- Reading shapes
- Broadcasting rules
- view, reshape, permute and contiguity
- einsum for anything complicated
- dtypes that matter
Nearly every PyTorch bug is a shape bug, and the dangerous ones raise no error at all: subtract a tensor of shape [N] from one of shape [N, 1] and PyTorch quietly compares every prediction with every target, trains happily, and hands you a useless model. This section is about reading and predicting shapes, the broadcasting rules that decide when two tensors may meet, the difference between view, reshape and permute, and einsum once more than two dimensions are in play. The habit it is trying to install is writing down the shape of every tensor as you create it.
Say the shape out loud#
Most PyTorch bugs are shape bugs. The habit that prevents them: write the shape of every tensor as a comment when you create it, in named dimensions such as batch, time and channels.
import torch
x = torch.randn(32, 128, 256) # [batch, time, d_model]
w = torch.randn(256, 64) # [d_model, d_head]
h = x @ w # [batch, time, d_head] -> [32, 128, 64]Broadcasting#
When two tensors of different shapes meet in an operation, PyTorch lines their shapes up from the right. Each pair of dimensions must be equal, or one of them must be 1, which is then stretched. Missing dimensions on the left count as 1.
[32, 128, 256] + [256] -> [32, 128, 256] bias added to every position
[32, 1, 256] + [1, 128, 256] -> [32, 128, 256]
[32, 128] + [32] -> error: 128 and 32 do not matchTwo rules produce all of it. Missing dimensions are added on the left as 1, and any dimension of size 1 is stretched to match the other. Nothing is copied while doing it, which is why broadcasting is fast and also why a [1000, 1] against a [1000] silently allocates a million elements without warning you that it has.
The bug that does not crash#
pred = model(x) # [N, 1]
y = targets # [N]
loss = ((pred - y) ** 2).mean() # [N, 1] - [N] broadcasts to [N, N]. No error.
loss = ((pred.squeeze(-1) - y) ** 2).mean() # what you meantThe first version compares every prediction with every target. It trains, the loss goes down a little, and the model is useless. Assert shapes at the boundaries of your functions to catch this.
The reason this survives review is that the number it prints is entirely plausible. Compare the two on random data:
torch.manual_seed(0)
pred = torch.randn(1000, 1)
y = torch.randn(1000)
print(tuple((pred - y).shape)) # (1000, 1000) a million elements
print(((pred - y) ** 2).mean().item()) # 1.9996
print(((pred.squeeze(-1) - y) ** 2).mean().item()) # 2.0374And it does train, which is the part that costs a day. Fitting a linear model to a problem it can solve exactly, with and without the squeeze:
for squeeze in (True, False):
torch.manual_seed(0)
X = torch.randn(500, 4); w = torch.randn(4, 1)
Y = (X @ w).squeeze(-1) + 0.1 * torch.randn(500)
m = nn.Linear(4, 1); opt = torch.optim.SGD(m.parameters(), lr=0.05)
for _ in range(300):
p = m(X)
loss = ((p.squeeze(-1) - Y) ** 2).mean() if squeeze else ((p - Y) ** 2).mean()
opt.zero_grad(); loss.backward(); opt.step()
print(squeeze, round(((m(X).squeeze(-1) - Y) ** 2).mean().item(), 4))
# squeeze=True true mse 0.0102 solved
# squeeze=False true mse 4.0588 400 times worse, no error, loss fell the whole waySo the defence cannot be watching the loss. It has to be an assertion at the point the shapes meet, which costs one line and no measurable time.
def mse(pred, target):
assert pred.shape == target.shape, f"{tuple(pred.shape)} vs {tuple(target.shape)}"
return ((pred - target) ** 2).mean()
# Better still, make the intent explicit rather than relying on a squeeze:
loss = mse(pred.reshape(-1), target.reshape(-1))Where else the silent broadcast appears
Anywhere a length-N vector meets an [N, 1] column. Computing per-example weights for a weighted loss, masking with a [N] boolean against [N, 1] scores, adding a bias of shape [d] to something of shape [d, 1], and subtracting a per-row mean you forgot to keep dims on. The rule that catches all of them: if an operation should produce N numbers and you cannot say why the result is not N by N, print the shape.
Reshaping#
- view changes the shape without copying, and only works when the memory layout allows it.
- reshape does the same when it can and copies when it must.
- permute and transpose reorder dimensions by changing strides, which leaves the tensor non-contiguous. A view after them fails; call contiguous() first or use reshape.
The underlying idea is the stride: the number of elements to step to move one place along each dimension. A tensor is a flat block of memory plus a shape and a stride, so transpose does not move a single number, it swaps two entries in the stride tuple. That is why it is free, and why the result no longer matches the layout view assumes.
q = torch.randn(32, 128, 8 * 64) # [batch, time, heads * d_head]
q = q.view(32, 128, 8, 64).transpose(1, 2) # [batch, heads, time, d_head]
out = q.transpose(1, 2).reshape(32, 128, 8 * 64) # back: reshape, since transpose broke contiguityt = torch.randn(32, 128, 512).view(32, 128, 8, 64).transpose(1, 2)
print(tuple(t.shape)) # (32, 8, 128, 64)
print(t.is_contiguous()) # False
print(t.stride()) # (65536, 64, 512, 1) <- dim 1 steps 64, dim 2 steps 512
t.view(32, 128, 512)
# RuntimeError: view size is not compatible with input tensor's size and stride
# (at least one dimension spans across two contiguous subspaces)
t.reshape(32, 8 * 128, 64) # works: copies when it must
t.contiguous().view(32, 8, 8192) # works: copy first, then a free viewSo reshape is not simply the safer view. reshape returns a view when it can and a copy when it cannot, which means it never raises but its cost is unpredictable, and the returned tensor may or may not share memory with the original. view always shares and always raises if it cannot. Use view where you want to be told about an accidental copy in a hot loop, and reshape where you only want the result.
einsum#
For anything beyond a plain matrix multiply, einsum names every dimension and says which are summed. It is often the clearest way to write attention.
scores = torch.einsum("bhqd,bhkd->bhqk", q, k) / d_head ** 0.5
# d is summed away; b, h, q, k remainThe rule is short: name every dimension of every input, then name the ones you want out. Any letter that appears on the left and not on the right is summed over. Nothing else is going on, and it makes the alternative, a sequence of transposes and unsqueezes whose correctness you have to reconstruct each time you read it, unnecessary.
q = torch.randn(2, 4, 16, 32); k = torch.randn(2, 4, 16, 32)
a = torch.einsum("bhqd,bhkd->bhqk", q, k)
b = q @ k.transpose(-2, -1)
print((a - b).abs().max().item()) # 0.0 the same operation, written twice
print(tuple(a.shape)) # (2, 4, 16, 16)einsum string What it does
"ij,jk->ik" plain matrix multiply
"bij,bjk->bik" batched matrix multiply
"bhqd,bhkd->bhqk" attention scores: sum over the head dimension d
"bhqk,bhkd->bhqd" apply those scores to the values
"bi,bi->b" row-wise dot product, without a sum and a multiply
"bij->bji" transpose
"ii->i" the diagonal
"bchw,bcHW->bhwHW" when you have gone too far and should draw it insteaddtypes#
- Class labels for cross_entropy must be integers (long), not floats.
- Default floats are 32-bit. Mixed precision uses 16-bit for speed; bfloat16 has the same range as 32-bit and rarely overflows, float16 can.
- Mixing dtypes in one operation raises an error or silently upcasts. Check with .dtype when something is odd.
Where people go wrong
Relying on the absence of errors. PyTorch will happily compute the wrong thing if the shapes broadcast.
Build this: Write multi-head attention three ways: with view and transpose, with einsum, and with scaled_dot_product_attention. Assert every intermediate shape, and check all three give the same output to within floating point tolerance.
Open this section on its own pageBuild and train a small transformer from scratch
By the end: You have written, trained and sampled from a character-level transformer, and can explain every line.
- Data as a stream of token ids
- Batches of context and next-token targets
- A transformer block in twenty lines
- The training loop with validation
- Sampling with temperature
You know roughly how a transformer works; this is writing one, in about eighty lines, and training it on a book until it produces text. Embeddings, a causal attention block, residual adds, a training loop with validation, and sampling with a temperature: the same structure as the models you call over an API, with the scale taken out. Writing it yourself is what turns the architecture diagram into something you can debug, and it is the piece a research interview is most likely to ask you to reproduce live.
Piece In this file In a production LLM
tokeniser one integer per char BPE, 100k+ merged subwords
context length T 128 128k to 1M
width D 256 4,000 to 16,000
heads H 4 32 to 128
layers L 4 32 to 120
parameters 3.2 million billions
position information learned embedding RoPE, applied inside attention
normalisation LayerNorm, pre-norm RMSNorm, pre-norm
attention dense causal grouped-query, with a KV cache
training data one book trillions of tokensData#
Use any plain text file of a few megabytes, such as a public domain book. Each character becomes an integer. The task is to predict the next character from the ones before it.
Characters rather than subwords because it removes a whole component: no tokeniser to train, no vocabulary file, and a vocabulary of about 65 that you can print. The cost is that the model spends its capacity learning spelling before it learns anything else, which is why the samples are gibberish for the first thousand steps and then abruptly contain real words.
import torch, torch.nn as nn, torch.nn.functional as F
text = open("book.txt", encoding="utf-8").read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
split = int(0.9 * len(data))
train, val = data[:split], data[split:]
V, T, D, H, L = len(chars), 128, 256, 4, 4 # vocab, context, width, heads, layers
def batch(src, n=32):
ix = torch.randint(len(src) - T - 1, (n,))
x = torch.stack([src[i:i + T] for i in ix]) # [n, T]
y = torch.stack([src[i + 1:i + T + 1] for i in ix]) # the same, shifted by one
return x.cuda(), y.cuda()The y line is the entire supervision signal and it is worth staring at. The target is the input shifted by one, so position 0 must predict what sits at position 1, position 1 predicts position 2, and so on. One sequence of 128 characters therefore provides 128 separate training examples, all computed in a single forward pass. That is why language models are efficient to train and why the causal mask is not optional.
The model#
class Block(nn.Module):
def __init__(self):
super().__init__()
self.ln1, self.ln2 = nn.LayerNorm(D), nn.LayerNorm(D)
self.qkv = nn.Linear(D, 3 * D)
self.proj = nn.Linear(D, D)
self.mlp = nn.Sequential(nn.Linear(D, 4 * D), nn.GELU(), nn.Linear(4 * D, D))
def forward(self, x): # [B, T, D]
B, T_, _ = x.shape
q, k, v = self.qkv(self.ln1(x)).split(D, dim=-1)
q, k, v = (t.view(B, T_, H, D // H).transpose(1, 2) for t in (q, k, v))
a = F.scaled_dot_product_attention(q, k, v, is_causal=True) # [B, H, T, D/H]
x = x + self.proj(a.transpose(1, 2).reshape(B, T_, D)) # attention, added
return x + self.mlp(self.ln2(x)) # MLP, added
class TinyGPT(nn.Module):
def __init__(self):
super().__init__()
self.tok, self.pos = nn.Embedding(V, D), nn.Embedding(T, D)
self.blocks = nn.Sequential(*[Block() for _ in range(L)])
self.ln, self.head = nn.LayerNorm(D), nn.Linear(D, V)
def forward(self, idx): # [B, T] of ids
pos = torch.arange(idx.shape[1], device=idx.device)
x = self.tok(idx) + self.pos(pos)
return self.head(self.ln(self.blocks(x))) # [B, T, V] logitsidx [32, 128] token ids
tok(idx) + pos(arange) [32, 128, 256] embeddings, added not concatenated
per block:
ln1(x) [32, 128, 256]
qkv(...) [32, 128, 768] one matmul for all three
.split(D, dim=-1) 3 x [32, 128, 256]
.view(B, T, H, D//H) [32, 128, 4, 64] split the width across heads
.transpose(1, 2) [32, 4, 128, 64] heads become a batch dimension
attention [32, 4, 128, 64] internally [32, 4, 128, 128] scores
.transpose(1,2).reshape(...) [32, 128, 256] heads rejoined
x + proj(...) [32, 128, 256] residual
x + mlp(ln2(x)) [32, 128, 256] residual
head(ln(x)) [32, 128, 65] a distribution per positionTwo design choices in that walkthrough are worth naming because they recur everywhere. Q, K and V come from one Linear of width 3D rather than three of width D, because a single large matrix multiply is faster than three small ones and the split afterwards is free. And multi-head attention is not several attention modules; it is one, with the width reshaped so the head index behaves like an extra batch dimension. Nothing in the attention code knows there are heads.
Map it to what you know
Every piece here appears in large language models: token and position embeddings, pre-norm blocks with a residual add after attention and after the MLP, a causal mask, a final norm and a projection to vocabulary logits. Scale changes; the structure does not.
token embedding 65 x 256 16,640
position embedding 128 x 256 32,768
per block:
2 LayerNorms 2 x 2 x 256 1,024
qkv 256 x 768 + 768 197,376
proj 256 x 256 + 256 65,792
mlp 256x1024 + 1024 + 1024x256 + 256 525,568
per block 789,760
x 4 blocks 3,159,040
final norm + head 512 + 256 x 65 + 65 17,217
TOTAL 3,225,665Memory, separately, is dominated by something that has no parameters at all. The attention scores are [batch, heads, T, T], which here is 32 by 4 by 128 by 128, or 2.1 million numbers per layer held for the backward pass. Double the context length and that term quadruples while the parameter count does not move. That single fact is why context length is expensive, why scaled_dot_product_attention is used here instead of writing the softmax by hand, and why FlashAttention exists.
Train it#
model = TinyGPT().cuda()
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for step in range(5001):
x, y = batch(train)
loss = F.cross_entropy(model(x).view(-1, V), y.view(-1))
opt.zero_grad(); loss.backward(); opt.step()
if step % 500 == 0:
model.eval()
with torch.no_grad():
vx, vy = batch(val)
vloss = F.cross_entropy(model(vx).view(-1, V), vy.view(-1))
model.train()
print(f"step {step:5} train {loss.item():.3f} val {vloss.item():.3f}")At step 0 the loss should be close to the natural log of the vocabulary size, about 4.2 for 65 characters, because an untrained model guesses uniformly. If it starts far from that, something is wrong before training even begins.
Run on a smaller corpus with a 28-character vocabulary, on CPU:
loss at init 3.4867 ln(28) = 3.3322 close enough: the model is uninformed
logits shape (32, 128, 28) one distribution per position
step 0 train 3.494 val 3.128
step 100 train 0.171 val 0.192
step 200 train 0.021 val 0.023
step 300 train 0.014 val 0.016
sample: 'the n liquor jugs. pack my box with five doze dozen liquor jugs. pack my...'Running it on something trivial first is a habit worth keeping. A model that cannot memorise two sentences has a bug, and you find that out in ninety seconds rather than after an overnight run on a real book. It is the same argument as overfitting a single batch, applied to the whole pipeline including sampling.
On a real book expect a different shape: validation loss around 1.5 after a few thousand steps, samples that are pronounceable nonsense at step 500, recognisable English words by 2,000, and something with the rhythm of the source and no meaning by 5,000. Validation loss rising while training loss falls means the model is large for the corpus, which at 3.2 million parameters on a single book it genuinely is.
Sample from it#
@torch.no_grad()
def sample(start, n=400, temperature=0.8):
model.eval()
idx = torch.tensor([[stoi[c] for c in start]], device="cuda")
for _ in range(n):
logits = model(idx[:, -T:])[:, -1] / temperature # last position only
nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return "".join(chars[i] for i in idx[0].tolist())
print(sample("The "))Four details in that function are the whole of inference. The model is run on idx[:, -T:] because it has only T position embeddings and cannot be given more; a longer prompt must be truncated. Only the last position's logits are used, because the others are predicting characters you already have. multinomial samples from the distribution rather than taking the argmax, which is what stops the output repeating a single phrase forever. And each new token is appended and the whole thing re-run, which is why generation is sequential and slow.
That last point is also the one real inefficiency here, and naming it is worth marks. Every step recomputes the keys and values for all the preceding positions, having computed exactly the same numbers on the previous step. Caching them is the KV cache, and it turns generation from quadratic in the output length into linear. This implementation has none, which is fine for 400 characters and hopeless for a chat model.
temperature effect when
0.0 (argmax) fully deterministic, repeats extraction, classification, code completion
0.7 to 0.8 coherent with some variety the usual default
1.0 the model's raw distribution when you want the true samples
> 1.2 increasingly incoherent creative tasks, and not much elseTemperature 0 is not a temperature
Dividing by zero is undefined, so every implementation special-cases it to mean argmax. Handle it with an if rather than passing 1e-8 and hoping. torch.softmax survives 1e-8 because it subtracts the row maximum internally, but a hand-written exp(logits / t) / sum at that temperature overflows to NaN across the whole row, which is the failure the max subtraction exists to prevent.
Where people go wrong
Copying a reference implementation and running it. The learning is in writing it yourself and fixing the shape errors you hit.
Build this: Type the model out yourself, train it for 5,000 steps on a public domain book, and sample. Then double the context length, then the layers, and record validation loss and training time for each change.
Open this section on its own pageWhen the loss will not go down
By the end: You have a checklist for a training run that is not learning, and can work through it in order.
- Check the loss at step zero
- Overfit one batch
- Look at the data the model sees
- Learning rate and gradient norms
- NaNs and where they come from
A training run that is not learning tells you almost nothing about why, and the temptation is to start changing the architecture. This section is the checklist to work through instead, in order: is the loss at step zero what uniform guessing would give, can the model overfit a single batch, does the data look right when you decode it back, is the learning rate within a factor of ten, and do gradients reach the parameters at all. Each step rules out a whole class of cause, and the answer is almost never the architecture.
In this order#
- 1.Initial loss: close to log of the number of classes for classification? If not, check initialisation and labels.
- 2.Overfit a single batch: train on the same small batch for a few hundred steps. The loss should go to nearly zero. If it cannot, the model or the loss is broken, and no amount of data will fix it.
- 3.Look at what goes in: decode a batch back into text or images and look at it with the labels. Misaligned labels, wrong normalisation and empty inputs show up immediately.
- 4.Learning rate: too high and loss jumps or explodes; too low and it barely moves. Try ten times higher and lower.
- 5.Gradients: log the total gradient norm. Zero means something is detached or not trainable; enormous means instability, fixed with clipping, warmup or a lower rate.
Step 1: the loss at initialisation#
An untrained classifier should be uninformed, not wrong. Uninformed means equal probability on every class, and the cross-entropy of that is the natural log of the number of classes. This number is free, takes one forward pass, and rules out a surprising amount.
classes expected initial loss if you see instead
2 0.6931 much higher: labels or initialisation
10 2.3026 exactly 0: the model is predicting the label
65 4.1744 nan: bad init, or a loss applied to a softmax
50,257 10.8249 much lower: leakage, or the target is the inputStep 2: overfit one batch#
Take sixteen examples and train on only those for a few hundred steps. A correctly wired model must drive the loss to essentially zero, because memorising sixteen examples requires no generalisation at all. If it cannot, no quantity of data will help and the fault is in the model, the loss or the optimiser.
import torch, torch.nn as nn, torch.nn.functional as F
X = torch.randn(16, 20); Y = torch.randint(0, 5, (16,)) # labels are random noise
model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 5))
opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
for _ in range(400):
loss = F.cross_entropy(model(X), Y)
opt.zero_grad(); loss.backward(); opt.step()
print(loss.item()) # 0.00009 memorised, as it must be
# Now the classic wiring bug: only some parameters handed to the optimiser.
model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 5))
opt = torch.optim.AdamW(model[0].parameters(), lr=1e-2) # layer 2 will never update
# ... same loop ...
print(loss.item()) # 0.00569 60x higher, and still falling. Suspicious, not obvious.Note how weak the failure signal is. The broken model still reached 0.00569, which on an unfamiliar problem reads as a success. The test is only decisive if you know what the healthy number looks like, which is why you run the healthy version first and write the number down.
Step 3: look at what actually goes in#
Decode a batch back into text or images and print it beside its label. This catches the bugs that no gradient statistic can: targets not shifted, images normalised twice, labels off by one from a category mapping, an empty string where a document should be, augmentation that has rotated the digits into different digits. It takes two minutes and it is the step people skip because it is not clever.
Step 4 and 5: learning rate and gradients#
x, y = batch(train) # one fixed batch
for step in range(300):
loss = F.cross_entropy(model(x).view(-1, V), y.view(-1))
opt.zero_grad(); loss.backward()
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # returns the norm before clipping
opt.step()
if step % 50 == 0:
print(step, f"{loss.item():.3f}", f"grad norm {norm:.2f}")The learning rate is swept by factors of ten, not by fine adjustment: try it a hundred times smaller and a hundred times larger and see which direction improves things. Anything finer is wasted before the order of magnitude is right. A useful default to start from is 3e-4 for AdamW on a transformer and 1e-3 for smaller networks.
Gradient norms distinguish causes that the loss curve cannot. Print them per parameter rather than in total, because a total norm of 0.4 is perfectly consistent with one layer receiving nothing.
for name, p in model.named_parameters():
print(name, None if p.grad is None else round(p.grad.norm().item(), 5))
# 0.weight 0.41170 healthy
# 0.bias 0.13331
# 2.weight 0.82731
# 2.bias 0.35066
# with requires_grad = False on layer 0:
# 0.weight None <- not "small". None. It is frozen.What you see Cause Fix
p.grad is None frozen, detached, or not in the graph check requires_grad
p.grad is exactly 0.0 dead ReLU, or a saturated activation lower lr, change init
early layers tiny, late layers big vanishing gradients residuals, normalisation
norm jumps 100x on some batches a handful of extreme examples clip, then inspect them
norm healthy, loss flat the loss does not measure the task re-read the lossNaNs#
- The log of zero or a division by zero in a custom loss. Add a small epsilon, or use the built-in stable functions such as cross_entropy with logits.
- float16 overflow in mixed precision. Use bfloat16 where the hardware supports it, or a gradient scaler.
- An exploding learning rate. The loss rises for a few steps before becoming NaN; look back in the log.
- torch.autograd.set_detect_anomaly(True) reports the operation that produced the first NaN, at a large speed cost. Use it only while debugging.
The first of those is worth seeing, because it is why you are told to hand the loss raw logits rather than probabilities. A hand-written loss takes the log of a probability that is allowed to be exactly zero; the built-in one never forms the probability at all.
p = torch.tensor([0.0, 1.0])
print(torch.log(p)[0].item()) # -inf then -inf * 0 in the backward pass is nan
print(torch.log(p + 1e-8)[0].item()) # -18.42 survivable
# The built-in version works in log space throughout and never forms p at all:
print(F.cross_entropy(torch.tensor([[-1e4, 1e4]]), torch.tensor([0])).item()) # 20000.0When a NaN appears mid-run rather than at the start, it is almost always instability rather than a formula. Look back a few hundred steps in the log: an exploding run rises before it breaks, and the step where the loss first started climbing is the one to explain. If the loss was flat and then NaN in a single step, suspect a single pathological batch and print it.
if not torch.isfinite(loss):
torch.save({"x": xb, "y": yb, "step": step}, "bad_batch.pt")
raise RuntimeError(f"non-finite loss at step {step}")Check the gradients for NaN, not only the loss
A finite loss with a NaN gradient poisons every parameter on the next optimiser step, and from then on the loss is NaN with no clue where it began. Assert on the gradient norm after backward rather than on the loss after forward, because that is the earlier of the two and names the step that did it.
Where people go wrong
Changing the architecture when the loss will not fall. Almost always the bug is in the data, the loss or the learning rate.
Build this: Take your working small transformer and deliberately introduce five bugs one at a time: unshifted targets, a detached block, a learning rate a hundred times too high, labels shuffled against inputs, and a missing zero_grad. Use the checklist to find each and note which step caught it.
Open this section on its own pageCheckpoints and reproducibility
By the end: You can save and resume a training run exactly where it stopped, and make runs as repeatable as the hardware allows.
- What a checkpoint must contain
- Loading on a different device
- Resuming mid-run
- Seeds and deterministic settings
A long run will be interrupted, by a preempted instance or a full disk or a crash at hour nine, and what you saved decides whether you resume or start again. This section is what a checkpoint has to contain to resume exactly: not only the weights, but the optimiser state, the scheduler, the step count and the random number generator states. It also covers how far repeatability actually goes, which is less far than setting seeds suggests, and why comparing several seeds beats trusting one run.
More than the weights#
import torch
torch.save({
"model": model.state_dict(),
"optimizer": opt.state_dict(), # Adam's moments: resuming without them restarts learning dynamics
"scheduler": sched.state_dict(),
"step": step,
"rng": torch.get_rng_state(),
"cuda_rng": torch.cuda.get_rng_state_all(),
}, f"ckpt_{step}.pt")What Why it must be there Left out, you get
model state_dict the weights, and buffers such as BN stats a model, missing its statistics
optimizer Adam's two moment estimates, per parameter a different loss curve, see below
scheduler where in the warmup or decay you were the learning rate resets to step 0
step / epoch where to resume counting logs that overlap and a wrong schedule
RNG states shuffling, dropout, augmentation a different data order from here on
data position which examples have been seen this epoch some examples twice, some never
config the hyperparameters this run used a checkpoint nobody can reproduceThe optimiser row is the one worth measuring, because the cost of omitting it is invisible. AdamW holds two running averages per parameter, exp_avg and exp_avg_sq, plus a step count used for bias correction. Discard them and the first few steps after a resume take badly scaled updates while those averages refill.
# Train 40 steps straight through, and compare against stopping at 20 and resuming.
#
# step | uninterrupted | resumed with opt | resumed without opt
# 20 | 1.039943 | 1.039943 | 1.039943
# 21 | 1.006236 | 1.006236 | 0.999241 <- already different
# 25 | 0.871761 | 0.871761 | 0.848588
# 39 | 0.432757 | 0.432757 | 0.399392
#
# max deviation with optimiser state: 0.0 bit-identical
# max deviation without optimiser state: 0.0338
print(list(opt.state_dict()["state"][0].keys())) # ['step', 'exp_avg', 'exp_avg_sq']Note which way it diverged. The run without optimiser state reached a lower loss, because emptying Adam's second moment produces larger effective steps for a while. That is the trap: the broken resume looks better, so nobody investigates, and the experiment you thought you were continuing is now a different experiment whose curve cannot be compared with anything.
Loading#
ckpt = torch.load("ckpt_4000.pt", map_location="cpu") # load anywhere, move later
model.load_state_dict(ckpt["model"])
opt.load_state_dict(ckpt["optimizer"])
sched.load_state_dict(ckpt["scheduler"])
start_step = ckpt["step"] + 1weights_only
Recent PyTorch versions load checkpoints with weights_only=True by default, refusing arbitrary pickled Python objects. It is a security feature, since loading a pickle can run code. Keep checkpoints to tensors and plain values, and never load untrusted checkpoints with it turned off.
torch.save({"w": torch.ones(2), "obj": SomeCustomClass()}, "x.pt")
torch.load("x.pt")
# UnpicklingError: Weights only load failed. This file can still be loaded, to do so
# you have two options ...
torch.load("x.pt", weights_only=False) # runs whatever the pickle says to run
torch.save({"w": torch.ones(2), "step": 7}, "y.pt")
torch.load("y.pt") # fine: tensors and plain values onlyTwo further habits. Save to a temporary name and rename afterwards, because a process killed midway through writing leaves a truncated file that loads as a corrupt checkpoint rather than as no checkpoint. And keep two: the latest, for resuming, and the best validation so far, for shipping. They are rarely the same file, and the run that overwrites one with the other has lost the model it was trying to produce.
import os
def save_atomic(state, path):
tmp = path + ".tmp"
torch.save(state, tmp)
os.replace(tmp, path) # atomic on the same filesystem
save_atomic(state, "last.pt")
if val_loss < best_val:
best_val = val_loss
save_atomic(state, "best.pt")Repeatable runs#
import random, numpy as np
def seed_everything(seed):
random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
seed_everything(0)
torch.use_deterministic_algorithms(True) # errors on ops without a deterministic versionEven with every seed set, some GPU operations are not bit-for-bit repeatable, and different hardware or library versions differ. Aim for runs that agree closely, record versions, and compare results across several seeds rather than trusting one.
What Repeatable? Why not
same machine, same versions usually yes with seeds set and deterministic algorithms on
DataLoader with num_workers > 0 needs care each worker needs its own seed; use worker_init_fn
atomicAdd on GPU (scatter, some
pooling and index operations) no float addition is not associative and the
thread order is not fixed between runs
different GPU model no different kernels chosen for the same operation
different library version no kernel selection and defaults change
across a cuDNN benchmark cycle no it times algorithms and picks the fastest,
which varies with machine loadThe practical conclusion is not to chase bit-exactness but to stop treating a single run as a measurement. A difference of 0.3 per cent between two configurations is noise unless you have run each three times and the gap is larger than the spread within each. Record the seed alongside the result, report the mean and the range, and reserve determinism for debugging, where you genuinely need two runs to differ in only the thing you changed.
The versions are part of the checkpoint
Write torch.__version__, the CUDA version, the Python version and your own commit hash into the checkpoint dict. It costs a few bytes and is the difference between reproducing a six-month-old result and guessing at it. Loading a checkpoint saved by a much older PyTorch usually works; loading one into a much older PyTorch usually does not.
Where people go wrong
Saving only the model weights every epoch, then discovering that resuming a crashed run changes the loss curve because the optimiser started from scratch.
Build this: Train your small transformer for 2,000 steps, kill it at 1,000, resume from the checkpoint, and confirm the loss curve matches an uninterrupted run closely.
Open this section on its own pageWhere you use this
A topic is easier to care about when you can see what it is for. These projects lean on it at a named stage.
- Stage 2, Feed it data: Batch and time dimensions are where almost every bug in this project lives.
- Stage 3, One head of self-attention: Transposing the wrong two dimensions here produces working code and nonsense results.
- Stage 6, Train it: The specific thing to check first when the loss sits flat, which it will.
- Stage 3, Set up LoRA, and find out what the adapter is: The adapter is two small matrices per targeted layer; the arithmetic only makes sense if you can see their shapes.
- Stage 4, Train it: The specific checks when the loss sits flat, which for a LoRA run is usually the prompt template or a frozen adapter.
- Stage 7, Merge, save and serve the adapter: Saving so that the thing you load next month is provably the thing you measured.
Now actually do it
Reading a topic is not learning it. These are the hours that count, and each one tells you when you are finished.
1Write the training loop by hand
12hDataset, DataLoader, model, optimiser, train and eval modes, per-epoch metrics, checkpointing. No trainer abstraction.
Done when: It trains correctly and you called model.eval() and torch.no_grad() in the right places without prompting.
2Fix a starving GPU
8hMeasure GPU utilisation during training. Raise num_workers, enable pin_memory, move preprocessing off the critical path, and measure after each change.
Done when: You raised utilisation substantially and can say which change did most of it.
3Find the synchronisation stall
5hAdd a .item() call inside the inner training loop and measure throughput. Then accumulate on device and read once per epoch.
Done when: You can state the throughput cost of per-batch synchronisation.
25 hours of practice
Words used here
3 terms this topic teaches, each with a one-line definition if one of them stopped you.
- JAX
- A numerical library built around transforming pure functions rather than objects that hold state: jit compiles one, grad differentiates it, vmap batches it, and pmap or shard_map spreads it across devices. Used heavily at some frontier labs, and the native stack for TPUs.
- Checkpoint
- A saved snapshot of weights, and usually optimiser state and step count too, so a run can be resumed or an earlier model recovered.
- Transpose
- Flipping a matrix over its diagonal, so rows become columns. In model code it is usually there to line up the two dimensions a multiply needs to share.
Interview questions
9 questions, grouped by the section that teaches the answer. An answer you have read is not an answer you can give, so use Test yourself and say it out loud first.
PyTorch in practice
What does model.eval() do?
Switches layers that behave differently at inference, notably dropout and batch normalisation, into evaluation mode. Forgetting it makes validation numbers noisy and pessimistic, and forgetting to switch back disables dropout for the rest of training.
Why wrap validation in torch.no_grad()?
It stops the framework recording operations for gradient computation, which is unnecessary during evaluation. It saves memory and time, and without it a large validation pass can run out of memory.
Your GPU utilisation is 20 percent during training. What is wrong?
Almost certainly data loading, not the model. Raise num_workers, enable pin_memory, and check for expensive per-item preprocessing that should be precomputed or moved onto the GPU.
Autograd, custom pieces and performance
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.
Shapes, broadcasting and the bugs they cause
A regression model's loss decreases slowly but predictions are nonsense. What would you check first?
Shapes. A prediction of shape [N, 1] subtracted from targets of shape [N] broadcasts to [N, N] and computes a meaningless loss without error. Assert shapes where the loss is computed, and squeeze or reshape explicitly.
Then they push further
Why does view sometimes fail after transpose?
Transpose only changes strides, so the data is no longer laid out contiguously for the new shape. view cannot reinterpret non-contiguous memory. Call contiguous() first, or use reshape, which copies if it has to.
What gets marked up. Marked up for naming the silent broadcast bug specifically. It is one of the most common real training bugs.
Build and train a small transformer from scratch
Walk me through a minimal GPT-style model.
Token ids go through token and position embeddings and are added. A stack of blocks follows, each with layer norm, causal multi-head self-attention added back to the residual, then layer norm and an MLP added back. A final norm and a linear layer produce logits over the vocabulary. Training minimises cross-entropy between logits and the next token at every position.
Then they push further
Why is the loss at initialisation a useful check?
An untrained model should predict roughly uniformly, so its loss should be close to the log of the vocabulary size. Much higher means a bad initialisation or a bug, such as targets not shifted correctly.
Why is the causal mask needed during training?
Training predicts every position at once. Without the mask, each position could attend to the tokens it is supposed to predict, and the model would learn to copy rather than to predict.
What gets marked up. Strong candidates can write this from memory and explain residual adds, pre-norm and the mask. Many applicants for research roles are asked to implement part of it live.
When the loss will not go down
Your model's training loss is flat from the first step. How do you debug it?
Check the initial loss against the expected value for uniform guessing, try to overfit one small batch, decode and inspect a batch with its labels, sweep the learning rate by factors of ten, and log gradient norms to see whether gradients reach the parameters at all.
Then they push further
It cannot overfit a single batch. What does that tell you?
The problem is in the model, the loss or the optimiser setup, not the data volume. Common causes are parameters not passed to the optimiser, a detached tensor, a wrong loss for the output format, or labels that do not match the inputs.
What gets marked up. Marked up for 'overfit one batch' as an early step and for a systematic order rather than random changes.
Checkpoints and reproducibility
What do you need to save to resume a training run exactly?
Model weights, optimiser state, scheduler state, the step or epoch counter, random number generator states, and the position in the data. Missing the optimiser state is the common mistake: Adam restarts with empty moment estimates and training behaves differently.
Resources
Free to watch and read, picked for this topic.
For pyTorch in practice
For autograd, custom pieces and performance
For shapes, broadcasting and the bugs they cause
For build and train a small transformer from scratch
For when the loss will not go down
For checkpoints and reproducibility
Full courses: the whole topic, start to finish
Read and explore