Repair a Broken Tiny Language Model, Then Make Decoding Incremental

You are handed PocketDecoder, a two-layer decoder-only transformer, plus a training step and a greedy generator. A teammate reports that training produces nan losses and, after a "quick fix", that the model trains but never learns anything useful.

Part 1 - Debug

The code below contains exactly four bugs. Find each one, explain the symptom it causes (crash/NaN vs. silent quality loss), and fix it. Do not restructure the model beyond what is needed.

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


class Attention(nn.Module):
    def __init__(self, dim, heads, ctx):
        super().__init__()
        self.heads, self.hd = heads, dim // heads
        self.to_q = nn.Linear(dim, dim)
        self.to_k = nn.Linear(dim, dim)
        self.to_v = nn.Linear(dim, dim)
        self.out = nn.Linear(dim, dim)
        causal = torch.triu(torch.ones(ctx, ctx, dtype=torch.bool), diagonal=1)
        self.register_buffer("future", causal)          # True = position in the future

    def forward(self, h):
        B, T, D = h.shape
        q = self.to_q(h).view(B, self.heads, T, self.hd)
        k = self.to_k(h).view(B, self.heads, T, self.hd)
        v = self.to_v(h).view(B, self.heads, T, self.hd)
        scores = q @ k.transpose(-1, -2) / math.sqrt(self.hd)
        scores = scores.masked_fill(~self.future[:T, :T], float("-inf"))
        weights = scores.softmax(dim=-1)
        mixed = (weights @ v).transpose(1, 2).reshape(B, T, D)
        return self.out(mixed)


class Layer(nn.Module):
    def __init__(self, dim, heads, ctx):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
        self.attn = Attention(dim, heads, ctx)
        self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

    def forward(self, h):
        h = h + self.attn(self.n1(h))
        return h + self.ff(self.n2(h))


class PocketDecoder(nn.Module):
    def __init__(self, vocab, dim=64, heads=4, layers=2, ctx=128):
        super().__init__()
        self.tok = nn.Embedding(vocab, dim)
        self.pos = nn.Embedding(ctx, dim)
        self.layers = nn.ModuleList(Layer(dim, heads, ctx) for _ in range(layers))
        self.norm = nn.LayerNorm(dim)
        self.head = nn.Linear(dim, vocab)

    def forward(self, ids):
        B, T = ids.shape
        h = self.tok(ids) + self.pos(torch.arange(B, device=ids.device))[:, None, :]
        for layer in self.layers:
            h = layer(h)
        return self.head(self.norm(h))


def train_step(model, opt, batch):              # batch: (B, T+1) token ids
    inputs, targets = batch[:, :-1], batch[:, :-1]
    logits = model(inputs)
    loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()
    return loss.item()


@torch.no_grad()
def generate(model, prompt, n_new):            # prompt: (1, T0)
    ids = prompt
    for _ in range(n_new):
        logits = model(ids)
        nxt = logits[:, -1].argmax(-1, keepdim=True)
        ids = torch.cat([ids, nxt], dim=1)
    return ids

A good sanity task: sequences that count upward modulo 10 from a random start (e.g. 7 8 9 0 1 2 ...). After a few hundred AdamW steps a correct model reaches near-zero loss and greedy generation from [3, 4, 5] continues 6 7 8 9 0 1 2 3.

buggy model:  final loss = nan, generation = [3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0]
fixed model:  final loss ~ 0.002, generation = [3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3]

Part 2 - Incremental decoding

generate re-runs the full sequence through every layer for each new token, which is O(T^2) work per step of attention and wasteful in the projections. Add a key/value cache:

• PocketDecoder.forward(ids, cache=None) returns (logits, new_cache), where the cache holds, per layer, the keys and values of every token processed so far.
• generate_cached(model, prompt, n_new) runs the prompt once ("prefill"), then feeds only the newest token on each step.
• Its output must be identical to the uncached generate for the same model and prompt.

Constraints

• Batch size 1 is enough for generation, but forward must still work for training batches with no cache.
• Positions must remain correct when only one token is passed in.
• Raise a clear error if the total length exceeds the context window.