Transformer Bug Hunt: Fix a Broken Mini-GPT - End-to-End Implementations

Implement Transformer Bug Hunt: Fix a Broken Mini-GPT in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

Transformer Bug Hunt: Fix a Broken Mini-GPT

What you'll build

A working gpt_forward by debugging a broken one. Reading and fixing someone else's model code is a common interview format, and these four bugs are the classic silent ones: the code runs and the shapes are right, but the numbers are wrong.

The setup

The starter code is a tiny GPT-style language model: token and position embeddings, one causal self-attention block, one MLP block, a final LayerNorm and an LM head. It runs without errors on short inputs, but its outputs are wrong.

There are exactly 4 bugs, all inside gpt_forward. The helpers init_params, layer_norm, softmax and gelu are correct. Find and fix the bugs; you shouldn't need to rewrite the model.

Reference architecture (what the code should compute)

For token ids tokens of length T (positions t = 0 .. T-1), with d = d_model:

x   = wte[tokens] + wpe[t]                               # learned positions, t starts at 0

h   = LN1(x)
q, k, v = h @ W_q, h @ W_k, h @ W_v
S   = q @ k.T / sqrt(d)
S[i, j] = -inf  for j > i                                # causal mask, applied BEFORE the softmax
att = softmax(S)                                         # row-wise, each row sums to 1
x   = x + (att @ v) @ W_o                                # residual 1

x   = x + GELU(LN2(x) @ W_1 + b_1) @ W_2 + b_2           # residual 2 (pre-LN)

logits = LN_f(x) @ wte.T                                 # tied LM head

LN(x) = g * (x - mean) / sqrt(var + 1e-5) + b over the last axis (population variance), and GELU uses the tanh approximation.

Expected Return Format

gpt_forward(tokens, params) returns a tuple (logits, att):

logits: (T, vocab_size) array
att: (T, T) attention probabilities, row i holds query i's weights

How it is tested

The tests check properties every correct decoder must have (causality, attention rows summing to 1), compare against an independent reference computation with the attention and MLP outputs switched off, and check exact logits and attention values for fixed init_params weights.

Debugging tips

Write small experiments in the editor: inspect att, try a single token, zero out W_o or W_2, compare positions.

Constraints

Edit only gpt_forward; keep the signature and the (logits, att) return.

Follow-up

Which of these bugs would still let training loss go down? (Most of them: that is why they are dangerous.)
What unit tests would have caught each bug before training?

Examples

Input: logits, att = gpt_forward([3, 1, 4, 1, 5], init_params())
Output: logits.shape == (5, 11), att.shape == (5, 5)

In the fixed model every row of att sums to 1 and att is lower-triangular.

Input: gpt_forward([6], init_params(seed=3))[1]
Output: array([[1.]])

A single token can only attend to itself.

Practice in Python

import numpy as np

def init_params(vocab_size=11, d_model=8, max_len=8, seed=0):
    """Deterministic random weights for a 1-layer, single-head model. (Correct.)"""
    rng = np.random.default_rng(seed)
    def w(*shape, scale=0.5):
        return rng.standard_normal(shape) * scale
    d = d_model
    return {
        'wte': w(vocab_size, d), 'wpe': w(max_len, d),
        'ln1_g': 1 + w(d, scale=0.1), 'ln1_b': w(d, scale=0.1),
        'W_q': w(d, d), 'W_k': w(d, d), 'W_v': w(d, d), 'W_o': w(d, d),
        'ln2_g': 1 + w(d, scale=0.1), 'ln2_b': w(d, scale=0.1),
        'W_1': w(d, 4 * d), 'b_1': w(4 * d, scale=0.1),
        'W_2': w(4 * d, d), 'b_2': w(d, scale=0.1),
        'lnf_g': 1 + w(d, scale=0.1), 'lnf_b': w(d, scale=0.1),
    }

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

def softmax(x):
    """Row-wise softmax over the last axis. (Correct.)"""
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def gelu(x):
    """GELU, tanh approximation. (Correct.)"""
    return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3)))

def gpt_forward(tokens, params):
    """
    Forward pass of a 1-layer, single-head, pre-LayerNorm GPT.

    Args:
        tokens: List of T token ids (T <= max_len)
        params: Dict from init_params

    Returns:
        logits: (T, vocab_size)
        att: (T, T) attention probabilities (row i = how query i attends)
    """
    p = params
    tokens = np.asarray(tokens)
    T = len(tokens)
    d = p['wte'].shape[1]

    # Token + position embeddings
    x = p['wte'][tokens] + p['wpe'][1:T + 1]

    # Causal self-attention block
    h = layer_norm(x, p['ln1_g'], p['ln1_b'])
    q, k, v = h @ p['W_q'], h @ p['W_k'], h @ p['W_v']
    scores = q @ k.T
    att = softmax(scores)
    mask = np.tril(np.ones((T, T)))
    att = att * mask
    x = x + (att @ v) @ p['W_o']

    # MLP block
    h = layer_norm(x, p['ln2_g'], p['ln2_b'])
    x = h + gelu(h @ p['W_1'] + p['b_1']) @ p['W_2'] + p['b_2']

    # Final LayerNorm and tied LM head
    x = layer_norm(x, p['lnf_g'], p['lnf_b'])
    logits = x @ p['wte'].T
    return logits, att

The interactive editor loads with Python and NumPy. Run your implementation against 8 tests, or use the progressive hints.

More exercises in this section →