Understanding Transformers: From First Principles to Production GPTs

# The simplest possible language model: bigrams
import random
from collections import Counter

text = "hello world hello there"
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}  # char -> int
itos = {i: c for c, i in stoi.items()}       # int -> char

# Count every pair of consecutive characters
bigrams = Counter()
for w in text:
    for ch1, ch2 in zip(w, w[1:]):
        bigrams[(ch1, ch2)] += 1

# Now we can "generate" by sampling from the distribution
# after each character. This is a language model!

import torch

# Suppose we have 4 tokens, each represented by a 2D vector
x = torch.tensor([
    [1.0, 0.0],   # token 0
    [0.0, 1.0],   # token 1
    [1.0, 1.0],   # token 2
    [0.5, 0.5],   # token 3
])

T = x.shape[0]  # sequence length

# Version 1: naive averaging with a loop
out = torch.zeros_like(x)
for i in range(T):
    out[i] = x[:i+1].mean(dim=0)  # average all tokens up to position i

print(out)
# token 0 sees only itself
# token 1 sees average of tokens 0,1
# token 2 sees average of tokens 0,1,2
# token 3 sees average of all four

# Version 2: the same thing as a matrix multiply
import torch.nn.functional as F

# Create a lower-triangular matrix of ones
tril = torch.tril(torch.ones(T, T))
# Normalize each row to sum to 1
weights = tril / tril.sum(dim=1, keepdim=True)

print(weights)
# tensor([[1.0000, 0.0000, 0.0000, 0.0000],
#         [0.5000, 0.5000, 0.0000, 0.0000],
#         [0.3333, 0.3333, 0.3333, 0.0000],
#         [0.2500, 0.2500, 0.2500, 0.2500]])

out2 = weights @ x  # matrix multiply!
# out2 is identical to out from the loop version

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

torch.manual_seed(42)

B, T, C = 1, 4, 8  # batch, time (sequence length), channels (embedding dim)
x = torch.randn(B, T, C)

head_size = 4
key   = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)

k = key(x)    # (B, T, head_size)
q = query(x)  # (B, T, head_size)
v = value(x)  # (B, T, head_size)

# Compute attention scores: how much should each token attend to each other?
# This is Q @ K^T — a dot product between every pair of query and key vectors
weights = q @ k.transpose(-2, -1)  # (B, T, T)

# Scale to prevent large values (which would make softmax too "peaky")
weights = weights / math.sqrt(head_size)

# Mask: prevent attending to future tokens (causal / autoregressive)
tril = torch.tril(torch.ones(T, T))
weights = weights.masked_fill(tril == 0, float('-inf'))

# Softmax: convert scores to probabilities (each row sums to 1)
weights = F.softmax(weights, dim=-1)

# Weighted sum of values
out = weights @ v  # (B, T, head_size)

# Demonstrating why scaling matters
d_k = 64
q = torch.randn(1, d_k)
k = torch.randn(100, d_k)

# Without scaling: dot products have high variance
unscaled = q @ k.T
print(f"Unscaled std: {unscaled.std():.2f}")  # ~8.0 for d_k=64

# With scaling: variance ≈ 1
scaled = unscaled / math.sqrt(d_k)
print(f"Scaled std: {scaled.std():.2f}")  # ~1.0

# Softmax on unscaled: almost one-hot (extreme)
print(F.softmax(unscaled, dim=-1).max())  # ~0.99

# Softmax on scaled: smooth distribution (learnable)
print(F.softmax(scaled, dim=-1).max())    # ~0.05

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"

        # Single linear projection for Q, K, V (more efficient than three separate ones)
        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape

        # Project to Q, K, V and reshape into multiple heads
        qkv = self.qkv(x)  # (B, T, 3*C)
        qkv = qkv.reshape(B, T, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, num_heads, T, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]

        # Scaled dot-product attention (per head, in parallel)
        scale = math.sqrt(self.head_dim)
        attn_scores = (q @ k.transpose(-2, -1)) / scale  # (B, H, T, T)

        # Causal mask: prevent attending to future tokens
        mask = torch.tril(torch.ones(T, T, device=x.device))
        attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = F.softmax(attn_scores, dim=-1)  # (B, H, T, T)

        # Apply attention to values
        out = attn_weights @ v  # (B, H, T, head_dim)

        # Concatenate heads and project
        out = out.transpose(1, 2).reshape(B, T, C)  # (B, T, C)
        return self.out_proj(out)

class TransformerBlock(nn.Module):
    """
    A single Transformer block: the fundamental repeating unit.
    Stack N of these and you have GPT.
    """
    def __init__(self, embed_dim: int, num_heads: int, ff_dim: int = None):
        super().__init__()
        ff_dim = ff_dim or 4 * embed_dim  # convention: FFN is 4x wider

        # Layer norms (pre-norm architecture — more stable than post-norm)
        self.ln1 = nn.LayerNorm(embed_dim)
        self.ln2 = nn.LayerNorm(embed_dim)

        # Multi-head self-attention
        self.attn = MultiHeadSelfAttention(embed_dim, num_heads)

        # Feed-forward network (the "thinking" part)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.GELU(),
            nn.Linear(ff_dim, embed_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Attention with residual connection
        # "x + ..." means the original signal passes through unchanged,
        # and attention only adds its contribution on top
        x = x + self.attn(self.ln1(x))

        # MLP with residual connection
        x = x + self.mlp(self.ln2(x))

        return x

# Without residual: signal must survive every transformation
x = layer_3(layer_2(layer_1(x)))  # x is deeply transformed, gradients struggle

# With residual: original signal preserved, layers learn corrections
x = x + layer_1(x)  # layer 1 adds its contribution
x = x + layer_2(x)  # layer 2 adds its contribution
x = x + layer_3(x)  # layer 3 adds its contribution

class SwiGLU(nn.Module):
    """SwiGLU activation: modern replacement for GELU in FFN"""
    def __init__(self, embed_dim, ff_dim):
        super().__init__()
        self.w1 = nn.Linear(embed_dim, ff_dim, bias=False)
        self.w2 = nn.Linear(embed_dim, ff_dim, bias=False)  # gate
        self.w3 = nn.Linear(ff_dim, embed_dim, bias=False)

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

# RMSNorm: simpler and slightly faster than LayerNorm
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

def sinusoidal_positional_encoding(max_len, d_model):
    pe = torch.zeros(max_len, d_model)
    position = torch.arange(0, max_len).unsqueeze(1).float()
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))

    pe[:, 0::2] = torch.sin(position * div_term)  # even dimensions
    pe[:, 1::2] = torch.cos(position * div_term)  # odd dimensions
    return pe

# Simple learned position embeddings (GPT-2 style)
position_embedding = nn.Embedding(block_size, embed_dim)
# Usage: pos_emb = position_embedding(torch.arange(T))

def precompute_rope_freqs(dim, max_len, theta=10000.0):
    """Precompute the complex exponentials for RoPE"""
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_len)
    freqs = torch.outer(t, freqs)  # (max_len, dim/2)
    return torch.polar(torch.ones_like(freqs), freqs)  # complex exponentials

def apply_rope(x, freqs):
    """Apply rotary embeddings to input tensor"""
    # Reshape x into pairs and treat as complex numbers
    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
    # Rotate by the position-dependent frequencies
    x_rotated = x_complex * freqs
    # Convert back to real
    return torch.view_as_real(x_rotated).reshape(*x.shape)

class GPT(nn.Module):
    """
    A complete GPT language model.
    This is essentially what Karpathy builds in nanoGPT, simplified.
    """
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 768,
        num_heads: int = 12,
        num_layers: int = 12,
        block_size: int = 1024,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.block_size = block_size

        # Token and position embeddings
        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
        self.position_embedding = nn.Embedding(block_size, embed_dim)
        self.drop = nn.Dropout(dropout)

        # Stack of Transformer blocks
        self.blocks = nn.Sequential(*[
            TransformerBlock(embed_dim, num_heads)
            for _ in range(num_layers)
        ])

        # Final layer norm and output projection
        self.ln_f = nn.LayerNorm(embed_dim)
        self.lm_head = nn.Linear(embed_dim, vocab_size, bias=False)

        # Weight tying: share weights between token embedding and output projection
        # This is a key insight from the "Using the Output Embedding" paper
        self.token_embedding.weight = self.lm_head.weight

        # Initialize weights (important for stable training!)
        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: torch.Tensor) -> torch.Tensor:
        B, T = idx.shape
        assert T <= self.block_size, f"Sequence length {T} exceeds block size {self.block_size}"

        # Create token + position embeddings
        tok_emb = self.token_embedding(idx)                          # (B, T, C)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))  # (T, C)
        x = self.drop(tok_emb + pos_emb)                            # (B, T, C)

        # Pass through all transformer blocks
        x = self.blocks(x)                                          # (B, T, C)

        # Final layer norm + project to vocabulary
        x = self.ln_f(x)                                            # (B, T, C)
        logits = self.lm_head(x)                                    # (B, T, vocab_size)

        return logits

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """Autoregressive generation: predict one token at a time"""
        for _ in range(max_new_tokens):
            # Crop to block_size if needed
            idx_cond = idx[:, -self.block_size:]

            # Forward pass
            logits = self(idx_cond)

            # Take logits at the last position
            logits = logits[:, -1, :] / temperature

            # Optionally apply 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')

            # Sample from the distribution
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)

            # Append to sequence
            idx = torch.cat([idx, idx_next], dim=1)

        return idx

def train_gpt(model, train_data, epochs=1, batch_size=32, lr=3e-4, block_size=128):
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)

    for epoch in range(epochs):
        # Sample a random batch of training sequences
        ix = torch.randint(len(train_data) - block_size, (batch_size,))
        x = torch.stack([train_data[i:i+block_size] for i in ix])       # inputs
        y = torch.stack([train_data[i+1:i+block_size+1] for i in ix])   # targets (shifted by 1)

        # Forward pass: compute logits
        logits = model(x)  # (B, T, vocab_size)

        # Cross-entropy loss: how wrong are our predictions?
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)),  # flatten to (B*T, vocab_size)
            y.view(-1),                         # flatten to (B*T,)
        )

        # Backward pass: compute gradients
        optimizer.zero_grad()
        loss.backward()

        # Gradient clipping: prevent exploding gradients
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        # Update parameters
        optimizer.step()

        if epoch % 100 == 0:
            print(f"step {epoch}: loss {loss.item():.4f}")

# The core idea: wrap scalars to track gradients
class Value:
    def __init__(self, data, children=(), local_grads=()):
        self.data = data
        self.grad = 0
        self._children = children
        self._local_grads = local_grads

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data + other.data, (self, other), (1, 1))

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data * other.data, (self, other), (other.data, self.data))
    # ... more operations, then backward() traverses the graph

# In PyTorch 2.0+, Flash Attention is a one-liner:
from torch.nn.functional import scaled_dot_product_attention

# This automatically uses Flash Attention when available
out = scaled_dot_product_attention(q, k, v, is_causal=True)

# GQA: Grouped-Query Attention (Llama 2 style)
# Instead of n_heads KV pairs, use n_kv_heads (e.g., 8 instead of 32)
class GroupedQueryAttention(nn.Module):
    def __init__(self, embed_dim, n_heads, n_kv_heads):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_groups = n_heads // n_kv_heads  # queries per KV group
        self.head_dim = embed_dim // n_heads

        self.wq = nn.Linear(embed_dim, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(embed_dim, n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(embed_dim, n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(n_heads * self.head_dim, embed_dim, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        # Repeat K, V for each group of query heads
        k = k.repeat_interleave(self.n_groups, dim=1)  # (B, n_heads, T, head_dim)
        v = v.repeat_interleave(self.n_groups, dim=1)

        # Standard scaled dot-product attention from here
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).reshape(B, T, -1)
        return self.wo(out)

# Without KV cache: recompute everything at each step (slow)
for i in range(max_tokens):
    logits = model(full_sequence)  # processes ALL tokens every time
    next_token = sample(logits[:, -1, :])
    full_sequence = torch.cat([full_sequence, next_token], dim=1)

# With KV cache: only compute the new token (fast)
kv_cache = None
for i in range(max_tokens):
    logits, kv_cache = model(new_token_only, kv_cache=kv_cache)
    next_token = sample(logits)