You are given a buggy PyTorch implementation of a GPT-like causal language model (transformer), along with a training loop that overfits the model on a short sequence of text. Your task is to identify and fix all bugs in the code so that: (1) training loss decreases (converges), and (2) the model can generate coherent text. Typically there are 4 bugs total (occasionally 5), and the problem statement or comments in the code indicate which sections contain bugs. Verified success criteria: training loss goes from NaN/non-decreasing to converging, and sample/generation output produces readable text. The standard four bugs are: (a) positional embedding initialization error (learnable embeddings initialized incorrectly — not a sin/cos swap, but wrong initial values), (b) causal attention mask not set to -inf before softmax (so masked positions are not properly zeroed out), (c) missing or misplaced loss.backward() / gradient descent step in training loop (next-token prediction loss label shift also commonly included), and (d) output projection nn.Linear dimensions swapped / shape mismatch. A fifth bug (seen in some variants) is a subtle typo (e.g., variable 'y' written as 'v'). A KV class is provided; no test harness needs to be written from scratch — a provided train/sample function validates the fix.
Implement KV cache for inference/decoding. A KV cache class is provided (returns length and (K, V) tensors). Modify the attention computation to use/update the cache, adjust positional encoding offsets, disable causal mask when cache is populated, and implement a generate() function. Verify output matches non-cached generation.
Modify the transformer to work as a text classifier (change the final output layer). A test harness is provided; make it pass.
What is the time complexity of the generation function? How does KV cache improve it?
Discuss tradeoffs of different approaches to convert the transformer to a classifier (CLS token, last token, mean pooling).
| Approach | Notes |
|---|---|
| Test-driven incremental fixing | Run training after each fix to isolate impact; slower overall but confirms individual bugs. Risky under time pressure given slow CoderPad execution. |
| Diff against known-good reference | Mentally compare against a memorized nanoGPT reference implementation; fast if architecture is well-known, but can miss novel or subtle bugs not in canonical reference. |
| Replace learnable positional embeddings with sinusoidal | Interviewers noted that some candidates tried to swap in sin/cos embeddings to fix the initialization bug; this was considered incorrect — the intent is to fix the initialization of the learnable embeddings, not replace them. Will likely be flagged as a wrong fix. |
| Using token mean pooling before logits for classifier follow-up | One interviewer explicitly preferred mean-pooling over the sequence dimension before computing classification logits; other approaches (CLS token, last token) are functionally valid but may be less preferred by some interviewers. |
Common mistakes: Missing the loss label-shifting bug — candidates from a masked-language-model background (BERT-style) did not think to shift labels for next-token prediction and left this bug unfixed.; Trying to replace the learnable positional embedding with a fixed sinusoidal one instead of correctly initializing the existing learnable embedding.; Finishing all bugs and the follow-up but running out of time to test, leaving no verified output — the interviewer could not confirm correctness.; Doing everything correctly but too slowly — completing the debug and classifier right at the time limit with only two or three minutes for questions was noted as a likely reason for rejection even when the code was correct.; The platform (CoderPad) runs code very slowly; candidates who ran the code frequently wasted significant time waiting.; Fixing the positional embedding bug by replacing learnable embeddings with sinusoidal embeddings (wrong intent — should fix initialization of learnable embeddings); Not recognizing that the causal mask must be set to -inf before softmax (setting it to 0 after softmax has no effect); Forgetting to reset/skip the causal mask when KV cache is present during inference; Missing the positional encoding offset when generating with KV cache (always starting from position 0 instead of cache_length); Not running/validating the code after fixes (finishing without checking that loss actually decreases); Spending too much time on one bug and running out of time before completing KV cache follow-up; Missing the typo bug ('y' written as 'v' or similar) due to not reading code carefully enough; Fixing output projection shape bug in wrong direction (swapping input/output dims incorrectly)
Interviewer hints: The code contains comments marking which sections contain bugs — candidates only need to inspect those regions and do not need to read the entire file.; The interviewer knows where each bug is but may not know which of multiple valid fixes is correct, so if a candidate's fix compiles but behaves unexpectedly the interviewer may not be able to guide them.; For the classifier follow-up, the interviewer explicitly prefers taking the mean of token representations before computing logits over other pooling strategies.; For the KV cache, correctness is checked by comparing output to the non-cached generate function — matching output is the pass criterion.; For KV cache validation, interviewer accepts matching output between cached and non-cached generation as proof of correctness (no separate test harness needed); Interviewer gave candidates a couple extra minutes to finish test cases when close to completion
What passers do: Candidates who had hand-written or carefully studied the core transformer components (positional embedding, causal mask, output projection, training loop) before the interview found all four bugs quickly and had time for follow-ups.; Successfully fixing all four bugs is verified by running the provided training loop and observing loss decrease and correct text output — candidates who confirmed both signals passed this stage cleanly.; For the KV cache follow-up, candidates who had practiced a minimal GPT with KV cache end-to-end beforehand found the wiring straightforward and verified correctness by matching output to the non-cached version.; The classifier follow-up was completed quickly by candidates who recognized it as a small parameter change (mean-pool then new linear layer) rather than a structural redesign.; Studying Karpathy's nanoGPT/YouTube course beforehand and hand-writing the transformer architecture before the interview; Using GPT to generate fake buggy transformers for practice; Finding all 4 bugs quickly (well before time limit) by reading code sections flagged by comments first; Implementing KV cache cleanly with correct positional offset and conditional causal masking; Discussing tradeoffs with the interviewer when implementing the classifier follow-up; Finishing ahead of time, leaving room for follow-up questions and testing
Why people fail: Wrote all four fixes correctly but ran out of time before running training to verify — the explicit success criterion is loss decrease + correct generated output, and at least one report attributes the loss to exactly this: "failed at the end - I finished writing it but had no time to test".; Background mismatch in one report: candidate had CV experience but no autoregressive-LM background, missed the loss-shifting bug because the next-token-prediction framing was unfamiliar. Single report — anecdotally informative on the prep gap that bites.; Got 3/4 architecture-related bugs but lost the round on the loss bug because it required understanding causal LM mechanics, not just transformer architecture; Wasted ~20 min on bug #1 (positional embedding initialization) trying to switch to sin/cos when the intended fix was a 1-line nn.init.normal_(self.pos_emb, std=0.02) — running out of clock for the harder bugs; On the KV cache follow-up: didn't get to the parity check with the no-cache reference because the cache implementation was buggy in subtle ways (forgot offset adjustment, applied mask when cache present) — the interviewer wants the parity check explicitly; Finishing all tasks but only at the deadline (no time to test/validate), which interviewers penalized; Fixing only the architecture bugs and missing the training loop bugs (e.g., loss label shift, missing backward pass); Being too slow overall — interviews described as 'time-tight' requiring practiced hand speed; Not testing the code after fixes, so bugs remain undetected; Replacing learnable positional embeddings with sin/cos instead of fixing initialization; Failing the KV cache follow-up due to time pressure after spending too long on debugging phase; Poor mindset/nerves causing degraded performance on otherwise known material
Edge cases probed: KV cache + positional encoding offset: positional encoding for the new token must use cache_length + 0 (or cache_length), not position 0 — getting this wrong silently breaks generation parity; KV cache + causal mask: for single-token decode with a non-empty cache, there are no future positions to mask, so the mask is unnecessary; for chunked decode (>1 new token at a time over a non-empty cache), an offset-aware mask sized to the new chunk is still required to prevent new tokens from attending to each other's future positions.; 5-bug variant with a subtle typo (e.g., wrong variable name, transposed argument): no canonical recipe, only careful reading catches it; Multi-head variant with bugs in head dimension reshape (B,T,n_head,head_dim swap order); Classifier conversion test harness: tests check specific tensor shape (batch, num_classes), not just convergence — candidates who change the loss but forget to change the output shape fail the harness; Time-complexity discussion for both training and inference. The interviewer typically asks about inference (autoregressive generation): without KV cache, scoring each step at context length T is O(T²·d) per step (O(T³·d) for generating T tokens); with cache, the per-step cost drops to O(T·d) (O(T²·d) total). For training/teacher-forced decoding, attention is O(T²·d) per example because all positions are computed in parallel.; Whether the candidate considers batch_size > 1 when implementing KV cache (or simplifies to batch_size=1); Whether the candidate correctly handles positional encoding offset when KV cache is non-empty (must offset by cache length, not start from 0); Whether causal mask is conditionally applied only when cache is absent/empty (not always); Typo bugs (e.g., variable 'y' written as 'v') that are subtle and easy to miss; Whether the candidate validates their fix by running the provided training loop (loss should go from NaN to converging); Label shifting in next-token prediction loss (input[:-1] vs input[1:]) — easy to miss for candidates from non-LM backgrounds (e.g., BERT/MLM background); Whether the learnable positional embedding is kept vs. replaced with sin/cos
What you just read — canonical solution, follow-up arc, what passing candidates actually did — exists for all 100 OpenAI questions, refreshed monthly from new candidate reports.