nanoGPT — the whole GPT,
in one file you can read

00 What a GPT actually does (one sentence)

A GPT is a next-token predictor. You give it a sequence of words and it outputs a probability for every possible next word. That's the whole job. "Writing" is just doing that over and over: predict a token, stick it on the end, predict again.

"Training" is the same idea as Walk 01 — show it real text, measure how wrong its prediction was, and nudge every knob to be a little less wrong. What changes here is theshape of the formula between input and prediction. That shape is the Transformer. Let's read it.

01 The map — one straight pipeline

The entire model is a straight line from token ids to a prediction:

token ids ──► + embeddings ──► Block × N ──► final LayerNorm ──► lm_head ──► logits
[the, cat]     word + position    attention + MLP                  a score for every
                                  (repeated 12×)                   possible next token

Five stages, and only two of them are interesting (embeddings and the Block). Everything inmodel.py is one of these boxes. Read the code for forward and you're reading this diagram:

def forward(self, idx, targets=None):
    b, t = idx.size()                              # idx = token ids, shape (batch, seq_len)
    pos = torch.arange(0, t)                       # positions 0,1,2,...,t-1

    tok_emb = self.transformer.wte(idx)            # look up a vector for each TOKEN
    pos_emb = self.transformer.wpe(pos)            # look up a vector for each POSITION
    x = self.transformer.drop(tok_emb + pos_emb)   # add them: "this word, at this spot"

    for block in self.transformer.h:               # the stack of Transformer blocks
        x = block(x)
    x = self.transformer.ln_f(x)                   # final normalization

idx is a grid of token ids. It becomes vectors, flows through a stack of identicalBlocks, gets a final cleanup, and (next section) turns into scores. Hold this shape in your head; the rest is zoom-ins.

Before we zoom in, watch the whole thing run once. Press play: the panel on the left is the real forward pass, and the line that lights up is the one executing right now. The panel on the right shows what the data is doing at that instant — four tokens turning into vectors, attending to each other, and finally choosing a fifth. Every section below is just one of these eighteen steps, slowed down.

def forward(self, idx):          # idx = (B, T) token ids
    tok = self.wte(idx)          # (B,T,C)  token embeddings
    pos = self.wpe(positions)    # (T,C)    position embeddings
    x   = tok + pos              # the residual stream starts here
 
    for block in self.blocks:    # × 12 identical layers
        h = block.ln_1(x)        # layernorm
        q, k, v = block.attn.c_attn(h).split(C, dim=2)
        att = (q @ k.transpose(-2,-1)) / sqrt(head_size)
        att = att.masked_fill(causal_mask == 0, -inf)
        att = softmax(att, dim=-1)   # attention weights
        y   = att @ v            # weighted sum of values
        x   = x + block.attn.c_proj(y)   # residual add
 
        h = block.ln_2(x)        # layernorm
        h = gelu(block.mlp.c_fc(h))      # up-project 4× + GELU
        x = x + block.mlp.c_proj(h)      # residual add
 
    x = self.ln_f(x)             # final layernorm
    logits = x @ self.wte.weight.T   # tied to the embeddings
    probs  = softmax(logits[:, -1])  # next-token distribution
    return probs

02 Words → vectors (and why position is added)

A computer can't do math on the word "cat." So the first step looks up a learned vector for each token (wte, the token embedding) — a list of numbers that captures its meaning. Similar words end up with similar vectors.

But there's a catch that matters a lot later: attention, by itself, has no sense of order — to it, "dog bites man" and "man bites dog" look identical. So we also look up a vector for eachposition (wpe) and add it in. Now each vector means "this word, at this spot." That's the single line x = tok_emb + pos_emb.

03 The Block — the unit that repeats 12 times

The whole depth of GPT is just this two-line block, stacked over and over:

def forward(self, x):
    x = x + self.attn(self.ln_1(x))   # 1) let tokens share information (attention)
    x = x + self.mlp(self.ln_2(x))    # 2) let each token think on its own (MLP)
    return x

Two ideas, and one trick:

• Attention lets tokens look at each other and share information ("this word depends on that earlier word").
• MLP lets each token think on its own after gathering context.
• The trick is x = x + (...) — the residual connection. Each blockadds its result to the input rather than replacing it. This keeps a clean "highway" for gradients to flow back through (Walk 01!) and is what makes very deep networks trainable.

ln_1 / ln_2 are LayerNorms — they just keep the numbers in a sane range before each step so training stays stable. Applied before the sub-layer ("pre-norm").

04 Attention — the heart of the whole thing

This is the idea that made modern LLMs possible, and it's simpler than its reputation. Every token asks: "of the tokens before me, which ones are relevant, and what do they tell me?"

To do that, each token produces three vectors (all learned projections of itself):

• Query (q) — "what am I looking for?"
• Key (k) — "what do I offer / what am I about?"
• Value (v) — "if you attend to me, here's the information I'll give you."

The mechanism is three steps — and the code is exactly three steps:

def forward(self, x):
    B, T, C = x.size()   # batch, sequence length, embedding dim

    # 1) for every token, produce a query, a key, and a value vector
    q, k, v = self.c_attn(x).split(self.n_embd, dim=2)

    # 2) score: how relevant is each token to each other token?
    att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))   # dot-product similarity
    att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))   # causal: can't see the future
    att = F.softmax(att, dim=-1)                                      # scores → weights that sum to 1

    # 3) mix: each token's output is a weighted blend of the value vectors
    y = att @ v
    y = y.transpose(1, 2).contiguous().view(B, T, C)                 # recombine the heads
    return self.resid_dropout(self.c_proj(y))

Step 3 is worth feeling directly — it's the same softmax the whole model runs. Drag the four scores and watch attention concentrate on one token or spread evenly across all of them. Sharp logits → the token attends to one place; flat logits → it averages everything:

Step 1's q @ kᵀ is a matmul — play with the matmul itself inStudy · Module 1, and the softmax's full story inModule 2.

Note line 64 in the real file: if PyTorch ≥ 2.0 is present, all of this is replaced by one call toscaled_dot_product_attention — the fused FlashAttention CUDA kernel. Same math, one highly-optimized GPU kernel instead of four separate ops. Remember that name; it's Walk 04.

05 The MLP — where each token "thinks"

After attention has gathered context, the MLP processes each token independently. It's deliberately boring: expand to 4× the width, apply a nonlinearity, shrink back.

def forward(self, x):
    x = self.c_fc(x)     # expand:  n_embd → 4 * n_embd
    x = self.gelu(x)     # nonlinearity (a smooth ReLU)
    x = self.c_proj(x)   # project back: 4 * n_embd → n_embd
    return x

A clean way to hold it: attention moves information between tokens; the MLP does computationwithin a token. Alternate the two, twelve times, and you have GPT-2. (Fun fact: the MLPs hold roughly two-thirds of the model's parameters — most of a model's "knowledge" lives here.)

06 Vectors → words, and the loss

After the last block, each position holds a rich vector. lm_head — one big matrix — turns it into logits: one raw score for every token in the vocabulary (~50,000). Softmax those and you have a probability distribution over the next token.

# weight tying: the same matrix maps words→vectors and vectors→words
self.transformer.wte.weight = self.lm_head.weight

# ... at the end of forward():
logits = self.lm_head(x)   # (batch, seq, vocab_size): a score for EVERY possible next token
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

Two things worth noticing. Weight tying (line 1): the same matrix that maps words→vectors at the start is reused to map vectors→words at the end — one matrix, learned once, used both directions. And the loss is cross_entropy: it's high when the model put low probability on the token that actually came next, low when it was confident and right. That single number is what we minimize.

07 Training is the exact same loop as Walk 01

Here's the payoff for reading micrograd first. Once forward gives us thatloss, training nanoGPT is the identical four-step loop — just with a real optimizer (AdamW) instead of hand-written p.data -= lr*p.grad:

1. forward — run the pipeline above, get logits, loss.
2. loss.backward() — the same autograd from Walk 01, now over millions of parameters.
3. optimizer.step() — nudge every weight down its gradient.
4. optimizer.zero_grad() — clear gradients (the += reason), and repeat.

That's it. The Transformer is only a new shape for the formula; the learning machinery underneath is unchanged from the 154 lines you already read.

08 Generation — how it "writes"

Sampling text is a tiny loop: predict, pick a token, append, repeat.

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    for _ in range(max_new_tokens):
        logits, _ = self(idx)              # run the model on what we have so far
        logits = logits[:, -1, :] / temperature   # take the last position's scores
        probs = F.softmax(logits, dim=-1)          # scores → probabilities
        idx_next = torch.multinomial(probs, num_samples=1)   # sample the next token
        idx = torch.cat((idx, idx_next), dim=1)    # append it, and loop
    return idx

• temperature — divides the scores before softmax. Low = safe and repetitive; high = varied and risky.
• top_k — only consider the k most-likely tokens, so it never picks something absurd.
• multinomial — actually samples from the distribution (not always the top pick), which is why the model isn't deterministic.