Description

Level 2 — Step-by-Step Progression (Following Karpathy's Build Order)

1. Step 1 — Data Prep & Offset Target Tensors (x and y):
• Given a 1D tensor of token IDs, slice random chunks of length block_size ().
• Input x = data[i : i + T] and Target y = data[i + 1 : i + T + 1].
• Crucial Insight: A single chunk of length actually packs separate training examples into one forward pass (predicting from , predicting from , ..., predicting from )!

2. Step 2 — The Simplest Baseline (Bigram Language Model):
• Directly maps token ID to next-token logits via a single nn.Embedding(vocab_size, vocab_size) table. Zero context history, establishes the training loop and baseline loss ( at random init!).

3. Step 3 — Adding Mathematical Self-Attention ("The Mathematical Trick"):
• Introduce causal weighted averaging using torch.tril, then replace uniform weights with data-dependent dot products scaled by .

4. Step 4 — Communication + Computation (The Transformer Block):
• Self-Attention is communication (tokens gather context from earlier tokens).
• The Feed-Forward MLP () is computation (each token independently "thinks about" what it just gathered).

5. Step 5 — Deep Network Stabilization:
• Add Residual Connections (x = x + sublayer(x)), Pre-LayerNorm (ln1, ln2, plus a final ln_f before lm_head), and Dropout so deep stacks converge smoothly.

6. Step 6 — Scaling Up & Hardware Optimizations (nanoGPT Tricks):
• Use F.scaled_dot_product_attention(..., is_causal=True) to invoke FlashAttention, torch.compile(model) for kernel fusion, BF16 mixed precision, and residual projection initialization.

Level 3 — Advanced Architecture Verification Checks (Karpathy's Debugging Rules)

• Sanity Check #1 (Initial Loss at Step 0): Before any gradient step, a well-initialized model should assign roughly uniform probability to every token. Therefore, initial Cross-Entropy loss MUST equal (e.g., for TinyShakespeare with chars, ; for GPT-2 with , ). If initial loss is , your output logits are too large at initialization!

• Sanity Check #2 (Residual Stream Variance Scaling): Because each Transformer block adds two residual branches (), variance in the residual stream grows as . GPT-2 scales the weights of the output projections (c_proj) by at initialization to keep residual variance at across depth.

Learning notes

Complete Data & Tensor Shape Flow Through Tiny GPT

Visualizing How 1 Sequence Chunk (B=1, T=4) Encodes 4 Training Examples

• x = [12, 45, 88, 19]

• y = [45, 88, 19, 73]

Example or implementation

Complete, Self-Contained Tiny GPT in ~100 Lines of PyTorch (Karpathy nanoGPT Style)

from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F

@dataclass
class GPTConfig:
    vocab_size: int = 65      # e.g., TinyShakespeare char vocab (or 50257 for GPT-2)
    block_size: int = 128     # Maximum context length T
    n_layer: int = 4          # Number of Transformer blocks
    n_head: int = 4           # Number of attention heads
    n_embd: int = 128         # Embedding dimension C (d_model)

class CausalSelfAttention(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.size()
        qkv = self.c_attn(x).split(self.n_embd, dim=2)
        q, k, v = [t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) for t in qkv]
        # Uses FlashAttention kernel automatically with is_causal=True!
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)

class MLP(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.c_fc   = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.gelu   = nn.GELU(approximate="tanh")
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.c_proj(self.gelu(self.c_fc(x)))

class Block(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp  = MLP(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln_1(x))  # Communication (Pre-LN + Residual)
        x = x + self.mlp(self.ln_2(x))   # Computation   (Pre-LN + Residual)
        return x

class TinyGPT(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.config = config
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.wpe = nn.Embedding(config.block_size, config.n_embd)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.wte.weight = self.lm_head.weight  # Weight tying!

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
        B, T = idx.size()
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        x = self.wte(idx) + self.wpe(pos)
        for block in self.blocks:
            x = block(x)
        logits = self.lm_head(self.ln_f(x))  # (B, T, vocab_size)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) if targets is not None else None
        return logits, loss

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 1.0) -> torch.Tensor:
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)
            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

# Verify Forward Pass, Expected Initial Loss (~ln(65) = 4.17), and Generation!
cfg = GPTConfig()
model = TinyGPT(cfg)
x = torch.randint(0, cfg.vocab_size, (2, 16))
y = torch.randint(0, cfg.vocab_size, (2, 16))
logits, loss = model(x, y)
print(f"Logits shape: {logits.shape} | Step-0 Loss: {loss.item():.2f} (Theoretical ln(65) = 4.17)")
generated = model.generate(torch.zeros((1, 1), dtype=torch.long), max_new_tokens=10)
print("Generated token IDs:", generated.tolist()[0])

Interview-ready explanation

90-Second Spoken Pitch for Interviews

"If you ask me to code a working GPT from scratch on a whiteboard following Karpathy's nanoGPT, the entire architecture fits into five clean modules.

First, in the Data Loader, we sample batches of token chunks x of shape (B, T) and target chunks y shifted right by token. Thanks to causal masking, a single tensor of shape (B, T) trains the model on distinct context-to-next-token predictions simultaneously.

Second, in the Input Stem, we sum token embeddings wte(idx) of shape (B, T, C) with positional embeddings wpe(pos) of shape (T, C).

Third, the tensor passes through Transformer Blocks. Each block has two residual sub-layers using Pre-LayerNorm: x = x + attn(ln1(x)) where tokens communicate across time via Causal Multi-Head Self-Attention, followed by x = x + mlp(ln2(x)) where each token independently computes features via a GELU feed-forward network.

Fourth, after the final block, we apply a final LayerNorm ln_f and project from C to vocab_size via lm_head (whose weights are tied to wte). During training, we flatten logits to (B*T, V) and compute F.cross_entropy against flattened targets—verifying that step- loss equals . During inference, our generate() loop slices the last token's logits logits[:, -1, :], scales by temperature, applies softmax, samples via torch.multinomial, and appends the new token autoregressively."

Questions to revisit

Q1: Why do we flatten logits from (B, T, V) to (B * T, V) when calling F.cross_entropy in PyTorch?

Q2: During generate(), why do we slice logits[:, -1, :] instead of using all T positions of logits?

Q3: What are the biggest GPU speedups Karpathy adds in his GPT-2 reproduction tutorial?