DL S040 · nanoGPT Rebuilt from Scratch

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 40 / 80 · Module M07 · ~2.3 hours

The story — the commit that launched a thousand GPUs

On January 8, 2023, Andrej Karpathy pushed the first commit to a new repo called nanoGPT. The commit message: "init: first version of nanoGPT." 306 lines in model.py, 275 lines in train.py, plus a README. Two weeks earlier he'd released the "Let's build GPT" YouTube video (2 hours, 55 million views as of 2026). nanoGPT was that video committed to git.

By the end of 2023, nanoGPT had 30k stars, ~200 forks that trained real models, and one particularly famous one: llama.cpp (Georgi Gerganov) started as a nanoGPT-style port of Llama inference to pure C. The modded-nanoGPT speedrun (Keller Jordan, 2024–2025) turned it into an ongoing competition to reduce GPT-2 (124M) training wall-clock on a single 8×H100 node — the record fell from 45 minutes (Karpathy's own baseline) to under 3 minutes by mid-2025 through a cascade of micro-optimisations: Muon optimizer, FP8 attention, custom Triton kernels, rotary variants. The lineage of every one of those hacks starts here, in the file you're about to type out.

A quick tour of the nanoGPT commit history worth reading yourself (git log --oneline in your fork):

• Jan 8 2023 — initial commit. post-LN, learned PE, three separate Q/K/V linears. It was already good.
• Jan 12 2023 — flash attention. Karpathy swapped manual attention for torch.nn.functional.scaled_dot_product_attention a week after PyTorch 2.0's preview added it. Same day: fused QKV projection.
• Feb 2 2023 — weight tying + scaled residual init. The two lines that made deep runs stable and cut params by 30%.
• Feb 4 2023 — configure_optimizers with weight-decay-on-2D-only. Straight lift from T5's tricks section, ~0.02 val-loss improvement.
• Apr 2023 — bfloat16 + torch.compile. The 2× throughput jump that made single-GPU GPT-2 replication a weekend project.

Every one of those five commits corresponds to a concrete lesson you're about to internalise. The rule for this session: you type the code, not copy-paste. Muscle memory is the goal. When you can, without notes, produce the Block class, the CausalSelfAttention class, and the GPT wrapper with the config-dataclass and the generate loop, you have earned "I understand transformers" as a phrase you can say without qualification.

We'll walk through the architecture (mostly review of S037–S039), then focus on the things nanoGPT does differently from the vanilla version in your head: fused QKV projection, causal mask as a registered buffer, weight tying, torch.nn.functional.scaled_dot_product_attention for flash-attention integration, configure_optimizers with weight decay for 2-D params only, and the sampling loop with top-k and temperature.

Prerequisites

• S036–S039 — all the pieces.
• Basic PyTorch training loop fluency (S018).

1 · The config

Modern PyTorch code uses a dataclass for hyperparameters. This is Karpathy's — copy it, and every subsequent piece slots in.

from dataclasses import dataclass
 
@dataclass
class GPTConfig:
    block_size: int = 1024        # max sequence length
    vocab_size: int = 50304       # rounded up to a multiple of 64 for GPU perf
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = False            # LayerNorm + Linear bias — False is faster and slightly better

vocab_size = 50304 looks weird — real GPT-2 has 50257 tokens. Rounded up to the next multiple of 64 for tensor-core-friendly matmul shapes. Zero-cost cheat you should copy.

bias: bool = False — every Linear and LayerNorm in the model runs bias-free. Marginally faster, marginally better val loss. Standard practice post-2022.

2 · CausalSelfAttention — the fused version

Instead of three separate nn.Linear(n_embd, n_embd) for Q, K, V, nanoGPT uses ONE nn.Linear(n_embd, 3 * n_embd) and splits the output. Faster: one big matmul beats three small ones on GPU.

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        # fused QKV
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
 
        # causal mask registered as buffer (moves with .to(device), doesn't train)
        self.register_buffer("bias",
            torch.tril(torch.ones(config.block_size, config.block_size))
                 .view(1, 1, config.block_size, config.block_size))

Then the forward:

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)                           # (B, T, 3C)
        q, k, v = qkv.split(self.n_embd, dim=2)        # each (B, T, C)
        d_k = C // self.n_head
        q = q.view(B, T, self.n_head, d_k).transpose(1, 2)   # (B, h, T, d_k)
        k = k.view(B, T, self.n_head, d_k).transpose(1, 2)
        v = v.view(B, T, self.n_head, d_k).transpose(1, 2)
 
        # scaled dot-product attention with causal mask
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(d_k))    # (B, h, T, T)
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v                                                  # (B, h, T, d_k)
 
        y = y.transpose(1, 2).contiguous().view(B, T, C)             # merge heads
        return self.resid_dropout(self.c_proj(y))

Study qkv.split(self.n_embd, dim=2) — one line to unfused Q/K/V from a (B, T, 3C) tensor. This is the single biggest speedup vs a naive three-projection implementation.

2.1 · The flash-attention swap-in

Post-PyTorch-2.0, replace the manual attention math with one call:

        y = F.scaled_dot_product_attention(q, k, v,
                                           attn_mask=None,
                                           dropout_p=self.dropout if self.training else 0.0,
                                           is_causal=True)

PyTorch dispatches to flash-attention (if your GPU supports it — Ampere and newer), memory-efficient attention, or fused math attention, in that priority order. is_causal=True skips building the mask tensor entirely.

3–5× speedup on long sequences, 5–10× memory reduction. We'll unpack the flash-attention trick in S043.

3 · The block

Pre-LN, exactly as in S039.

class LayerNorm(nn.Module):
    """LayerNorm with optional bias. PyTorch's own LN can't disable bias."""
    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
 
    def forward(self, x):
        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)
 
class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu = nn.GELU()
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)
 
    def forward(self, x):
        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))
 
class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
        self.mlp = MLP(config)
 
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

Nothing new — this is S039 with nanoGPT's naming (ln_1, ln_2, c_fc, c_proj, mlp — chosen to match OpenAI's original GPT-2 checkpoint parameter names so you can state_dict load).

4 · The GPT wrapper

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = LayerNorm(config.n_embd, bias=config.bias),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
 
        # weight tying
        self.transformer.wte.weight = self.lm_head.weight
 
        # Init all weights
        self.apply(self._init_weights)
        # Special scaled init on the residual projections (per GPT-2 paper)
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))
 
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

Two subtleties to internalise:

• Weight tying — self.transformer.wte.weight = self.lm_head.weight makes the input embedding and output head share the same tensor. Saves vocab_size · n_embd params (~38M for GPT-2) and empirically improves loss.
• Scaled residual init — the residual-path Linear (c_proj) is initialised with std shrunk by 1/sqrt(2 · n_layer). This keeps the residual stream's variance stable as depth grows. Skip this and deep models diverge in the first hundred steps.

Now the forward:

    def forward(self, idx, targets=None):
        device = idx.device
        B, T = idx.size()
        assert T <= self.config.block_size
        pos = torch.arange(0, T, dtype=torch.long, device=device)
 
        tok_emb = self.transformer.wte(idx)      # (B, T, C)
        pos_emb = self.transformer.wpe(pos)      # (T, C)
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
 
        if targets is not None:
            logits = self.lm_head(x)             # (B, T, vocab)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                   targets.view(-1), ignore_index=-1)
            return logits, loss
        else:
            # inference-time: only need the LAST position's logits
            logits = self.lm_head(x[:, [-1], :]) # (B, 1, vocab)
            return logits, None

Note the inference optimisation: during generation you only need x[:, [-1], :] since you're predicting one token. Computing logits for all T positions wastes vocab_size × T FLOPs per step.

5 · The optimizer — weight decay on 2-D params only

    def configure_optimizers(self, weight_decay, learning_rate, betas):
        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
        optim_groups = [
            {'params': decay_params, 'weight_decay': weight_decay},
            {'params': nodecay_params, 'weight_decay': 0.0}]
        return torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=True)

Why? Weight decay on 1-D params (LayerNorm gains, biases, embeddings' bias-like behaviour) tends to hurt. It's a widely-known trick from the T5 paper and nanoGPT bakes it in. Copy this exactly for any transformer you train.

6 · Sampling — temperature + top-k

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            # crop context if it exceeds block_size
            idx_cond = idx if idx.size(1) <= self.config.block_size \
                       else idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)                    # (B, 1, vocab)
            logits = logits[:, -1, :] / temperature        # (B, vocab)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

The three sampling knobs you get for free:

• Temperature — logits / T. Low T (e.g. 0.3) → deterministic, safe, boring. High T (e.g. 1.3) → creative, diverse, sometimes incoherent. T = 1.0 → sample from the raw distribution.
• Top-k — keep only the top k most likely tokens, zero out the rest, renormalise. Prevents rare-junk tokens from ever being picked. Classic value: 40.
• Greedy — argmax instead of multinomial. Deterministic. Boring but reproducible. Useful for debugging.

We'll add top-p (nucleus) sampling in S063.

7 · Parameter count sanity check

For GPTConfig(n_layer=6, n_head=6, n_embd=384, block_size=256, vocab_size=65) (Tiny Shakespeare char-level from S041):

• Embed: 65 · 384 + 256 · 384 = 25K + 98K = 123K
• Per block: 4 · 384² + 2 · 384 · 1536 = 590K + 1180K = 1.77M. Times 6 = 10.6M.
• Head (tied) = 0.
• Total ~10.7M.

That's the size we'll train on Tiny Shakespeare tomorrow. Trains to val loss ~1.48 in ~15 minutes on an M2 Mac, ~3 minutes on an A100.

cfg = GPTConfig(n_layer=6, n_head=6, n_embd=384, block_size=256, vocab_size=65)
 
model_tied  = GPT(cfg)                          # tying on
# In another instance, comment out the tying line
 
print("tied  params:", sum(p.numel() for p in model_tied.parameters()) / 1e6, "M")
# Untied should be ~25K larger (65 * 384) — the vocab × d_model of the extra head

8 · What nanoGPT does that we haven't done

Full nanoGPT has a few more tricks we've skipped for clarity:

9 · Common bugs

10 · Modern-2025 twist — the nanoGPT speedrun

The modded-nanoGPT leaderboard is the best free master-class in modern LLM optimisation. Each PR shaves seconds off GPT-2 (124M) training-to-3.28-val-loss on 8×H100, with the record trajectory looking like:

Every entry is <300 diff lines from vanilla nanoGPT. The educational value is enormous: read the PR descriptions, they're mini-papers.

Other nanoGPT-adjacent projects to know in 2026:

• nanoGPT-speedrun-jax — the JAX/TPU port, ~2× faster on TPU v5e.
• nanoDeepSeek (community, 2025) — same 300-line spirit but with MLA + MoE + multi-token prediction from DeepSeek-V3.
• llm.c — Karpathy's 2024 pure-C port. Same architecture, 3000 lines, no PyTorch dependency. Runs GPT-2 training on any CUDA GPU.

Further reading:

• Karpathy's "State of GPT" talk, Microsoft Build 2023 — the philosophical companion to nanoGPT.
• modded-nanoGPT README — pinned leaderboard.

Retention scaffold

One-line summary (write it in your own words): _______________________________________________

Spaced review: re-read §2 (fused QKV) and §5 (configure_optimizers) in 24 hours. On day 7, type the entire model.py from memory as a self-test.

Next session (S041): hook this model up to a training loop, feed it Tiny Shakespeare (1MB of the Bard), and watch it go from producing garbage to producing plausible pseudo-Shakespearean pastiche in ~15 minutes on a laptop GPU.

Sticky note (keep on your desk): nanoGPT = pre-LN decoder-only + fused QKV + weight tying + scaled residual init + F.scaled_dot_product_attention. 300 lines. Every 2025 open model is a small delta from this.

Previous: ← DL S039 · Next: DL S041 →