The transformer debugging round, step by step

In this article

1. Before you start
2. Why interviewers ask this
3. The transformer debugging round from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one part of the PyTorch and machine learning (ML) coding round guide, and it assumes you have seen attention once.

Why interviewers ask this

Reading and fixing code shows an interviewer what writing code from a blank page cannot: whether you know what correct looks like. Candidates report rounds in which a transformer implementation has several annotated regions and each region hides a bug, rounds where a pair-programming engineer asks you to fix a bug in a transformer block, and a pre-norm bug in a block with residuals. These are candidate reports, not published formats, so treat the details as examples.

A weak answer reruns the code and hopes. A strong answer states what each region must satisfy, such as "each row of attention weights sums to one" or "earlier tokens never change when a later token changes", and tests that. A Staff-level answer also says which bugs a loss curve hides, because several planted bugs leave the loss looking normal for a while.

The transformer debugging round from first principles

Transformer debugging round

A transformer debugging round is a timed interview exercise in which a candidate reads a transformer implementation that runs but is wrong, finds the planted bugs and fixes them, usually in PyTorch, while explaining the reasoning aloud.

Planted bug

A planted bug is a deliberate, usually one-line defect in otherwise working code, such as a mask applied in the wrong place, chosen so that the program still runs and returns plausible numbers.

Invariant test

An invariant test is a short check of a property that must always hold, such as attention weights summing to one, that fails when a bug breaks the property and does not depend on the exact output values.

In plain words

Think of a bank statement that must balance. You do not re-add every line to find a mistake. You check totals that must match, then narrow down. A transformer has totals like that. Each row of attention weights sums to one. A token never reads the future. The output keeps the input's shape. The analogy breaks in one place. A statement that fails to balance always shows it. A transformer with a bug still returns numbers, so you must go looking for the failed total.

In one precise sentence, debugging a transformer means checking, in the order the data flows, the invariants each stage promises, and fixing the first stage that breaks one.

Step 1: Read the shapes and layout before the math

Check the shape of every tensor at every stage, because many bugs show as a wrong axis. Write down the expected shape beside each line: input (batch, sequence, features), scores (batch, heads, sequence, sequence), output back to the input shape. PyTorch's multi-head attention module defaults undefined to False, so its layout is (sequence, batch, feature) unless you set it, and a batch-first tensor passed in silently treats the batch as time.

Input: the code and a sample tensor. Output: a shape annotation per line. A square sequence-by-sequence score matrix hides axis mix-ups, so use a batch size and a sequence length that differ, such as 2 and 5 (illustrative). You now know what each stage should receive, which the next step needs.

Step 2: Mask the scores before the softmax, never after

A causal mask is a lower-triangular pattern that stops a token from reading later tokens, and it must change the scores before the softmax. Vaswani et al. (2017) mask by setting illegal entries to minus infinity at the softmax input. PyTorch's reference code builds the keep-pattern with tril and fills the rest with minus infinity.

Masking after the softmax is the planted bug: the weights of the allowed tokens no longer sum to one, and the denominator still contained the future. Take the illustrative scores 1, 2, 0.5 and 3 with the last key masked. Masking before the softmax gives the weights 0.2312, 0.6285 and 0.1402, which sum to 1. Masking after gives 0.0854, 0.2320 and 0.0518, which sum to only 0.3692.¹ Worse, an earlier token's output still changes when the last token changes, so the future leaks. The invariant test is to edit the last token and assert that every earlier output is unchanged. See also what the causal mask does.

causal = torch.ones(T, T, dtype=torch.bool).tril()
scores = scores.masked_fill(~causal, float("-inf"))   # before softmax
weights = torch.softmax(scores, dim=-1)

model.eval()                                          # dropout off, or the test fails on correct code
x2 = x.clone(); x2[:, -1] += 5.0                      # invariant test
assert torch.allclose(f(x)[:, :-1], f(x2)[:, :-1])

Step 3: Check the scale and the softmax axis

The scores must be divided by the square root of the per-head size, and the softmax must run over the key axis. Vaswani et al. (2017) divide each dot product by the square root of the key size and suspect that large sizes push the softmax into regions with extremely small gradients. PyTorch's default scale is one over the square root of the last dimension of the query, which after the head split is the head dimension, not the model width. Dividing by the model width's root is a common planted bug.

The softmax runs over the last, key axis (dim=-1 in the reference code). Over the query axis, the rows no longer sum to one. The invariant test: weights.sum(-1) must be all ones. A missing scale passes that test, so also check the spread of the scores.

Step 4: Check the head split and the merge

The head dimension is the size of the slice each head receives, equal to the model width divided by the number of heads, so the width must divide evenly. The PyTorch documentation states that the embedding size is split across heads and the head outputs are concatenated and passed through an output projection. With a width of 512 and eight heads, the head dimension is 64 (illustrative sizes; the original base model is described by Vaswani et al.).

The bug is a reshape to (heads, sequence, head size) when the layout is (sequence, features): the numbers stay the same, but each head reads a mix of tokens. Reshape to (sequence, heads, head size), then transpose. In the toy, this bug changes the output by up to 1.643 while no future leaks, so only a comparison with a reference catches it.¹ The same care applies on the way back: transpose before merging heads.

Step 5: Check the norm placement and the residual

A residual connection adds a block's input back to its output, so x leaves as x plus the sublayer output. Layer normalization (norm for short) rescales each token's vector to zero mean and unit variance, then applies a learned scale and shift. The norm either follows the residual add or sits inside the branch. Post-norm, the layout of the original paper, computes the norm of x plus the sublayer output.

Pre-norm is the layout in which each sublayer reads a normalized copy of x and its output is added to the raw x. Xiong et al. (2020) show that post-norm has large gradients near the output at initialization, which makes warm-up practically helpful, while with pre-norm the gradients are well-behaved at initialization, which motivates removing warm-up. They describe the pre-norm variant as adding a final layer normalization before prediction.

A typical pre-norm bug is normalizing x and then adding the sublayer output to the normalized copy, which changes what the residual carries. Fix: x = x + attn(norm1(x)), then x = x + mlp(norm2(x)), and check the final norm exists.

1 · Shapes2 · Mask3 · Scale and axis4 · Heads5 · NormReplay

Step-by-step description

1. Write the expected shape beside each tensor.
2. Mask the scores before the softmax.
3. Scale by head size; softmax over keys.
4. Reshape, then transpose.
5. Keep the raw residual in pre-norm.

• Check shapes, then the mask, then scale and softmax axis, then head split, then norm placement.
• Debug in the order the data flows.

Worked example

A tiny attention layer shows how each planted bug behaves. The inputs are random (illustrative): four tokens, model width 8, two heads, so the head dimension is 4. Compared with a correct reference, the output of the layer changes by up to 0.336 without the scale, 1.192 with the mask after the softmax, 1.538 with no mask, 1.378 with the softmax over the query axis and 1.643 with the wrong head reshape.¹ The missing scale looks mild here only because the head size is small.

At a head dimension of 64, with unit-variance query and key entries (illustrative), the standard deviation of a dot product is 8.04 unscaled and about 1.0 after dividing by √64 = 8, which is the same as multiplying by 0.125.¹ Over 16 keys, the largest softmax weight averages 0.843 unscaled and 0.245 scaled.¹ The bad case is nearly one-hot, which is the saturation Vaswani et al. describe.

The values come from the script in the footnote.¹

Text description of the diagram

• The output changes by 0.336 without the scale and by 1.643 with the wrong head reshape.
• The mask after softmax changes the output by 1.192.

Trade-offs and edge cases

Reading order and tests trade breadth against time. Checking every stage costs minutes; a good order puts the cheap, high-yield checks first: shapes, then the mask and row sums, then scale. Two conventions need care. In scaled_dot_product_attention a boolean True means "take part"; in nn.MultiheadAttention a boolean True means "masked out" (PyTorch 2.13 docs). Passing one module's mask to the other inverts it. The scaled-dot-product reference code also refuses to combine is_causal with an explicit mask. The is_causal docs wording was reported (and the issue has since been closed) as misleading about the mask orientation, so trust the lower-triangular construction.

Edge cases: a positional-encoding bug (see positional encodings from scratch) shows up as a model that ignores word order; a weight-tying or initialization bug is hard to see in a single forward test, so ask for a few training steps and watch the loss; the debugging training loops guide covers that side. Where you have time, compare against a trusted reference implementation, as in building generative pre-trained transformer (GPT) from scratch.

Text description of the diagram

• Edit the last token to catch a mask bug.
• Check that attention rows sum to one to catch an axis bug.

How to answer it in an interview

The spoken structure: name the order, give the invariant at each stage, and say the test out loud before you run it.

Common mistakes

• "I'll run it and look at the loss." (wrong) Several planted bugs leave the loss plausible for a while. (right) "I'll test invariants: row sums, a future-leak check and shapes."
• "The mask can go after the softmax, it's the same thing." (wrong) The weights then do not sum to one and the denominator still held the future. (right) "Mask with minus infinity before the softmax."
• "The scale is the square root of the model width." (wrong) Each head attends over its own slice. (right) "I divide by the square root of the head dimension."
• "True always means masked." (wrong) The convention depends on the module. (right) "True means attend in scaled_dot_product_attention and masked in nn.MultiheadAttention; I check which module I am in."
• "Reshape to heads is enough." (wrong) It mixes tokens when the layout is sequence-first. (right) "Reshape, then transpose, and reverse it when merging."

Follow-up questions

How do you debug a transformer implementation in an interview?

State the expected shape of each tensor, then test one invariant per stage in data-flow order: shapes, mask and row sums, scale, head split, norm placement. Say each test aloud before running it.

What are the most common bugs in a transformer implementation?

The ones in this walkthrough: a mask applied after the softmax, a wrong softmax axis, a missing or wrong scale, a head reshape without a transpose and a misplaced norm or residual. Candidates report such planted bugs, so the list is advice, not a published catalogue.

Which is the bug if the loss looks fine but the model sees the future?

A missing or late mask. Training loss drops easily when the model can read the next token, so the loss looks good while generation fails. The edit-the-last-token test exposes it.

How do you implement KV caching in this transformer?

Store each layer's keys and values for past tokens, compute only the new token's query, key and value, append the new key and value, and attend over the cache. The position given to the new token must continue from the cache length, which is a bug candidates should check.

Can you use an LLM in a debugging round?

Rules differ by interviewer. A candidate reports a round at one company where an assistant was allowed on a toy codebase with unit tests; ask the interviewer at the start. Either way, you must explain every fix.

Sources

1. Attention Is All You Need (arXiv 1706.03762) — arXiv (Vaswani et al.), 12 Jun 2017
2. On Layer Normalization in the Transformer Architecture (arXiv 2002.04745) — arXiv (Xiong et al., ICML 2020), 29 Jun 2020
3. torch.nn.functional.scaled_dot_product_attention (PyTorch 2.14) — PyTorch documentation, 26 Oct 2026
4. torch.nn.MultiheadAttention (PyTorch 2.13) — PyTorch documentation, 26 Oct 2026
5. scaled_dot_product_attention is_causal description is misleading (issue 126873) — GitHub (pytorch/pytorch), 26 Oct 2026