Lesson 1: Six prerequisites, each checked by running code

Module 1 · Lesson 1 · About 45 minutes · Needs: Python, NumPy, matrix products and the chain rule; having trained one neural network

The situation

It is sol 400 at Aurora Station. The station record up to sol 399 is now your corpus, and today you start tinyformer-lab, the transformer you will build, train and take apart over fifteen modules.

First you check, by running code, six facts the rest of the course leans on. Kofi Mensah, who reviews the lab's claims, will ask where each number came from, so every number is labelled from today.

A real case, first

In 2022 Horace He, then on the PyTorch team, published "Making Deep Learning Go Brrrr From First Principles" (horace.io/brrr_intro.html). He splits a model's time into compute, memory bandwidth and overhead, and starts from hardware. On an Nvidia A100, the stated 312 teraflops is reached only by matrix multiplication: "if you aren't doing matrix multiplication, you'll only be able to achieve 19.5 teraflops instead of the stated 312."

He then quotes a FLOP count for BERT: non-matrix operations are only 0.2% of the FLOPs, yet normalisation and pointwise operations achieved 250 and 700 times fewer FLOPS than the matrix multiplications, so they cost more time than their arithmetic suggests. The lesson: count the matrix multiplications, and measure the rest.

What you will be able to do

• Turn text into ids, embeddings and batches, naming every shape.
• Count a matrix product's FLOPs from shapes and measure your machine.
• Use dot products and softmax as similarity, and spot saturation.
• State the language-model contract and check it on GPT-2.
• Derive and check the gradients of a linear layer and cross-entropy.

1. Tokens, embeddings and sequences as tensors (Essential)

A model never sees text, only integers and then vectors, and bugs are easier to find when you know which you hold. Start with the torch setup cell.

import torch
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"   # every lesson runs on "cpu"; a GPU is optional
torch.manual_seed(11)                                     # the course seed (see course-guide/06)
print("torch", torch.__version__, "| device", DEVICE, "| CPU threads", torch.get_num_threads())

Code walkthrough

• torch.manual_seed(11) sets the course seed. The author's Mac printed torch 2.14.1 | device cpu | CPU threads 5.

The running example is the first line of the first document in the corpus: "Bruno Ferreira works in the water plant on the coolant loop." The lab's tokeniser gives each character an id from a vocabulary of 69. A token is whatever unit a tokeniser produces; here, one character.

SPECIALS = ["<pad>", "<doc>", "<unk>"]   # padding, document boundary, unseen character
CHARS = list("\n ,-.0123456789:?ABCDEFGHIJKLMNOPQRSTUVWYZabcdefghiklmnopqrstuvwxy")
VOCAB = SPECIALS + CHARS
STOI = {s: i for i, s in enumerate(VOCAB)}
PAD, DOC, UNK = 0, 1, 2

def encode(text):
    return [STOI.get(c, UNK) for c in text]

def decode(ids):
    return "".join(VOCAB[i] if i >= len(SPECIALS) else "" for i in ids)

LINE = "Bruno Ferreira works in the water plant on the coolant loop."
ids = torch.tensor([encode(LINE)])                  # (batch 1, time 60)
print(len(VOCAB), tuple(ids.shape), ids[0, :8].tolist(), repr(decode(ids[0, :8].tolist())))

Code walkthrough

• Copied from course-guide/06, block 7, except CHARS: that block computes it from the whole corpus; here it is written out, so ids match every other lesson without generating the corpus.
• ids has shape (1, 60): one sequence of 60 characters.

GPT-2, the course's reference model, uses learned word pieces: 17 tokens for the same line, from a vocabulary of 50,257. Fewer tokens means less work per line. Module 8 trains such a tokeniser.

Ids become vectors through an embedding table, one learned row per vocabulary entry: (69, 128) at the base model's width. Looking up id 21 returns row 21, exactly what a one-hot row (zeros and a single 1) times the table gives.

import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(11)
embed = nn.Embedding(len(VOCAB), 128)              # table (69, 128), random until trained
x = embed(ids)                                     # (1, 60) ids -> (1, 60, 128) vectors
onehot = F.one_hot(ids, len(VOCAB)).float()        # (1, 60, 69)
print(tuple(x.shape), "max |lookup - one-hot @ table| =", (onehot @ embed.weight - x).abs().max().item())

batch = torch.stack([ids[0, :20], ids[0, 20:40], ids[0, 40:60]])   # three 20-character slices
print(tuple(batch.shape), tuple(embed(batch).shape), [decode(r.tolist()) for r in batch])

Code walkthrough

• nn.Embedding(69, 128) makes a random table; training would shape it.
• onehot @ embed.weight matches the lookup: difference 0.0.
• torch.stack piles three slices into a batch of shape (3, 20), (3, 20, 128) once embedded.

Memorise (batch, time, width): how many sequences, how many tokens each, how long each token's vector. Almost every tensor in a transformer has this shape or a relative of it. Now load GPT-2.

%pip install -q transformers==5.19.0

from transformers import AutoModelForCausalLM, AutoTokenizer

REF = "openai-community/gpt2"          # the course's reference model: GPT-2 small, MIT licence
tok = AutoTokenizer.from_pretrained(REF)
ref = AutoModelForCausalLM.from_pretrained(REF, attn_implementation="eager").eval()
g = tok(LINE, return_tensors="pt")
print(tuple(g.input_ids.shape), tok.convert_ids_to_tokens(g.input_ids[0])[:6])
print(tuple(ref.transformer.wte(g.input_ids).shape), "embedding table", tuple(ref.transformer.wte.weight.shape))

Code walkthrough

• The two from_pretrained calls download GPT-2 small once (548 MB); .eval() switches off dropout.
• Output: (1, 17), pieces ['B', 'run', 'o', 'ĠFer', 're', 'ira'] (Ġ marks a space), embeddings (1, 17, 768) from a (50257, 768) table, wte.

2. Matrix multiplication as the unit of compute (Essential)

Nearly all of a transformer's arithmetic is matrix multiplication, so you can work out its cost from shapes before running anything. That is how this course budgets experiments.

Put the 60 embedded characters through a (128, 128) layer. Each of the 7,680 outputs is a dot product of length 128, and each step is a multiply and an add: 2 FLOPs (floating-point operations). So (m, k) times (k, n) costs 2 × m × k × n: here 1,966,080 FLOPs (arithmetic).

import time

def matmul_flops(m, k, n):
    return 2 * m * k * n                           # one multiply and one add per step

def median_ms(f, runs=10):
    for _ in range(3):
        f()                                        # warm-up calls are not timed
    times = []
    for _ in range(runs):
        t0 = time.perf_counter(); f(); times.append(time.perf_counter() - t0)
    return 1000 * sorted(times)[runs // 2]

print("running example:", f"{matmul_flops(60, 128, 128):,}", "FLOPs")
for n in [64, 256, 1024, 2048]:
    a, b = torch.randn(n, n), torch.randn(n, n)
    ms = median_ms(lambda: a @ b)
    print(f"{n:5d}  {matmul_flops(n, n, n):>14,} FLOPs  {ms:8.3f} ms  {matmul_flops(n, n, n) / ms / 1e6:7.0f} GFLOP/s")

Code walkthrough

• matmul_flops is the figure's formula.
• median_ms warms up three times, times ten runs and returns the middle one in milliseconds, ignoring odd slow runs.
• The loop divides FLOPs by time: billions of FLOPs per second (GFLOP/s).

On the author's Mac (Apple silicon, 5 threads; Colab is slower), small products are inefficient and large ones level off at 1,800 to 1,900 GFLOP/s (measured here, two runs). Doubling every side multiplies FLOPs by 8, and the time from 1,024 to 2,048 grew by close to 8. Cost follows shapes.

Now He's point on your CPU: an add and a GELU (an activation function) on the same (2048, 2048) matrices, one FLOP per element. The product did about 130 times more arithmetic per second than the add and 250 times more than the GELU (measured here).

a, b = torch.randn(2048, 2048), torch.randn(2048, 2048)
for name, f, flops in [("matmul", lambda: a @ b, matmul_flops(2048, 2048, 2048)),
                       ("add", lambda: a + b, 2048 * 2048),
                       ("gelu", lambda: F.gelu(a), 2048 * 2048)]:
    ms = median_ms(f)
    print(f"{name:7s} {ms:7.3f} ms  {flops / ms / 1e6:8.1f} GFLOP/s")

Code walkthrough

• Each entry pairs an operation with its FLOP count.
• The add reads two 16 MB matrices and writes a third for one FLOP per number, so memory speed limits it, not arithmetic.

3. Dot products and softmax as similarity (Essential)

Attention, in Module 2, is a dot product followed by a softmax. Predict what these two do and attention holds no surprises.

A dot product multiplies two vectors entry by entry and adds the results; vectors pointing the same way score high. Divided by both lengths it is the cosine similarity, between -1 and 1. In GPT-2's table, " water" is closest to " liquid" and furthest from " Bruno".

table = ref.transformer.wte.weight.detach()        # (50257, 768)
def emb(word):
    (i,) = tok.encode(word)                          # each word below is one GPT-2 token
    return table[i]

q = emb(" water")
for w in [" liquid", " oxygen", " plant", " Bruno"]:
    k = emb(w)
    print(f"{w:8s} dot {float(q @ k):6.3f}  cosine {float(F.cosine_similarity(q, k, dim=0)):.3f}")

for d in [16, 128, 1024]:
    u, v = torch.randn(10_000, d), torch.randn(10_000, d)
    print(f"d {d:5d}  spread of random dot products {float((u * v).sum(-1).std()):6.2f}  sqrt(d) {d ** 0.5:6.2f}")

Code walkthrough

• (i,) = ... fails loudly if a word splits into two tokens.
• Cosines: liquid 0.460, oxygen 0.415, plant 0.324, Bruno 0.143.
• The second loop prints the standard deviation of 10,000 random dot products per length.

Random dot products spread out as the square root of the length: about 4 at length 16, about 32 at 1,024 (measured here, matching the arithmetic). Longer vectors give bigger scores for no other reason.

Softmax turns scores into positive weights that sum to 1: raise e to each score, divide by the total. Dividing scores by a temperature first sharpens the result (below 1) or flattens it (above 1).

s = torch.tensor([2.0, 1.0, 0.0])
for T in [1.0, 0.25, 4.0]:
    print(f"temperature {T}: {[round(float(p), 3) for p in (s / T).softmax(-1)]}")

torch.manual_seed(11)
q, K = torch.randn(1024), torch.randn(6, 1024)     # one query, six keys, length 1,024
raw = K @ q                                        # six dot products
print("raw scores", [round(float(r), 1) for r in raw])
print("softmax(raw)      max weight", round(float(raw.softmax(-1).max()), 3))
print("softmax(raw / 32) weights   ", [round(float(p), 3) for p in (raw / 32).softmax(-1)])

Code walkthrough

• (s / T).softmax(-1) applies a temperature, then softmax along the last axis.
• The second part scores six random keys against one query at length 1,024 and compares softmax with and without dividing by 32, the square root of 1,024.

Raw scores run from -19 to 24 and softmax gives 0.983 to one key; divided by 32, every key keeps a share. A saturated softmax passes almost no gradient to the losers, which is why Module 2 scales attention scores.

4. The contract of a language model (Essential)

Every model in this course keeps the same promise. Knowing it exactly tells you what to test and what a loss means.

The contract: ids (batch, time) in; logits, one score per vocabulary entry at every position, (batch, time, vocab), out. Softmax turns each position's row into a distribution over the next token. The smallest model that keeps it is an embedding table and a linear head.

Targets are the inputs moved one place left, so 60 characters give 59 training examples, no labels needed. The loss is cross-entropy: minus the log probability of the true next token, averaged.

import math

class TinyLM(nn.Module):
    """The smallest model that keeps the contract: embedding, then a linear head."""
    def __init__(self, vocab=len(VOCAB), width=128):
        super().__init__()
        self.embed = nn.Embedding(vocab, width)
        self.head = nn.Linear(width, vocab, bias=False)
    def forward(self, ids):                        # (B, T) -> (B, T, vocab)
        return self.head(self.embed(ids))

torch.manual_seed(11)
lm = TinyLM()
inp, tgt = ids[:, :-1], ids[:, 1:]                 # targets: the next character
logits = lm(inp)
probs = logits.softmax(-1)
loss = F.cross_entropy(logits.flatten(0, 1), tgt.flatten())
print(tuple(inp.shape), "->", tuple(logits.shape), "| rows sum to", round(float(probs.sum(-1).mean()), 4))
print(repr(decode(inp[0, :6].tolist())), "->", repr(decode(tgt[0, :6].tolist())))
print(f"untrained loss {loss.item():.3f}  vs ln(69) = {math.log(69):.3f}")

Code walkthrough

• TinyLM sees only the current character (a bigram model) but keeps the contract.
• flatten(0, 1) gives cross-entropy 59 separate predictions.
• Guessing evenly among 69 characters scores ln 69 = 4.234 (arithmetic); the untrained model measured 4.355.

GPT-2 keeps the same contract. Up to "on the coolant", it gives " plant" 0.13; the true " loop" gets 0.0007, rank 174 (measured here). A distribution, never an answer.

The third clause breaks silently: position t may use only tokens up to t. Swap the last token and earlier logits must not move. In GPT-2 they move by exactly 0.0; the last moves 9.7. Module 5 makes this a test.

prompt = "Bruno Ferreira works in the water plant on the coolant"
p_ids = tok(prompt, return_tensors="pt").input_ids                    # (1, 15)
with torch.no_grad():
    out = ref(input_ids=p_ids).logits                                 # (1, 15, 50257)
dist = out[0, -1].softmax(-1)
top = dist.topk(3)
print(tuple(out.shape), [(tok.decode(int(i)), round(float(v), 3)) for v, i in zip(top.values, top.indices)])
loop = tok.encode(" loop")[0]
print(f"' loop': {float(dist[loop]):.4f}, rank {int((dist > dist[loop]).sum()) + 1}")

changed = p_ids.clone()
changed[0, -1] = tok.encode(" fan")[0]                                # a different last token
with torch.no_grad():
    out2 = ref(input_ids=changed).logits
print("earlier positions moved", (out2[0, :-1] - out[0, :-1]).abs().max().item(),
      "| last position moved", round((out2[0, -1] - out[0, -1]).abs().max().item(), 1))

Code walkthrough

• out[0, -1] is the last position's row; topk(3) reads the three largest probabilities, and the sum gives the rank of " loop".
• changed swaps the last token, "ant", for " fan"; the comparison checks the causal clause.

5. Chain rule through matrix operations (Essential)

Training is the chain rule applied to matrix products, millions of times. Derive it once and you need never trust a gradient you cannot explain.

Take TinyLM's head: y = xW, x (59, 128), W (128, 69). The loss hands back g, its gradient with respect to y. Entry by entry, dW = xᵀg and dx = gWᵀ. Shortcut: a gradient has its tensor's shape, and only one arrangement of transposes gives it.

x = lm.embed(inp)[0].detach().requires_grad_()            # (59, 128)
W = lm.head.weight.detach().T.clone().requires_grad_()    # (128, 69)
y = x @ W                                                 # (59, 69)
g = torch.randn_like(y)                                   # a made-up upstream gradient
(g * y).sum().backward()                                  # a loss whose dL/dy is exactly g
print("dW: max |autograd - x.T @ g| =", (W.grad - x.T @ g).abs().max().item())
print("dx: max |autograd - g @ W.T| =", (x.grad - g @ W.T).abs().max().item())

Code walkthrough

• requires_grad_() asks autograd, PyTorch's automatic differentiation, to track a tensor.
• (g * y).sum() has derivative exactly g, so only the matrix rule is tested. Both differences print 0.0.

For cross-entropy, the gradient on the logits is probabilities minus a one-hot. The true character gets p - 1, negative, so it is pushed up; every other gets its own p and is pushed down. At position 0, "B" should predict "r", given 0.037.

Then check autograd itself with finite differences: nudge one weight up and down by a millionth and watch the loss.

z = (x.detach() @ W.detach()).requires_grad_()            # logits (59, 69)
F.cross_entropy(z, tgt[0]).backward()
by_hand = (z.softmax(-1) - F.one_hot(tgt[0], 69).float()) / 59   # mean over 59 positions
print("cross-entropy: max |autograd - (p - onehot)/59| =", (z.grad - by_hand).abs().max().item())
t0 = tgt[0, 0]                                            # position 0: input "B", target "r"
print("p('r') =", round(float(z.softmax(-1)[0, t0]), 3), "| 59 x gradient on 'r' =", round(float(59 * z.grad[0, t0]), 3))

x64, W64 = x.detach().double(), W.detach().double().requires_grad_()
loss_of = lambda Wm: F.cross_entropy(x64 @ Wm, tgt[0])
loss_of(W64).backward()
nudge = torch.zeros_like(W64); nudge[5, 7] = 1e-6
numeric = (loss_of(W64.detach() + nudge) - loss_of(W64.detach() - nudge)) / 2e-6
print(f"W[5,7]: autograd {W64.grad[5, 7].item():.10f}  finite difference {numeric.item():.10f}")

Code walkthrough

• by_hand divides by 59 because cross_entropy averages; the difference, 2e-9, is rounding.
• t0 is position 0's target: 59 × gradient = 0.037 - 1 = -0.963.
• .double() uses 64-bit floats so the nudge survives rounding; the two agree to ten decimal places.

In a deep stack the input's gradient is g times every layer's Wᵀ in turn, a product that can vanish or explode. Lesson 2 measures that in a recurrent network.

6. Tensor shapes as the debugging target (Essential)

You will spend more time on shapes than on any idea here. Errors that raise are easy; the expensive kind trains and is wrong.

Ids (1, 59), embeddings (1, 59, 128), logits (1, 59, 69). Cross-entropy wants classes on axis 1, so unflattened logits raise an error; flattened, they work. A shape comment on every line that changes one is the cheapest habit here.

try:
    F.cross_entropy(logits, tgt)                       # (1, 59, 69) and (1, 59): classes on the wrong axis
except RuntimeError as err:
    print("RuntimeError:", err)
print("flattened:", tuple(logits.flatten(0, 1).shape), tuple(tgt.flatten().shape),
      round(F.cross_entropy(logits.flatten(0, 1), tgt.flatten()).item(), 3))

Code walkthrough

• The real message: Expected target size [1, 69], got [1, 59]: PyTorch read 59 as the class count. Flattened, the loss is 4.355 again.

Break it: view where transpose was meant

Andrej Karpathy, in "A Recipe for Training Neural Networks" (2019), names a bug he has met several times: using view instead of transpose or permute, which "inadvertently mix[es] information across the batch dimension". The network, he notes, "will typically still train okay". His test: make the loss depend on example i only, and check that only input i receives gradient.

Average each sequence of the (3, 20, 128) batch over time, via a time-first layout. transpose moves the data; view relabels the same memory, so its "sequence 0" mixes all three. Both return (3, 128); only the gradient test tells them apart.

h = embed(batch).detach().requires_grad_()            # (3, 20, 128)

def mean_right(h):
    return h.transpose(0, 1).mean(0)                  # (20, 3, 128) -> (3, 128)

def mean_wrong(h):
    return h.view(20, 3, 128).mean(0)                 # same shape, scrambled rows

for f in (mean_right, mean_wrong):
    h.grad = None
    out = f(h)
    out[0].sum().backward()                           # a loss that uses example 0 only
    print(f.__name__, tuple(out.shape), "gradient reaching each example:",
          [round(float(h.grad[i].abs().sum()), 1) for i in range(3)])

Code walkthrough

• Both reshape to (20, 3, 128); only transpose keeps each sequence's data together.
• out[0].sum() uses example 0 only. mean_right sends gradient [128.0, 0.0, 0.0]; mean_wrong sends [44.8, 44.8, 38.4].

In production

Karpathy's nanoGPT, a compact GPT implementation, opens its attention with B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) and flattens logits before cross-entropy, as you did. Before optimising, He's advice is to find which cost dominates: compute, memory or overhead.

His test for overhead: if doubling the batch barely changes run time, the hardware is waiting on Python and the framework. Module 5 turns the causal and gradient checks into the lab's tests; Module 9 budgets training runs with the same FLOP formula.

What you added to tinyformer-lab

• The tokeniser and running example, every shape named.
• matmul_flops and median_ms, with your machine's GFLOP/s.
• TinyLM, the smallest model that keeps the contract, and a causal check on GPT-2.
• Three gradient checks and a batch-isolation test for Module 5.

Check yourself

1. How many FLOPs does a (32, 128) by (128, 512) product cost?
2. Why do random dot products get larger as vectors get longer, and why does that matter to softmax?
3. State the three clauses of the language-model contract.
4. For cross-entropy, what is the gradient on the true token's logit, and on the others?
5. Two functions return tensors of the same shape. How can you tell whether one mixes examples across the batch?

1. 2 × 32 × 128 × 512 = 4,194,304.
2. It sums d terms, so its spread grows as the square root of d; large scores saturate softmax.
3. (batch, time) ids in, (batch, time, vocab) logits out; each row's softmax is the next-token distribution; position t sees only tokens up to t.
4. p - 1 on the true token, p on each other, averaged.
5. Put a loss on one example and check only its input gets gradient.

Exercise

Turn the break-it test into a lab tool: batch_isolation_check(f, h) returns True when, for every i, a loss on f(h)[i] sends gradient only to h[i]. Try it first; one solution is below.

def batch_isolation_check(f, h, tol=0.0):
    """True when output i of f depends on input i only, for every i in the batch."""
    for i in range(h.shape[0]):
        h = h.detach().requires_grad_()
        f(h)[i].sum().backward()
        others = torch.cat([h.grad[:i], h.grad[i + 1:]])
        if others.abs().max() > tol:
            return False
    return True

Code walkthrough

• Each pass detaches h so gradients start fresh; others is the gradient that reached other examples, and any non-zero value means mixing.

Check (it must print True False):

print(batch_isolation_check(mean_right, embed(batch)), batch_isolation_check(mean_wrong, embed(batch)))

Code walkthrough

• Runs your check on both averaging functions from the break-it step.

Key terms

Sources

• Horace He, "Making Deep Learning Go Brrrr From First Principles", 2022. https://horace.io/brrr_intro.html
• Andrej Karpathy, "A Recipe for Training Neural Networks", 25 April 2019. https://karpathy.github.io/2019/04/25/recipe/
• Andrej Karpathy, nanoGPT model.py (read 9 October 2026). https://github.com/karpathy/nanoGPT/blob/master/model.py
• OpenAI, GPT-2 small model card, openai-community/gpt2 on Hugging Face (MIT licence, checked October 2026). https://huggingface.co/openai-community/gpt2