PyTorch in One Sitting

You now know that a model is matrix multiplies, a softmax, a loss, and gradient descent. Doing that in plain Python is possible and hopeless: a single GPT-2 forward pass is about a hundred billion multiplications, and you would have to write the derivative of every one by hand. You need a tool that does big matrix math fast and computes gradients for you.

That tool is PyTorch. It has two halves, and this chapter is organized around them. The first half is tensors: NumPy-like arrays that can live on a GPU. The second is autograd: PyTorch records every operation you do on a tensor and can replay the chain rule backwards through the record. Everything else (modules, optimizers, data loaders) is convenience built on those two.

Every code block here runs on PyTorch 1.12 or newer with no GPU. Type them in rather than reading them; the sitting is meant to be at a keyboard.

Tensors: creation, shapes, dtype, device

What is a tensor? A rectangular block of numbers with a shape, plus two labels: what kind of number (dtype) and where it lives (device). Vectors, matrices and the $(B, L, d)$ blocks from the math chapter are all tensors of different rank.

import torch

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])   # from a nested list
print(x.shape, x.dtype, x.device)

z = torch.zeros(2, 3)            # shape given as separate ints
o = torch.ones(2, 3)
r = torch.randn(2, 3)            # standard normal, mean 0, variance 1
a = torch.arange(6).reshape(2, 3)   # 0..5, integers
print(a.dtype)

Three habits to build right now. First, print .shape constantly; it is the single most informative thing about a tensor. Second, notice the dtype: torch.tensor([1, 2]) is integer, torch.tensor([1., 2.]) is float32, and models want float32 (or float16/bfloat16 later, for speed). Third, .device is cpu until you move it, and two tensors on different devices cannot be combined.

Shapes are read left to right, outermost first. x.shape == (2, 3) means 2 rows of 3. A 1-d tensor has shape (3,), with the comma, and a scalar has shape (); x.sum() returns one, and .item() converts it to a plain Python number.

v = torch.tensor([1., 2., 3.])
print(v.shape, v.ndim)             # (3,)  1
s = v.sum()
print(s.shape, s.item())           # ()  6.0
h = torch.randn(4, 8, 16)          # batch=4, seq=8, dim=16
print(h.shape, h.numel())          # 512 numbers in total

Broadcasting: how shapes stretch to fit

You want to add a bias vector of shape (4,) to every row of a (3, 4) matrix. Writing a loop is slow and ugly. Should you copy the bias three times first? No. PyTorch does it for you, virtually, under a rule called broadcasting. The rule is short and you must know it exactly, because it silently produces wrong shapes when misapplied.

Align the two shapes on the right. Walk the dimensions from the right. At each position the sizes must either be equal, or one of them must be 1 (or missing, which counts as 1). A size-1 dimension is stretched to match the other. The result has the larger size at every position.

The code matches the picture exactly:

A = torch.tensor([[10.], [20.], [30.]])     # (3, 1)
B = torch.tensor([[1., 2., 3., 4.]])        # (1, 4)
print((A + B).shape)                        # (3, 4)

X = torch.randn(3, 4)
bias = torch.tensor([1., 2., 3., 4.])       # (4,) -> treated as (1, 4)
print((X + bias).shape)                     # (3, 4), bias added to every row

col = torch.tensor([1., 2., 3.])            # (3,)
# X + col  -> error: (3,4) vs (3,) aligns 4 against 3
print((X + col[:, None]).shape)             # (3,1) broadcasts down the columns

Indexing, and views versus copies

How do you get at one token's vector inside a $(B, L, d)$ block? Indexing, and it works like NumPy: integers pick, slices a:b take ranges, ... means "all the remaining axes", None inserts a new axis of size 1.

h = torch.randn(4, 8, 16)         # (B, L, d)
print(h[0].shape)                 # first sequence          (8, 16)
print(h[0, 3].shape)              # its 4th token           (16,)
print(h[:, -1].shape)             # last token of every seq (4, 16)
print(h[:, :4, :].shape)          # first 4 tokens          (4, 4, 16)
print(h[..., :8].shape)           # first 8 features        (4, 8, 8)
print(h[:, None].shape)           # insert an axis          (4, 1, 8, 16)

ids = torch.tensor([2, 5, 5])
print(h[0, ids].shape)            # gather tokens 2,5,5     (3, 16)
mask = torch.tensor([True, False, True, False])
print(h[mask].shape)              # boolean select on B     (2, 8, 16)

Here is the catch that bites everyone once. Slicing, .view(), .reshape() (usually), .transpose() and .T do not copy the numbers. They return a view: a new shape laid over the same memory. Modify the view and you modify the original.

x = torch.arange(6.)
v = x.view(2, 3)        # a view: same 6 numbers, new shape
v[0, 0] = 100.
print(x)                # x changed too
c = x.clone().view(2, 3)   # .clone() copies; now they are independent
c[0, 0] = -1.
print(x[0])

Why does PyTorch do this? Speed. Reshaping a billion-entry tensor should not take time or memory, and views make it free. The price is that you have to remember. Rule of thumb: if you are about to modify something in place and are not sure whether it aliases something else, .clone() it.

One more wrinkle: .view() only works when the underlying memory is laid out contiguously in the order the new shape needs. After a .transpose() it is not, and .view() raises an error. .reshape() does the same thing but silently copies when it must, and .contiguous() forces the copy explicitly. In the attention code you will see .transpose(1, 2).contiguous().view(B, L, d): swap axes, make it contiguous, reshape. Now you know why all three are there.

Matrix multiplication and @

The one operation that is 90% of a model's arithmetic. In PyTorch it is the @ operator (or torch.matmul), and it follows the shapes rule from the math chapter with one addition: extra leading dimensions are treated as a batch and broadcast.

X = torch.randn(5, 3)             # 5 tokens, 3 features
W = torch.randn(3, 2)             # maps 3 features -> 2
print((X @ W).shape)              # (5, 2)

q = torch.randn(4, 8, 16)         # (B, L, d)
k = torch.randn(4, 8, 16)
scores = q @ k.transpose(-2, -1)  # (4, 8, 16) @ (4, 16, 8) -> (4, 8, 8)
print(scores.shape)

v = torch.tensor([1., 2., 3.])
print((v @ v).item())             # vector @ vector = dot product = 14

Read the middle example carefully; it is the attention score computation. k.transpose(-2, -1) swaps the last two axes, turning $(B, L, d)$ into $(B, d, L)$ so the inner dimensions match, and the batch axis $B$ rides along. The result is $B$ separate $L \times L$ tables, one per sequence. The -2, -1 convention (count from the end) is what keeps the same line working whether or not there is a batch dimension in front.

Autograd: PyTorch computes the gradients

You know from the math chapter that gradients come from the chain rule applied backwards through a pipeline. Writing that by hand for a whole model is exactly the tedium we want to avoid. How does PyTorch avoid it? By watching.

Mark a tensor with requires_grad=True. From then on, every operation involving it is recorded into a computation graph: a bookkeeping structure that remembers what was computed from what. When you call .backward() on a scalar at the end, PyTorch walks the graph in reverse, multiplying local derivatives, and deposits the result for each marked tensor in its .grad attribute.

x = torch.tensor(2.0)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(-1.0, requires_grad=True)

u = w * x          # 6
v = u + b          # 5
y = v ** 2         # 25
print(y, y.grad_fn)   # grad_fn shows the recorded op

y.backward()       # run the chain rule backwards
print(w.grad, b.grad)   # dy/dw = 2v * x = 20,  dy/db = 2v = 10

Check it against the math: $y = (wx + b)^2$, so $dy/dw = 2(wx+b)\cdot x = 2 \cdot 5 \cdot 2 = 20$ and $dy/db = 2(wx+b) = 10$. The graph that produced those numbers looks like this.

Step through the same graph with the numbers. Forward values first, then the backward pass one edge at a time.

Three details that matter in practice. Gradients accumulate: a second .backward() adds to .grad rather than replacing it, which is why training loops zero the gradients every step. Only leaf tensors you marked keep a .grad; intermediate results like u do not, to save memory. And .backward() needs a scalar; for a loss that is a vector you first reduce it with .mean() or .sum().

Linear regression by hand: the whole training idea in 20 lines

Before touching any convenience classes, train something with just tensors and autograd. The task: points $(x, y)$ that roughly follow $y = 1.5x + 0.5$; find $w$ and $b$ from the data. This is a two-parameter language model with no language, and every training loop in the course has exactly this skeleton.

import torch
torch.manual_seed(0)

# data: y = 1.5 x + 0.5 + noise
x = torch.linspace(-2, 2, 32)
y = 1.5 * x + 0.5 + 0.3 * torch.randn(32)

w = torch.zeros(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.1

for step in range(100):
    pred = w * x + b                    # forward
    loss = ((pred - y) ** 2).mean()     # mean squared error, a scalar
    loss.backward()                     # gradients land in w.grad, b.grad
    with torch.no_grad():               # update without recording a graph
        w -= lr * w.grad
        b -= lr * b.grad
    w.grad.zero_(); b.grad.zero_()      # gradients accumulate: reset them
    if step % 25 == 0:
        print(f"step {step:3d}  loss {loss.item():.3f}  w {w.item():.2f}  b {b.item():.2f}")

Look at the four things inside the loop, because they are the four things inside every loop: forward, loss, backward, update. The torch.no_grad() block matters: without it the update w -= ... would itself be recorded into the graph, and the next backward pass would try to differentiate through your optimizer. The .zero_() calls matter because gradients accumulate.

Play with the same problem below. The learning rate and step count are the knobs; the fitted line and the loss curve are the readout. Push the learning rate past about 0.5 and watch the loss climb instead of fall, the divergence you saw on the 1-d bowl, now in two parameters.

nn.Module: packaging parameters and forward

The manual loop works, but a model has hundreds of weight tensors. Tracking them by hand, moving them all to the GPU, saving them, feeding them to an optimizer: that needs a container. nn.Module is that container. You subclass it, create sub-modules and parameters in __init__, and write the computation in forward.

import torch, torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self, d_in, d_hidden, d_out):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)   # weight (d_hidden, d_in), bias (d_hidden,)
        self.fc2 = nn.Linear(d_hidden, d_out)
    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

m = TinyMLP(16, 64, 4)
out = m(torch.randn(8, 16))      # call the module, not .forward(), so hooks run
print(out.shape)
for name, p in m.named_parameters():
    print(f"{name:12s} {tuple(p.shape)}")
print(sum(p.numel() for p in m.parameters()), "parameters")

Notice nn.Linear(d_in, d_out) stores its weight as (d_out, d_in) and computes $xW^\top + b$. That transpose surprises people reading the math chapter's $XW$; it is just a storage convention. Anything assigned as an attribute that is itself a Module or an nn.Parameter is registered automatically, which is how .parameters() finds everything, including inside nested modules and nn.ModuleLists.

The three layers a transformer is made of

Besides nn.Linear, a GPT uses exactly two other parameterized layers, and both are simple.

emb = nn.Embedding(num_embeddings=50257, embedding_dim=768)   # a lookup table (V, d)
ids = torch.tensor([[464, 3290, 3332]])                       # (B=1, L=3) token ids
print(emb(ids).shape)                                         # (1, 3, 768): row lookup

ln = nn.LayerNorm(768)          # per-token: subtract mean, divide by std, then scale & shift
h = torch.randn(1, 3, 768) * 5 + 2
out = ln(h)
print(out.mean(-1), out.std(-1))   # each token now ~mean 0, std 1 (before learned scale/shift)

nn.Embedding is a matrix of shape (vocabulary, $d$) indexed by token id; that is the "embed" box in the model map from the first chapter, and it is trained like any other weight. nn.LayerNorm normalizes each token's vector to mean 0 and variance 1 across its features and then applies a learned per-feature scale and shift; it keeps activations at a sane scale through dozens of layers. Both get their own sections in the embedding and layers chapters; here you only need to know their shapes.

Optimizers: torch.optim.AdamW

In the manual loop, the update was w -= lr * w.grad for each parameter. An optimizer object does that for all parameters at once, and the good ones do it smarter. The one used for essentially every language model is AdamW: it keeps a running average of each parameter's gradient and of its square, scales each step by their ratio (so parameters with noisy gradients take smaller steps), and applies weight decay directly to the weights.

opt = torch.optim.AdamW(m.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)

# the three calls you make every step:
opt.zero_grad()        # clear .grad on every parameter
loss = ...             # forward + loss
loss.backward()        # fill .grad
opt.step()             # update every parameter from its .grad

Those hyperparameters (lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1) are the GPT-style defaults and are what code/lumen/train.py uses. The optimizers chapter builds AdamW from scratch; for now, treat it as "gradient descent with per-parameter step sizes".

The canonical training loop

Every training script in this course, from linear regression to the capstone, is this loop. Learn the five moves in order.

import torch.nn.functional as F

model = TinyMLP(16, 64, 4)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

for step in range(200):
    x = torch.randn(32, 16)                 # a batch of 32 inputs
    y = (x[:, 0] > 0).long() * 2 + (x[:, 1] > 0).long()   # a made-up 4-class label
    opt.zero_grad()
    logits = model(x)                       # (32, 4) raw scores
    loss = F.cross_entropy(logits, y)       # softmax + -log p(true), averaged
    loss.backward()
    opt.step()
    if step % 50 == 0:
        print(step, round(loss.item(), 3))

F.cross_entropy takes raw scores (logits) and integer class ids, applies log-softmax and picks out $-\log p$ of the true class, averaged over the batch. It is the loss from the math chapter, in nats. You never apply softmax yourself before it; doing so is a classic bug that trains badly without erroring. For a language model the "classes" are vocabulary entries and the logits have shape $(B, L, V)$, which is flattened to $(B \cdot L, V)$ before the call.

DataLoader basics

Where do the batches come from? For toy problems, from torch.randn. For real data you want shuffling, batching and (on a GPU) background loading. Dataset is anything with __len__ and __getitem__; DataLoader turns it into an iterator over batches.

from torch.utils.data import Dataset, DataLoader

class TokenWindows(Dataset):
    """Every item is (context of L tokens, the next L tokens shifted by one)."""
    def __init__(self, ids, L):
        self.ids, self.L = ids, L
    def __len__(self):
        return len(self.ids) - self.L - 1
    def __getitem__(self, i):
        x = self.ids[i : i + self.L]
        y = self.ids[i + 1 : i + self.L + 1]
        return x, y

ids = torch.randint(0, 100, (10_000,))          # pretend these are token ids
dl = DataLoader(TokenWindows(ids, L=8), batch_size=4, shuffle=True)
x, y = next(iter(dl))
print(x.shape, y.shape)
print(x[0]); print(y[0])

Look at x[0] and y[0]: the target is the input shifted one position left. That shift is next-token prediction, and the same eight tokens give eight training examples at once. code/lumen/data.py is this class plus a tiny corpus.

Saving and loading: state_dict

A trained model is its parameter values. model.state_dict() is an ordered dictionary from parameter name to tensor, and it is the thing you save. The architecture (the Python class) is not saved; you rebuild it and load the values in.

torch.save(model.state_dict(), "tiny.pt")

model2 = TinyMLP(16, 64, 4)                          # same architecture
model2.load_state_dict(torch.load("tiny.pt", map_location="cpu"))
print(list(model2.state_dict().keys())[:2])

This is exactly how you will load OpenAI's GPT-2 weights in month 3: build your own GPT2 class in code/lumen/gpt2.py, then copy their tensors into your state_dict name by name, with a few transposes where their layout convention differs from nn.Linear's. Also save the optimizer's state_dict() if you want to resume a training run; AdamW's running averages are state too.

GPU and MPS: moving things

Everything so far ran on the CPU. To use a GPU, pick a device string and move both the model and every batch to it. Tensors on different devices cannot interact, so the pattern is always "model once, batches every step".

if torch.cuda.is_available():
    device = "cuda"
elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
    device = "mps"             # Apple silicon
else:
    device = "cpu"

model = TinyMLP(16, 64, 4).to(device)      # moves all parameters in place
x = torch.randn(32, 16).to(device)          # returns a copy on the device
out = model(x)
print(out.device)
print(out.cpu().shape)                      # bring back for numpy / plotting

Two things to remember. model.to(device) modifies the module in place, but tensor.to(device) returns a new tensor, so you must assign it. And .item() or .cpu() on a GPU tensor forces the GPU to finish and copy back, which is slow if done every step; log the loss every 50 steps, not every one.

Common shape bugs and how to read the error

Most of your PyTorch time in month 2 will be spent on shape errors. They are not mysterious once you know the three messages and what each is telling you.

The debugging move that works: put print(name, t.shape) after every line in forward, run one batch, and compare against the shapes you expected. Write the expected shape as a comment on each line, the way the code in code/lumen/ does. Those comments are not decoration; they are the spec.

torch.no_grad and eval mode

When you generate text, you do not need gradients. Recording the graph costs memory (it keeps every intermediate activation, as the math chapter explained) and time. Wrap inference in torch.no_grad(), and it records nothing.

model.eval()                     # switch off dropout etc. (mode flag on the module)
with torch.no_grad():            # do not build a graph
    logits = model(x)
    probs = torch.softmax(logits, dim=-1)
    next_id = torch.multinomial(probs, num_samples=1)   # sample one per row
model.train()                    # back to training mode before the next step

These are two different switches, and you need both. model.eval() changes the behavior of certain layers: dropout stops dropping and batch-norm uses its running statistics. torch.no_grad() changes the bookkeeping: no graph is recorded. Forgetting eval() makes generation noisy; forgetting no_grad() makes it slow and eventually runs out of memory. The sampling helpers in code/lumen/sampling.py use both.

The 20 tensor operations this course uses

You will see these again and again. Each line is a one-sentence spec; the shapes in comments are the point.

Practice

Further reading

• PyTorch documentation. Broadcasting semantics. The rule from this chapter, stated formally, in one page.
• PyTorch documentation. Autograd mechanics. What is recorded, when, and how in-place operations interact with it.
• PyTorch tutorials. Learn the Basics. Official walkthrough covering tensors, datasets, modules and the training loop with runnable notebooks.
• Andrej Karpathy. micrograd. A 100-line scalar autograd engine; reading it demystifies backward() completely.
• Loshchilov & Hutter (2019). Decoupled Weight Decay Regularization. The AdamW paper: why weight decay is applied to the weights rather than folded into the gradient.