07_mini_gpt_forward.py

"""Challenge 07 - Mini GPT Forward Pass (Hard)

PROBLEM
-------
Implement the full forward pass of a decoder-only transformer (GPT-style) in
pure numpy. No autograd, no training - just inference-time math.

Architecture (GPT-2 style, pre-norm):

    token embeddings + learned position embeddings
    -> N x [ LayerNorm -> causal multi-head self-attention -> residual add
             LayerNorm -> MLP (4x expansion, GELU)         -> residual add ]
    -> final LayerNorm
    -> logits via tied embedding matrix (x @ W_emb.T)

Required signatures:

    init_params(cfg: GPTConfig, seed: int = 0) -> dict
    gpt_forward(params: dict, cfg: GPTConfig, token_ids, return_attn=False)
        token_ids: int array of shape (batch, seq_len)
        returns logits of shape (batch, seq_len, vocab_size)
        (and, if return_attn, a list of per-layer attention weight tensors
        of shape (batch, n_heads, seq_len, seq_len))

Constraints:
- Tiny config: vocab 64, d_model 32, 2 layers, 4 heads, max_seq_len 16.
- Attention must be causal: position i may attend only to positions <= i.
- Softmax must be numerically stable (subtract the row max).

INTERVIEW NOTES
---------------
A strong solution demonstrates:
- Correct pre-norm residual wiring (norm INSIDE the residual branch; post-norm
  is the original 2017 layout and a common mix-up).
- Head split/merge done with reshape+transpose, not Python loops over heads.
- Scaling scores by 1/sqrt(head_dim), not 1/sqrt(d_model).
- The causality argument: masking scores (not weights) before softmax so each
  row renormalizes over the visible prefix only.
Common mistakes: forgetting the mask entirely (tests below catch it), masking
with 0 instead of -inf/-1e9, applying LayerNorm over the sequence axis instead
of the feature axis, and unstable softmax overflowing in float32.
Follow-up variations: add a KV cache and show decode is O(T) per step; swap
learned positions for RoPE; make attention weights optional to save memory;
implement grouped-query attention (fewer K/V heads than Q heads).
"""

from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class GPTConfig:
    vocab_size: int = 64
    d_model: int = 32
    n_layers: int = 2
    n_heads: int = 4
    max_seq_len: int = 16


def layer_norm(x: np.ndarray, g: np.ndarray, b: np.ndarray, eps: float = 1e-5) -> np.ndarray:
    mu = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    return g * (x - mu) / np.sqrt(var + eps) + b


def gelu(x: np.ndarray) -> np.ndarray:
    # tanh approximation, as used in GPT-2.
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))


def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)  # stability: exp of <=0 only
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)


def init_params(cfg: GPTConfig, seed: int = 0) -> dict:
    rng = np.random.default_rng(seed)

    def w(*shape: int) -> np.ndarray:
        return rng.standard_normal(shape) * 0.02

    def ln() -> dict:
        return {"g": np.ones(cfg.d_model), "b": np.zeros(cfg.d_model)}

    blocks = []
    for _ in range(cfg.n_layers):
        blocks.append({
            "ln1": ln(),
            "attn": {
                "w_qkv": w(cfg.d_model, 3 * cfg.d_model),
                "b_qkv": np.zeros(3 * cfg.d_model),
                "w_o": w(cfg.d_model, cfg.d_model),
                "b_o": np.zeros(cfg.d_model),
            },
            "ln2": ln(),
            "mlp": {
                "w1": w(cfg.d_model, 4 * cfg.d_model),
                "b1": np.zeros(4 * cfg.d_model),
                "w2": w(4 * cfg.d_model, cfg.d_model),
                "b2": np.zeros(cfg.d_model),
            },
        })
    return {"wte": w(cfg.vocab_size, cfg.d_model), "wpe": w(cfg.max_seq_len, cfg.d_model),
            "blocks": blocks, "ln_f": ln()}


def causal_self_attention(x: np.ndarray, p: dict, n_heads: int) -> tuple[np.ndarray, np.ndarray]:
    """Returns (output, attention_weights). x: (B, T, D)."""
    B, T, D = x.shape
    head_dim = D // n_heads

    qkv = x @ p["w_qkv"] + p["b_qkv"]              # (B, T, 3D)
    q, k, v = np.split(qkv, 3, axis=-1)

    def split_heads(t: np.ndarray) -> np.ndarray:  # (B, T, D) -> (B, H, T, hd)
        return t.reshape(B, T, n_heads, head_dim).transpose(0, 2, 1, 3)

    q, k, v = split_heads(q), split_heads(k), split_heads(v)

    scores = q @ k.transpose(0, 1, 3, 2) / np.sqrt(head_dim)   # (B, H, T, T)
    future = np.triu(np.ones((T, T), dtype=bool), k=1)         # True above diagonal
    scores = np.where(future, -1e9, scores)                    # mask BEFORE softmax
    weights = softmax(scores, axis=-1)

    out = weights @ v                                          # (B, H, T, hd)
    out = out.transpose(0, 2, 1, 3).reshape(B, T, D)           # merge heads
    return out @ p["w_o"] + p["b_o"], weights


def gpt_forward(params: dict, cfg: GPTConfig, token_ids, return_attn: bool = False):
    ids = np.asarray(token_ids)
    B, T = ids.shape
    assert T <= cfg.max_seq_len, "sequence longer than max_seq_len"

    x = params["wte"][ids] + params["wpe"][:T]     # (B, T, D)
    attn_maps = []
    for blk in params["blocks"]:
        a, w = causal_self_attention(layer_norm(x, **blk["ln1"]), blk["attn"], cfg.n_heads)
        x = x + a                                  # pre-norm residual
        h = layer_norm(x, **blk["ln2"])
        m = blk["mlp"]
        x = x + gelu(h @ m["w1"] + m["b1"]) @ m["w2"] + m["b2"]
        attn_maps.append(w)

    x = layer_norm(x, **params["ln_f"])
    logits = x @ params["wte"].T                   # weight tying
    return (logits, attn_maps) if return_attn else logits


if __name__ == "__main__":
    cfg = GPTConfig()
    params = init_params(cfg, seed=0)
    rng = np.random.default_rng(42)
    B, T = 2, 10
    ids = rng.integers(0, cfg.vocab_size, size=(B, T))

    # 1. Output shape.
    logits, attn_maps = gpt_forward(params, cfg, ids, return_attn=True)
    assert logits.shape == (B, T, cfg.vocab_size)
    assert len(attn_maps) == cfg.n_layers
    assert attn_maps[0].shape == (B, cfg.n_heads, T, T)

    # 2. Attention softmax rows sum to 1, and future positions get zero weight.
    for w in attn_maps:
        assert np.allclose(w.sum(axis=-1), 1.0)
        assert np.allclose(w[..., np.triu_indices(T, k=1)[0], np.triu_indices(T, k=1)[1]], 0.0)

    # 3. Causality: changing a FUTURE token must not change past logits.
    t = 6
    ids2 = ids.copy()
    ids2[:, t] = (ids2[:, t] + 1) % cfg.vocab_size
    logits2 = gpt_forward(params, cfg, ids2)
    assert np.allclose(logits[:, :t], logits2[:, :t], rtol=0.0, atol=1e-12), \
        "past logits changed when a future token changed - attention is not causal"
    # ...and the changed position itself SHOULD produce different logits.
    assert not np.allclose(logits[:, t], logits2[:, t])

    # 4. Next-token distribution is a valid probability distribution.
    probs = softmax(logits, axis=-1)
    assert np.allclose(probs.sum(axis=-1), 1.0)
    assert (probs >= 0).all()

    # 5. Works at max context length.
    ids_full = rng.integers(0, cfg.vocab_size, size=(1, cfg.max_seq_len))
    assert gpt_forward(params, cfg, ids_full).shape == (1, cfg.max_seq_len, cfg.vocab_size)

    print("All tests passed.")