Debug a PyTorch Decoder-Only Transformer: Find Four Silent Bugs

You are handed the PyTorch code below for a small decoder-only transformer language model trained with next-token prediction, together with its loss function. It runs without raising an error, but it does not behave like a correctly working language model. Find the bugs, fix them, and explain how each one would show up. The report says only that the candidate was asked to debug a transformer; it does not include the code. The code below is an illustrative stand-in with exactly four planted bugs, none of which raises an exception for valid inputs. ```python import math import torch import torch.nn as nn import torch.nn.functional as F class CausalSelfAttention(nn.Module): def __init__(self, d_model: int, n_heads: int, dropout: float): super().__init__() assert d_model % n_heads == 0 self.n_heads = n_heads self.d_head = d_model // n_heads self.qkv = nn.Linear(d_model, 3 * d_model) self.proj = nn.Linear(d_model, d_model) self.attn_drop = nn.Dropout(dropout) def _split_heads(self, t: torch.Tensor) -> torch.Tensor: # (B, T, C) -> (B, n_heads, T, d_head) B, T, C = t.shape return t.reshape(B, self.n_heads, T, self.d_head) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) q, k, v = self._split_heads(q), self._split_heads(k), self._split_heads(v) scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head) # (B, n_heads, T, T) future = torch.triu( torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=1 ) scores = scores.masked_fill(future, float("-inf")) weights = F.softmax(scores, dim=-2) weights = self.attn_drop(weights) out = weights @ v # (B, n_heads, T, d_head) out = out.transpose(1, 2).contiguous().view(B, T, C) return self.proj(out) class Block(nn.Module): def __init__(self, d_model: int, n_heads: int, dropout: float): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = CausalSelfAttention(d_model, n_heads, dropout) self.ln2 = nn.LayerNorm(d_model) self.mlp = nn.Sequential( nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model), nn.Dropout(dropout), ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.ln1(x)) x = self.mlp(self.ln2(x)) return x class TinyLM(nn.Module): def __init__(self, vocab_size: int, max_len: int, d_model: int = 256, n_heads: int = 8, n_layers: int = 4, dropout: float = 0.1): super().__init__() self.tok_emb = nn.Embedding(vocab_size, d_model) self.pos_emb = nn.Embedding(max_len, d_model) self.drop = nn.Dropout(dropout) self.blocks = nn.ModuleList( [Block(d_model, n_heads, dropout) for _ in range(n_layers)] ) self.ln_f = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size, bias=False) def forward(self, idx: torch.Tensor) -> torch.Tensor: # idx: (B, T) token ids -> logits: (B, T, vocab_size) B, T = idx.shape pos = torch.arange(T, device=idx.device) x = self.drop(self.tok_emb(idx) + self.pos_emb(pos)) for block in self.blocks: x = block(x) return self.head(self.ln_f(x)) def lm_loss(model: TinyLM, idx: torch.Tensor) -> torch.Tensor: # idx: (B, T) a batch of training sequences logits = model(idx) return F.cross_entropy(logits.reshape(-1, logits.size(-1)), idx.reshape(-1)) ``` ### Constraints and Clarifications - The intended architecture is a standard pre-layer-norm decoder-only transformer: each block adds the output of causal multi-head self-attention, and then the output of a feed-forward network, to its input, each sublayer reading a layer-normalized copy of that input. - `lm_loss` should return the average cross-entropy of predicting each token from the tokens before it. - Training batches are unpadded, every sequence in a batch has the same length, and lengths run from 2 to `max_len` tokens; token ids are valid. - These are design choices, not bugs: learned positional embeddings, the final layer norm, the output head without a bias, the GELU activation, dropout on the attention weights and PyTorch's default initialization. ### Clarifying Questions - May I run the code on small inputs, or should I find the bugs by reading only? - Will the model be used for autoregressive generation, so that causality must hold exactly? - Should the fixes keep the architecture as written, or may I replace parts with library calls such as PyTorch's fused attention? ### Part 1 — Find and fix the bugs Identify the four bugs, explain exactly what each one does wrong, and give the corrected code. ```hint Shapes are not enough Every tensor here has the expected shape. Write down what each axis means after each line of the attention module. ``` ```hint Read the loss too For each position, ask what the target is and whether the model can already see it. ``` #### What This Part Should Cover - The location of each bug and the mechanism by which it breaks the model - A minimal, correct fix for each - Which bugs let information from later tokens reach earlier positions ### Part 2 — Symptoms and tests For each bug, describe how it would show up during training, evaluation or text generation, and write down a quick test that would catch it without reading the code. ```hint Properties of a correct model List the properties every correct causal language model must have, then turn each one into a check that runs in seconds. ``` #### What This Part Should Cover - The expected symptom of each bug, including the direction of its effect on the training loss - A concrete, fast test for each property the bugs violate - A general debugging order for a model that trains but does not work ### What a Strong Answer Covers - All four bugs found, with no correct code flagged as a bug - Reasoning about what each axis means, not only about tensor shapes - Which bugs break causality or change the training objective, and why a low loss cannot reveal them - Tests that would stop these bugs from coming back - Clear, minimal fixes ### Follow-up Questions - Why are attention scores divided by the square root of the head dimension, and what happens to training if that scaling is dropped? - How would you replace the hand-written attention with PyTorch's fused `F.scaled_dot_product_attention`, and how would you confirm the outputs match? - This block applies layer norm before each sublayer. How does that compare with applying it after the residual addition, especially in deep stacks?