Skip to content

ML coding: implementing an LLM block in PyTorch

  • Coding
  • Asked often
  • About 3 h to learn
  • 4 vetted resources
Build a plan that includes this

What this topic covers

The ML-coding round asks you to write a working piece of a transformer in PyTorch in about 30 minutes, usually a Llama-style decoder block: RMSNorm, rotary position embeddings, grouped-query attention, a SwiGLU feed-forward layer and pre-norm residuals. Interviewers want to see that you think in tensor shapes, can apply a causal mask correctly and can prove your code works with a quick test. A strong answer states the shape at every step, keeps heads as a batch dimension rather than looping, checks that changing a future token leaves earlier outputs unchanged, and explains why each Llama change replaced its GPT-2 counterpart.

Builds on

You’re ready when you can

What interviewers check, from the resources below.

  • Writes RMSNorm, RoPE, GQA and SwiGLU from memory with correct shapes
  • Builds the causal mask and shows a future token can’t change earlier outputs
  • Treats heads as a batch dimension instead of looping over them
  • Finishes a full pre-norm block with residuals in about 30 minutes
  • Explains masked attention as a weighted average using a lower-triangular matrix
  • Says why scores are divided by √head_size
  • Adds residual connections and layernorm in the right places
  • Merges per-head modules into one batched multi-head attention
  • Lists each GPT-2 to Llama change and the reason for it
  • Applies RoPE to queries and keys only, never to values

What to study, in order

Each link opened and checked by hand; the date says when.

  1. Start here

    TorchCode (opens in a new tab)

    Hugging Face Spaces (GitHub: duoan/TorchCode) · duoan · 90 min · Free · checked 21 Oct 2026

    Why
    The ML-coding round is a timed “write this layer in PyTorch” exercise, and TorchCode is a set of exactly those drills with an automatic grader that checks outputs and gradients.
    What to read
    Open the Space (it starts JupyterLab in the browser, or use the Colab badge on any problem) and do these in order, timing each: #8 RMSNorm, #9 Causal Self-Attention, #10 Grouped Query Attention, #24 RoPE, #15 SwiGLU MLP, then #13 GPT-2 Block as the full-block rehearsal. Try every one without hints first.
    What to skip
    Skip the non-LLM problems (ReLU, BatchNorm, Conv2d, ViT, Linear Regression) and the RL losses (DPO, GRPO, PPO); those belong to other topics.

  2. Free alternative

    Let's build GPT: from scratch, in code, spelled out. (opens in a new tab)

    YouTube · Andrej Karpathy · 75 min · Free · checked 21 Oct 2026

    Why
    Watching someone build a decoder step by step, checking shapes as he goes, is the best model for how to talk through your code in the interview.
    What to read
    Watch from 42:13 (“version 1: averaging past context with for loops”) through 1:42:39: the masked-softmax trick, “version 4: self-attention” at 1:02:00, the six attention notes, then multi-head attention, feed-forward, residuals, layernorm and scaling up. Code along in the linked Colab, then do his exercise EX1 (all heads in one batched class).
    What to skip
    Skip the first 42 minutes (data loading and the bigram baseline) if you know them, and the closing section from 1:42:39 on encoder/decoder variants, nanoGPT and ChatGPT.

  3. Go deeper

    Converting a From-Scratch GPT Architecture to Llama 2 (and Llama 2 to Llama 3.2) (opens in a new tab)

    GitHub (LLMs-from-scratch) · Sebastian Raschka · 60 min · Free · checked 21 Oct 2026

    Why
    Interviewers often follow up with “how does this differ from Llama?”, and these notebooks change a GPT block into a Llama block one piece at a time.
    What to read
    In converting-gpt-to-llama2.ipynb, read sections 1.1 to 1.6 (RMSNorm, SiLU, the gated feed-forward, RoPE, adding RoPE to attention, the updated block). In converting-llama2-to-llama3.ipynb, read 1.2 (modified RoPE) and 1.3 (grouped-query attention). Retype each module yourself rather than just reading it.
    What to skip
    Skip sections 2 to 5 of both notebooks (tokenizer, loading pretrained weights, chatting with the model) and the Llama 3.1/3.2 sections; they cover weight loading, not the block.

  4. Go deeper

    nanoGPT (opens in a new tab)

    GitHub · karpathy · 60 min · Free · checked 22 Oct 2026

    Why
    Interviewers still treat nanoGPT’s two ~300-line files as the reference shape of a GPT training loop and model, so being able to talk through them is a quick signal you’ve trained one yourself.
    What to read
    Read model.py (CausalSelfAttention, MLP, Block, GPT.forward and configure_optimizers) and the training loop in train.py (gradient accumulation, the learning-rate schedule, estimate_loss). Then follow the README’s “quick start”: run data/shakespeare_char/prepare.py and the CPU command for config/train_shakespeare_char.py, then sample.py.
    What to skip
    Skip “reproducing GPT-2”, “baselines”, multi-node torchrun and “efficiency notes”. The repo is now marked deprecated in favour of nanochat, so treat it as a readable reference, not a project to extend.

Companies that test Coding most

Put ML coding in your plan

Tell us your interview date and hours a day; it lands on the right day, in the right order.

Build my plan