Machine Learning

Let's Code a GPT2

3 min read
559 words
# deep-learning

Hyperparameters

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

@dataclass
class GPTConfig:
    block_size: int = 1024    # Maximum sequence length (context window)
    vocab_size: int = 50257   # GPT-2 vocab size (Byte-Pair Encoding)
    n_layer: int = 12         # Number of Transformer blocks
    n_head: int = 12          # Number of attention heads
    n_embd: int = 768         # Embedding dimension (768 / 12 = 64 per head)
    dropout: float = 0.1      # Dropout probability

Causal Multi-Head Self-Attention

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.head_dim = config.n_embd // config.n_head

        # We combine Q, K, V projections into a single linear layer for speed
        # Input: (B, T, n_embd) -> Output: (B, T, 3 * n_embd)
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        
        # Output projection back to the residual stream
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        
        # Regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        # Causal mask: lower triangular matrix of ones
        # Registered as buffer so it's not a parameter, but tracked in model state
        self.register_buffer(
            "bias",
            torch.tril(torch.ones(config.block_size, config.block_size))
                 .view(1, 1, config.block_size, config.block_size)
        )

torch.ones(T,T) create a T × T matrix filled with 1s. torch.tril() stands for triangular lower. It zeros out everything above main diagonal.

Why 4 dimensions (1, 1, block_size, block_size)? Our attention tensor has shape (Batch, Heads, T, T). Because of PyTorch broadcasting, a mask of shape (1, 1, T, T) automatically broadcasts across all batches and all heads without copying memory.

Why register_buffer? A buffer is tensor state that is not a trainable parameter (it has no gradients, requires_grad=False\text{requires\_grad}=\text{False}), but should automatically move to the GPU when you do model.to('cuda') and be saved with torch.save(model.state_dict()).

Why is it called bias? A mask isn’t a bias! Why name the variable bias? This is a quirk of history. In OpenAI’s original TensorFlow release of GPT-2, they created this mask variable and named it bias (or b). When researchers ported GPT-2 to PyTorch (notably in early Hugging Face and Karpathy’s minGPT), they kept the name “bias” so that the keys in their PyTorch state_dict matched OpenAI’s pre-trained checkpoint weights exactly.

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

        # 1. Project to Q, K, V in one batch
        qkv = self.c_attn(x) # (B, T, 3 * C)
        q, k, v = qkv.split(self.n_embd, dim=2) # Each is (B, T, C)

        # 2. Reshape into multi-head format: (B, n_head, T, head_dim)
        # Transpose so that head dimension is processed in parallel
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # 3. Scaled dot-product attention
        # (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        
        # Apply causal mask: mask out positions where bias == 0
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))

Distribute lower-triangular mask on demand according to T.

        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)

        # 4. Weighted aggregation of values
        # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = att @ v

        # 5. Re-assemble all head outputs side-by-side
        # (B, nh, T, hs) -> (B, T, nh, hs) -> (B, T, C)
        y = y.transpose(1, 2).contiguous().view(B, T, C)

        # 6. Final projection back to residual stream
        y = self.resid_dropout(self.c_proj(y))
        return y

MLP

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') # GPT-2 historically used the tanh approximation
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x

Transformer Block

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):
        # Attention sub-layer with Pre-LN and residual stream
        x = x + self.attn(self.ln_1(x))
        # MLP sub-layer with Pre-LN and residual stream
        x = x + self.mlp(self.ln_2(x))
        return x

GPT2

class GPT(nn.Module):
    def __init__(self, config: GPTConfig):
        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 = nn.LayerNorm(config.n_embd),
        ))
        # Projection from latent space back to vocabulary logits
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Weight tying: share parameters between input embeddings and lm_head
        self.transformer.wte.weight = self.lm_head.weight

        # Initialize weights following GPT-2 scheme
        self.apply(self._init_weights)

    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)

    def forward(self, idx, targets=None):
        device = idx.device
        B, T = idx.size()
        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is {self.config.block_size}"

        # 1. Generate position indices [0, 1, ..., T-1]
        pos = torch.arange(0, T, dtype=torch.long, device=device) # shape (T)

        # 2. Combine embeddings
        tok_emb = self.transformer.wte(idx) # (B, T, n_embd)
        pos_emb = self.transformer.wpe(pos) # (T, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb)

        # 3. Pass through all Transformer blocks
        for block in self.transformer.h:
            x = block(x)

        # 4. Final LayerNorm
        x = self.transformer.ln_f(x)

        # 5. Compute Logits and Loss (if targets are provided)
        if targets is not None:
            logits = self.lm_head(x) # (B, T, vocab_size)
            # Flatten to compute Cross Entropy loss: (B*T, vocab_size) vs (B*T)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            return logits, loss
        else:
            # Inference optimization: compute logits only for the last position
            logits = self.lm_head(x[:, [-1], :]) # (B, 1, vocab_size)
            return logits, None
@torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        idx is (B, T) tensor of indices in the current context
        """
        for _ in range(max_new_tokens):
            # Crop context if it exceeds maximum block_size
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            
            # Forward the model to get logits
            logits, _ = self(idx_cond) # (B, 1, vocab_size)
            logits = logits[:, -1, :] / temperature # Scale by temperature

            # Optional Top-K filtering
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')

            # Softmax to get probabilities
            probs = F.softmax(logits, dim=-1)
            
            # Sample next token
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            
            # Append to running sequence
            idx = torch.cat((idx, idx_next), dim=1)

        return idx
if __name__ == "__main__":
    # Use a miniature test configuration
    test_config = GPTConfig(
        block_size=64,
        vocab_size=1000,
        n_layer=4,
        n_head=4,
        n_embd=128
    )
    
    model = GPT(test_config)
    print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")

    # Dummy batch of 2 sequences, each of length 8 tokens
    dummy_input = torch.randint(0, 1000, (2, 8))
    dummy_targets = torch.randint(0, 1000, (2, 8))

    # Test forward pass with loss
    logits, loss = model(dummy_input, targets=dummy_targets)
    print(f"Logits shape: {logits.shape}")  # Expected: [2, 8, 1000]
    print(f"Calculated loss: {loss.item():.4f}")

    # Test generation
    prompt = torch.tensor([[10, 20, 30]]) # Batch of 1, 3 tokens
    generated = model.generate(prompt, max_new_tokens=5)
    print(f"Generated sequence shape: {generated.shape}") # Expected: [1, 8]
    print("Generated token IDs:", generated.tolist()[0])

Causal Mask

Step 1: The Raw Attention Matrix (Before Masking)

Let’s demystify exactly how causal masking works mathematically and in code.

We will trace a tiny example with a 3-token sentence: ["The", "robot", "speaks"] (T=3T = 3).

When we compute the dot product between queries and keys: scores=QKTdk\text{scores} = \frac{Q K^T}{\sqrt{d_k}}

We get a square matrix of size (T,T)(T, T)—in our case, (3×3)(3 \times 3).

Here, Rows = Queries (Tokens doing the looking) and Columns = Keys (Tokens being looked at):

                  Keys:  "The" (k0)   "robot" (k1)   "speaks" (k2)
Queries:
"The"    (q0)  ──►      [  2.3             1.1            0.5   ]
"robot"  (q1)  ──►      [  0.8             3.2            1.4   ]
"speaks" (q2)  ──►      [  1.2             0.9            2.7   ]

Look at the first row ("The"):

  • It has scores for "The" (allowed), "robot" (future!), and "speaks" (future!).
  • In this raw state, every token can see every other token.

Step 2: The Mask Pattern (Lower Triangular Matrix)

To block the future, we define a rule: A token at position ii can only look at tokens at position jj if j≤ij \le i.

In PyTorch, we build this using torch.tril (triangular lower):

mask = torch.tril(torch.ones(3, 3))

This creates a grid of 1s and 0s:

        Col 0   Col 1   Col 2
Row 0 [   1       0       0   ]   <-- Token 0 can only see Token 0
Row 1 [   1       1       0   ]   <-- Token 1 can see Tokens 0, 1
Row 2 [   1       1       1   ]   <-- Token 2 can see Tokens 0, 1, 2

Any cell with a 0 is a peek into the future that must be erased.

Step 3: The Mathematical Trick: Why −∞-\infty?

Why can’t we just replace the forbidden scores with 00?

Remember how attention probabilities are calculated: through the Softmax function: softmax(z)i=ezi∑jezj\text{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}

Let’s see what happens if we replace the forbidden values with 00 versus −∞-\infty:

If we put 0: e0=1e^0 = 1 If a raw score becomes 00, its numerator is still 11! The model will allocate probability mass to future tokens: softmax([2.3,0,0])=[e2.3e2.3+1+1,1e2.3+1+1,1e2.3+1+1]≈[0.83,0.08,0.08]\text{softmax}([2.3, \mathbf{0}, \mathbf{0}]) = \left[\frac{e^{2.3}}{e^{2.3} + 1 + 1}, \frac{1}{e^{2.3} + 1 + 1}, \frac{1}{e^{2.3} + 1 + 1}\right] \approx [0.83, \mathbf{0.08}, \mathbf{0.08}] The model is still cheating!

If we put −∞-\infty: lim⁡z→−∞ez=0\lim_{z \to -\infty} e^z = 0 Because e−∞=0e^{-\infty} = 0, that position completely vanishes: softmax([2.3,−∞,−∞])=[e2.3e2.3+0+0,0e2.3+0+0,0e2.3+0+0]=[1.0,0.0,0.0]\text{softmax}([2.3, \mathbf{-\infty}, \mathbf{-\infty}]) = \left[\frac{e^{2.3}}{e^{2.3} + 0 + 0}, \frac{0}{e^{2.3} + 0 + 0}, \frac{0}{e^{2.3} + 0 + 0}\right] = [1.0, \mathbf{0.0}, \mathbf{0.0}]

Future tokens get strictly zero probability.

Layer Normalization

In modern Transformer(GPT2), the LayerNorm is placed before sub-layers (called PreLN)

Why is this so powerful? Because of the residual connection, there is an uninterrupted “gradient highway” from the very last layer straight back to the input embeddings.

The LayerNorm ensures that the inputs to the Attention and MLP blocks are always well-scaled, preventing gradient explosions in deep architectures (e.g., 32 to 96 layers deep).

If you inspect modern architectures like LLaMA, Mistral, or Gemma, you might notice they replaced nn.LayerNorm with RMSNorm (Root Mean Square Normalization).