« All Roles

Senior AI Engineer — LLMs, VLMs & Agentic Systems — Curriculum

Target Role: Senior AI Engineer at a Robotics Research Center — design, fine-tune, align, and deploy Large Language Models (LLMs), Vision-Language Models (VLMs), and autonomous agentic systems as scalable, low-latency production software (the reference JD — jd.md).

Also prepares you for:

• Frontier-lab & applied-research roles — Member of Technical Staff / Research Engineer on training, inference, alignment, or multimodal teams.
• Any company productionizing LLMs/VLMs — staff/senior AI engineer owning the model stack end-to-end: tokenizer → transformer → fine-tune → align → quantize → serve → agent → evaluate → observe.
• Embodied-AI / robotics-foundation-model teams — the differentiator of this JD: agents that perceive, plan, and act, with the sim-to-real and VLA-action layer built explicitly.
• The "I can use the API" floor is not the ceiling here: this track makes you build the mechanisms the frameworks only expose — PagedAttention, the LoRA delta, the DPO loss, the ReAct loop, the all-reduce ring — so you can debug them at 2 a.m. and defend them in a principal interview.

Duration: ~36 weeks core (≈9 months) — extendable with the capstone and the C++/CUDA track Seniority Target: Senior → Staff IC — the person who owns the model stack, makes the quality/latency/cost/safety tradeoff, and is the escalation point when generation is wrong, slow, or unsafe. Goal: Make you the engineer who can architect a multimodal agentic AI system at the system level and debug it at the token-logit, attention-mask, KV-cache-block, LoRA-rank, quantization-scale, gradient-shard, reward-margin, and tool-call level — across training, alignment, inference serving, retrieval, agents, evaluation, and embodied control.

Why This Curriculum Exists

Most "learn LLMs" material teaches you to call an API: which SDK method, which prompt, which framework abstraction. That gets you to "AI-adjacent engineer." It does not get you to Senior AI Engineer at a research center, because that job is to make correct tradeoffs under conflicting requirements (quality vs latency vs cost vs memory vs safety) and to debug the layer everyone else treats as magic — and you cannot do either if the transformer, the sampler, the scheduler, the optimizer, and the agent loop are black boxes.

So this track is built on two non-negotiables:

1. Every mechanism in the JD gets implemented, not just described. Scaling-law and FLOPs/KV-cache arithmetic, a BPE tokenizer, scaled-dot-product + multi-head attention with RoPE and a KV-cache, a reverse-mode autograd engine that trains a tiny LM, a CLIP-style contrastive VLM with a LLaVA-style projector, the LoRA/QLoRA delta and a distillation loss, affine + GPTQ/AWQ-style quantization, a Bradley-Terry reward model and the DPO/PPO losses, temperature/top-k/top-p/min-p sampling and a grammar-constrained decoder, a PagedAttention block manager with continuous batching and speculative decoding, data/tensor/pipeline parallelism and ZeRO memory accounting, an HNSW/IVF ANN index with reranking, a ReAct agent with typed tools and memory, a multi-agent orchestrator, a neurosymbolic verifier, an embodied RL planner with a VLA action decoder, an eval harness with LLM-as-judge and calibration, and a GPU roofline/occupancy model — you build a working, tested miniature of each.
2. Everything is taught to the depth interviewers and incidents actually demand — the 6ND training-FLOPs rule, why softmax overflows without the max-subtraction trick, the merge-priority off-by-one in BPE, why a Durable… (sorry, wrong track) — why a Durable orchestrator is replay-safe but an agent loop is not, the R+W of the KV-cache block table, the LoRA α/r scaling, the symmetric-vs-affine zero-point, the DPO β temperature, the ring-all-reduce 2(N−1) step count, the speculative-decoding acceptance criterion. We assume you are building the system that serves millions of multimodal requests with quality, latency, and safety guarantees — not a notebook demo.

How Each Phase Teaches (the "many ways of explaining")

Every phase deliberately explains the same material through several lenses, because senior-level understanding is the intersection of all of them:

The lab engineering standard is documented in LAB-STANDARD.md. Read it before you start the first lab.

The Phases

Use the sidebar for the full clickable table of contents, and see GLOSSARY.md / CHEATSHEET.md for the running references.

How to Use This Track

1. Start here, then read Phase 00 end to end — it builds the mental model (scaling laws, the FLOPs/memory/latency dials, the tradeoff calculus) that every later phase plugs into.
2. For each phase: read WARMUP.md (professor), skim HITCHHIKERS-GUIDE.md (the numbers), do the lab (lab.py red → green against test_lab.py), then read BROTHER-TALK.md.
3. Work phases roughly in order — the transformer (P02), autograd (P03), and the FLOPs/memory math (P00) underpin everything after them; agents (P12–P13) assume the serving and tool layers below them.
4. Use interview-prep/ and system-design/ as running references; the capstone (P19) ties the whole stack together.

Cross-Cutting Themes (the senior AI engineer's lenses)

• Training compute vs inference compute — 6ND to train once; 2N per token forever. The bill, the latency, and the carbon all live on the inference side at scale, which is why quantization (P06), serving (P09), and the KV-cache (P02/P09) get their own phases.
• Quality is a distribution, not a number — every generation is a sample. Sampling (P08), evaluation (P16), and alignment (P07) all exist because the model is a probability distribution over tokens, and "it worked once" is not a guarantee.
• The KV-cache is the memory wall — autoregressive decoding is memory-bandwidth-bound, not compute-bound; PagedAttention, continuous batching, and quantized KV exist to fight that one fact. Internalize it in P00 and you understand half of inference engineering.
• Agents are loops over an unreliable oracle — the LLM is a stochastic function that can hallucinate a tool call. Every agent design (P12–P14) is about constraining and verifying that oracle: typed tools, memory, multi-agent critique, symbolic verification.
• Quantify the tradeoff — FLOPs, KV-cache bytes, tokens/sec, $/1M-tokens, recall@k, reward margin, ECE — these are arithmetic you do before the design review, not after the incident. Every phase ends with a number you can defend.
• Embodiment changes the contract — when the agent's action moves a motor, latency is a safety property and "regenerate" is not an option. P15 makes the perception→plan→act loop and the sim-to-real gap explicit.

This role is part of TII’s Robotics Research Center.

Position Overview We are seeking a highly skilled Senior AI Engineer with deep expertise in Large Language Models (LLMs), Vision-Language Models (VLMs), and agentic model architectures. The ideal candidate will have a strong foundation in both research and engineering, with hands-on experience developing, fine-tuning, and deploying advanced AI systems. You will contribute to building scalable, production-ready AI applications, integrating multimodal reasoning, and pushing the boundaries of autonomous intelligent agents.

Key Responsibilities • LLM/VLM Development & Integration: Design, train, fine-tune, and optimize LLMs and VLMs for real-world • Agentic AI Systems: Develop and orchestrate autonomous agent frameworks capable of multi-step reasoning, planning, and tool use. • Engineering & Deployment: Build scalable, low-latency inference systems for large models using frameworks like DeepSpeed, vLLM, TensorRT, or ONNX Runtime. Implement distributed training, model parallelism, and efficient inference pipelines; also optimize deployment for edge devices, GPUs, and cloud-based platforms. • Research & Innovation: Stay up to date with the latest advancements in LLMs, multimodal models, and autonomous agents. Core Competencies • AI/ML Expertise o Strong understanding of LLMs, VLMs, transformers, and multimodal architectures. o Experience with fine-tuning, LoRA/QLoRA, quantization, distillation, and evaluation. o Knoledge of neurosymbolic methodologi o Knowledge of reinforcement learning (RLHF, RLAIF) and alignment techniques. • Agentic Frameworks o Experience with frameworks such as LangChain, LlamaIndex, AutoGPT, CrewAI, OpenAI Agents, Hugging Face Transformers/Agents. o Ability to design reasoning loops, memory systems, and multi-agent coordination. • Development Tools & Libraries o Core AI frameworks: PyTorch, TensorFlow, Hugging Face, OpenAI APIs, DeepSpeed, vLLM. o Supporting tools: Weaviate, Pinecone, FAISS, Milvus (vector databases), Redis, Kafka. o Evaluation/monitoring: Weights & Biases, MLflow, TensorBoard, Evals frameworks. • Programming Skills o Python – for AI research, prototyping, and deployment pipelines. o C++ – for performance-critical components, model inference optimization, and system integration. • Systems & Infrastructure o Proficiency with Docker, and AI distributed training systems. o Strong knowledge of CUDA, GPU optimization, and high-performance computing. o Familiarity with cloud platforms (AWS, GCP, Azure) and edge deployment strategies. Qualifications • Master’s, or PhD in Computer Science, AI/ML, Robotics, or related field. • Proven track record of hands-on work with LLMs, VLMs, or agentic frameworks. • Experience in productionizing AI systems at scale. • Excellent communication and collaboration skills.

Preferred (Nice-to-Have) • Experience with reinforcement learning • Background in robotics, simulation environments, or embodied AI. • Publications in AI conference

AI Engineering Lab Standard

Every lab in this track builds a runnable, test-verified miniature of a real AI mechanism — the attention block, the autograd tape, the LoRA delta, the DPO loss, the PagedAttention block table, the all-reduce ring, the ANN index, the ReAct loop — so the internals become muscle memory instead of framework folklore. You do not "call model.generate()"; you build the sampler that turns logits into the next token. That is what makes the knowledge defensible in a senior interview and useful at 2 a.m. when generation is wrong, slow, or unsafe.

Why miniatures (and not a 70B model on a GPU cluster)

A real training/serving stack is non-deterministic, costs thousands of dollars, needs GPUs and credentials, and hides the mechanism behind CUDA kernels and a framework API. A lab that reimplements the algorithm — the softmax-with-max-subtraction, the reverse-mode chain rule, the NF4 quantization grid, the Bradley-Terry gradient, the KV-block allocator, the ring-reduce schedule, the HNSW greedy descent — is offline, deterministic, free, and teaches the part interviewers actually probe. Each lab README ends with a "How this maps to the real stack" section that ties the miniature back to PyTorch / vLLM / DeepSpeed / FAISS / LangChain, its limits, and the real API equivalent.

Required files (per lab)

The runnable core is Python (stdlib + pytest), offline, deterministic. No GPU, no CUDA, no network, no model download, no pip install torch, no Date.now()/unseeded random. Where a lab needs vectors or matrices, implement them with stdlib lists/array/math so the mechanism is visible — NumPy is allowed only if a lab explicitly declares it in its requirements.txt, and even then the algorithm (not a one-line library call) must be the thing the learner writes. Multi-language work (the C++/CUDA kernel, the real torch/vllm run) appears only as build specs in the README "Extensions" section for the learner's own hardware.

Determinism rules (these labs touch probability — be careful)

• Any randomness goes through an explicit, seeded random.Random(seed) passed in or defaulted; same seed → same bytes. Tests assert this.
• Floating-point comparisons in tests use pytest.approx (or an explicit abs-tolerance), never == on floats.
• softmax, log-sum-exp, and any loss must use the max-subtraction / log-space trick so the tests can include large-logit inputs without overflow — that numerical-stability detail is part of the lesson.

The test contract

import importlib, os
lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab"))

Run both ways; the reference must pass, and your lab.py passes once the TODOs are filled:

pytest test_lab.py -v                       # against your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # against the reference (must be green)
python solution.py                          # the worked example

Test taxonomy every flagship lab includes: happy path; malformed / out-of-range input (raises ValueError); boundary cases — the off-by-one interviewers probe (empty sequence, single token, top_k=1 ⇒ greedy, temperature→0, LoRA r=0, a full KV-block, an empty retrieval set, a degenerate preference pair); numerical-stability cases (large logits don't overflow, log(0) is guarded); invariants (softmax sums to 1, probabilities are non-negative, a merged-LoRA forward equals base+delta, dequant∘quant is within the error bound, all-reduce result equals the serial sum); and deterministic output (same seed → same bytes).

The four teaching docs (per phase)

Every phase explains the same material through four lenses, because senior-level understanding is their intersection:

Doc conventions (mdBook-compatible)

• WARMUP.md opens with a Table of Contents of working anchor links, kept in sync with the headings.
• MathJax for any math (\( … \) inline, $$ … $$ block). Standard Markdown only.
• No bare angle brackets in prose — wrap <like-this> in backticks so mdBook does not eat them as HTML. (This bites constantly in this domain: <bos>, <image>, <|endoftext|>, generic types — always fence them.)
• File references use relative links. Code fences are language-tagged.

Definition of done

A lab is complete only when: the reference suite passes (LAB_MODULE=solution pytest), the learner lab.py passes after the TODOs are filled in, python solution.py prints a sensible worked example, the README's success criteria are testable, and the "How this maps to the real stack" section is accurate about the production framework and its limits. A phase is complete when all four teaching docs exist, the WARMUP's ToC anchors resolve, and every lab in it is done.

Glossary — Senior AI Engineer

The running vocabulary for the whole track. Terms are grouped by area; each entry is the one-sentence definition plus the why it matters a senior is expected to know. Where a phase builds the mechanism, the phase number is noted.

Foundations, math & scaling (P00)

• Parameter (N) — a learned weight. Model "size" is the parameter count; memory for weights ≈ N × bytes_per_param (2 for fp16, 1 for int8, ~0.5 for int4).
• Token (D in training, T in a sequence) — the atomic unit the model reads/writes (a sub-word, byte, or image patch). Training cost scales with total tokens seen.
• Training FLOPs ≈ 6 × N × D — the rule of thumb: forward+backward over D tokens for an N-param dense model. Forward alone ≈ 2N per token.
• Inference FLOPs ≈ 2 × N per generated token — one forward pass per token; prefill processes the whole prompt at once, decode emits one token at a time.
• Chinchilla-optimal — for a fixed compute budget, the loss-minimizing split is roughly D ≈ 20 × N tokens per parameter; most "big" models were under-trained on data.
• Scaling law — loss falls as a power law in compute/params/data: ( L(N) = L_\infty + (N_c/N)^{\alpha} ). The reason "make it bigger" worked.
• FLOPs vs memory-bandwidth bound — prefill and training are compute-bound; autoregressive decode is memory-bandwidth bound (you re-read the whole weight matrix per token).
• Arithmetic intensity — FLOPs per byte moved; where it sits relative to the hardware roofline ridge tells you whether you're compute- or bandwidth-limited (P17).
• KV-cache — the stored keys/values of every past token so decode is O(1) per step instead of O(T). Size ≈ 2 × layers × T × d_model × bytes; the real inference memory wall.
• MFU (Model FLOPs Utilization) — achieved FLOPs ÷ peak hardware FLOPs; 40–55% is good for training, far lower for memory-bound decode.

Tokenization & input (P01)

• BPE (Byte-Pair Encoding) — learn a merge table by greedily merging the most frequent adjacent pair; encode by applying merges in learned priority order.
• Byte fallback — represent any unknown character as its raw UTF-8 bytes so the vocab is closed (no [UNK] data loss).
• Special tokens — <bos>, <eos>, <pad>, <|im_start|>, <image> — reserved IDs the model is trained to treat structurally; mishandling them is the #1 chat bug.
• Chat template — the exact string format (roles, turn markers) the model was trained on; the prompt you send must match it byte-for-byte or quality collapses.
• Context window — max tokens the model attends over; the KV-cache and attention cost scale with it.

Transformer internals (P02)

• Embedding — lookup table mapping token ID → vector of size d_model.
• Self-attention — ( \text{softmax}(QK^\top/\sqrt{d_k})V ); each token mixes information from others, weighted by query-key similarity.
• Causal mask — set future positions to -∞ before softmax so a token can't see the future (required for autoregressive LMs).
• Multi-head attention (MHA) — run attention in h parallel subspaces and concatenate; lets the model attend to several relations at once. MQA/GQA share K/V heads to shrink the KV-cache.
• RoPE (Rotary Position Embedding) — encode position by rotating Q/K by an angle proportional to position; relative, extrapolates better than learned absolute embeddings.
• RMSNorm / LayerNorm — normalize activations for stable training; RMSNorm drops the mean subtraction (cheaper, ~same quality).
• SwiGLU / GeLU FFN — the per-token MLP (~2/3 of params); SwiGLU is a gated activation used in Llama-class models.
• Residual stream — the additive x + sublayer(x) highway every block reads/writes; the conceptual "bus" of the transformer.
• Softmax stability — subtract max(logits) before exp to avoid overflow; mathematically identical, numerically safe.

Training (P03, P10)

• Autograd / reverse-mode AD — build a computation graph on the forward pass, then apply the chain rule backward to get every gradient in one pass.
• Cross-entropy loss — ( -\sum y \log \hat{p} ); the LM objective (next-token prediction).
• Adam / AdamW — adaptive optimizer tracking 1st/2nd moment estimates; AdamW decouples weight decay. Optimizer state is 2× the params in memory (the ZeRO motivation).
• LR schedule — warmup then cosine/linear decay; the single biggest non-architecture knob.
• Gradient clipping — cap the gradient norm to stop loss spikes / divergence.
• Mixed precision (AMP) — compute in fp16/bf16, keep an fp32 master copy + loss scaling.
• Gradient accumulation — sum grads over micro-batches to simulate a big batch on small memory.
• Data parallel (DP) — replicate the model, split the batch, all-reduce gradients.
• Tensor parallel (TP) — shard a single matmul across GPUs (Megatron); needs all-reduce inside the layer.
• Pipeline parallel (PP) — split layers across GPUs, stream micro-batches; the idle gap is the bubble (1F1B scheduling shrinks it).
• ZeRO — partition optimizer state (stage 1), gradients (2), and params (3) across data- parallel ranks to cut per-GPU memory; the core of DeepSpeed.
• All-reduce — every rank ends with the sum of all ranks' tensors; ring all-reduce does it in 2(N−1) steps moving ~2× the data, bandwidth-optimal.
• FSDP — PyTorch's native ZeRO-3-style fully-sharded data parallel.

Multimodal / VLMs (P04)

• ViT (Vision Transformer) — split an image into patches, linearly embed each as a "token," run a transformer; images become sequences.
• CLIP — train an image encoder and text encoder so matching pairs have high cosine similarity (contrastive InfoNCE loss over a batch); enables zero-shot classification and retrieval.
• Projector / connector — the MLP that maps vision-encoder features into the LLM's token embedding space (LLaVA pattern), so images become "tokens" the LLM reads.
• InfoNCE / contrastive loss — cross-entropy over a similarity matrix where the diagonal (matched pairs) is the target; temperature τ sharpens it.
• Cross-attention fusion — alternative to projection: the LLM attends to vision features via added cross-attention layers (Flamingo).
• VLM — Vision-Language Model: takes interleaved image+text, outputs text (captioning, VQA, grounding).

Fine-tuning & compression (P05, P06)

• PEFT — Parameter-Efficient Fine-Tuning: adapt a model by training a tiny fraction of weights.
• LoRA — freeze W, learn a low-rank update ( \Delta W = \frac{\alpha}{r} BA ) with B∈ℝ^{d×r}, A∈ℝ^{r×k}; trains <1% of params, merges back at inference for zero overhead.
• QLoRA — LoRA on top of a 4-bit (NF4) quantized frozen base with double quantization and paged optimizers; fine-tune a 65B model on one GPU.
• Knowledge distillation (KD) — train a small student to match a large teacher's soft probabilities; loss = KL of softened logits at temperature T (scaled by T²).
• Quantization — store/compute weights (and sometimes activations) in fewer bits. Affine: q = round(x/scale) + zero_point; symmetric: zero_point = 0.
• Per-tensor vs per-channel — one scale for the whole tensor vs one per output channel; per-channel is far more accurate for weights.
• PTQ vs QAT — Post-Training Quantization (calibrate, no retrain) vs Quantization-Aware Training (simulate quant during training).
• GPTQ — second-order, layer-wise weight quantization minimizing output error.
• AWQ — Activation-aware Weight Quantization: scale up salient weight channels before quantizing so important weights keep precision.
• Outlier / salient channels — a few activation channels with huge magnitude that dominate quant error; the thing AWQ/SmoothQuant protect.

Alignment & RL (P07, P15)

• SFT — Supervised Fine-Tuning on instruction/response pairs; the first alignment step.
• RLHF — Reinforcement Learning from Human Feedback: train a reward model on human preferences, then optimize the policy against it with PPO + a KL penalty to the SFT model.
• RLAIF — same loop, with an AI model (a "constitution") generating the preferences instead of humans.
• Reward model — predicts a scalar preference; trained with the Bradley-Terry loss ( -\log \sigma(r_{\text{chosen}} - r_{\text{rejected}}) ).
• DPO (Direct Preference Optimization) — skip the RL loop; optimize the policy directly on preference pairs with a closed-form loss using the SFT model as reference and temperature β. Simpler and more stable than PPO.
• KL penalty / reference model — keep the aligned policy from drifting too far from the base, preventing reward hacking and gibberish.
• PPO — Proximal Policy Optimization: clipped policy-gradient update; the classic RLHF optimizer.
• Reward hacking — the policy exploits flaws in the reward model to score high while getting worse; the KL term and good eval fight it.
• MDP / Q-learning / policy gradient — the RL substrate (states, actions, rewards); Q-learning learns action values, policy gradient learns the policy directly (P15).

Decoding & generation (P08)

• Logits — the model's raw pre-softmax scores over the vocab for the next token.
• Greedy / argmax — always take the highest-probability token; deterministic, repetitive.
• Temperature — divide logits by T before softmax; T<1 sharpens, T>1 flattens, T→0 ⇒ greedy.
• Top-k — sample only from the k highest-probability tokens.
• Top-p (nucleus) — sample from the smallest set whose cumulative probability ≥ p.
• Min-p — keep tokens with prob ≥ min_p × max_prob; adapts the cutoff to confidence.
• Repetition / presence penalty — down-weight already-emitted tokens to reduce loops.
• Constrained / structured decoding — mask logits to only-valid tokens via a grammar/FSM so output is guaranteed-valid JSON or matches a schema (the backbone of reliable tool calls).
• Beam search — keep the b best partial sequences; better for translation, worse for open-ended/creative generation.

Inference serving (P09, P17)

• Prefill vs decode — prefill processes the prompt in one parallel pass (compute-bound); decode emits tokens one at a time (memory-bound). Different bottlenecks, different tuning.
• PagedAttention — vLLM's trick: store the KV-cache in fixed-size blocks (like OS paging) so memory isn't fragmented and sequences can share prefixes; enables high batch occupancy.
• Continuous (in-flight) batching — add/remove sequences from the running batch every step instead of waiting for the slowest; 2–4× throughput over static batching.
• Speculative decoding — a small draft model proposes k tokens, the big model verifies them in one pass; accept the longest correct prefix. 2–3× speedup with identical output distribution.
• TTFT / TPOT / ITL — Time-To-First-Token (prefill latency), Time-Per-Output-Token, Inter-Token Latency; the two SLOs that matter for serving.
• Throughput vs latency — bigger batches raise tokens/sec (throughput) but raise per- request latency; the core serving tradeoff.
• vLLM / TensorRT-LLM / DeepSpeed-Inference / ONNX Runtime — the production serving stacks the JD names; all implement variants of the above.

Retrieval & RAG (P11)

• Embedding — a dense vector encoding semantic meaning; similarity = cosine / dot product.
• ANN (Approximate Nearest Neighbor) — sub-linear similarity search trading a little recall for huge speed. HNSW (navigable small-world graph) and IVF (inverted-file clusters) are the two dominant families.
• Recall@k / nDCG — fraction of true neighbors found in the top-k / rank-weighted gain; how you measure a retriever.
• Chunking — splitting documents into passages; size and overlap trade recall vs precision and cost.
• Reranking — re-score retrieved candidates with a stronger (cross-encoder) model; MMR re-ranks for relevance and diversity.
• Vector DB — FAISS (library), Pinecone/Weaviate/Milvus (managed/serving); store embeddings
  • metadata, serve ANN queries with filters.
• RAG — Retrieval-Augmented Generation: retrieve relevant context, put it in the prompt, generate grounded answers; the dominant pattern for factuality.

Agents (P12, P13, P14)

• Agent — an LLM in a loop that can take actions (tool calls) and observe results until a goal is met.
• ReAct — Reason+Act: interleave a "thought," an "action" (tool call), and an "observation" each step; the canonical agent loop.
• Tool / function calling — the model emits a structured call (name + JSON args) that the runtime executes and feeds back; typed schemas + validation keep it safe.
• Memory — short-term (the running transcript), summarized (compressed history), and long-term/semantic (vector-recalled facts); how an agent persists beyond the context window.
• Planning — decomposing a goal into steps (Plan-and-Execute, Tree-of-Thoughts, hierarchical task networks).
• Multi-agent — multiple specialized agents (planner/executor/critic) coordinating via a shared blackboard or message bus; frameworks: CrewAI, AutoGen, LangGraph.
• Reflection / critique-revise — an agent (or critic agent) reviews and improves a draft; the cheapest reliable quality boost.
• Neurosymbolic — couple a neural proposer (LLM) with a symbolic verifier (rule engine, solver, type checker) so answers are checked, not just generated (P14).

Evaluation, observability & safety (P16)

• Exact match / F1 — string-overlap metrics for closed-form answers.
• pass@k — probability that at least one of k samples passes (code/agent eval).
• LLM-as-judge — use a strong model with a rubric to score open-ended outputs; cheap, scalable, biased — calibrate against humans.
• Calibration / ECE — does a model's confidence match its accuracy? Expected Calibration Error bins predictions and measures the gap.
• Guardrails — input/output filters (PII, jailbreak, toxicity, schema validation) wrapping the model.
• Run tracking (W&B / MLflow) — log params, metrics, artifacts per run for reproducibility and comparison; the experiment system of record.
• Hallucination — confident, fluent, wrong; the failure mode RAG, grounding, and verification fight.
• Drift — input/output distribution shifts over time; monitored in production to catch silent quality decay.

Systems & hardware (P17)

• CUDA / kernel / thread / warp / block / SM — the GPU execution hierarchy; a kernel runs on a grid of blocks of threads scheduled in 32-thread warps on Streaming Multiprocessors.
• Occupancy — active warps per SM ÷ max; too-low occupancy starves the latency-hiding that makes GPUs fast.
• Memory coalescing — adjacent threads reading adjacent memory = one transaction; the #1 GPU performance lever.
• Roofline model — plot of attainable FLOPs vs arithmetic intensity; the ridge point is where you flip from bandwidth- to compute-bound.
• Operator fusion — combine ops (e.g. matmul+bias+GeLU) into one kernel to cut memory round-trips; what torch.compile/FlashAttention do.
• FlashAttention — fused, IO-aware attention that never materializes the full T×T score matrix; the reason long context is feasible.

Cheat Sheet — Senior AI Engineer

The numbers, formulas, and one-liners to tattoo on your arm. Pair with GLOSSARY.md.

The arithmetic that opens every design review (P00)

Decode is memory-bandwidth bound. tokens/sec ≈ mem_bandwidth / bytes_per_token_read. A 7B fp16 model on 1 TB/s HBM: 1e12 / 14e9 ≈ 70 tok/s single-stream ceiling — batching amortizes the weight read across requests, which is the whole point of P09.

Softmax & attention (P02) — the numerically-safe form

softmax(x)_i = exp(x_i - max(x)) / Σ_j exp(x_j - max(x))      # subtract max → no overflow
Attention(Q,K,V) = softmax( Q Kᵀ / √d_k  +  causal_mask ) V    # mask future = -inf

• √d_k keeps the dot products from saturating softmax as d_k grows.
• MHA: split d_model into h heads of d_k = d_model/h, attend, concat, project.
• GQA/MQA: fewer K/V heads than Q heads ⇒ smaller KV-cache (the modern default).
• RoPE: rotate (q,k) pairs by θ_pos; relative position falls out of the dot product.

Autograd & training (P03, P10)

• Reverse-mode: forward builds the graph; backward does grad_input = grad_output · ∂out/∂in in reverse topological order. One backward pass ⇒ all gradients.
• Cross-entropy + softmax gradient simplifies to (p − y) — memorize this.
• AdamW step: m=β₁m+(1−β₁)g; v=β₂v+(1−β₂)g²; θ −= lr·(m̂/(√v̂+ε)) + lr·wd·θ.
• Ring all-reduce: 2(N−1) steps, each moving tensor/N; total ≈ 2·tensor bytes/GPU regardless of N — bandwidth-optimal.
• ZeRO memory: stage 1 shards optimizer (÷N), stage 2 +gradients, stage 3 +params.
• Pipeline bubble fraction ≈ (P−1)/(M + P−1) for P stages, M micro-batches — more micro-batches ⇒ smaller bubble.

Fine-tuning & quantization (P05, P06)

• LoRA: W' = W + (α/r)·B·A; trainable params = r·(d+k) vs d·k — orders of magnitude fewer. r=8, α=16 is a common default; merge B·A into W for zero inference cost.
• QLoRA: NF4 4-bit base (frozen) + double-quant scales + paged optimizer + LoRA adapters.
• KD loss: T² · KL(softmax(z_t/T) ‖ softmax(z_s/T)) + α·CE(y, z_s) — the T² keeps gradient magnitude scale-stable.
• Affine quant: q = clamp(round(x/s) + z, qmin, qmax), x̂ = s·(q − z). Symmetric ⇒ z=0.
• Per-channel (weights) beats per-tensor; keep activation outliers in higher precision (SmoothQuant/AWQ) or accuracy tanks.

Alignment (P07)

• Bradley-Terry reward loss: −log σ(r(chosen) − r(rejected)).
• DPO loss: −log σ( β·[ (logπ_θ(y_w) − logπ_ref(y_w)) − (logπ_θ(y_l) − logπ_ref(y_l)) ] ). No reward model, no RL loop, uses the frozen reference policy. β≈0.1 typical.
• PPO objective: clipped min(ratio·Â, clip(ratio,1±ε)·Â) − λ·KL(π_θ ‖ π_ref).
• Always keep the reference/KL term or the policy reward-hacks into gibberish.

Decoding (P08)

top_k=1 ≡ greedy. temperature → 0 ≡ greedy. Constrained decode > "please return JSON."

Serving (P09)

• PagedAttention: KV in fixed blocks (e.g. 16 tokens) ⇒ ~0 fragmentation, prefix sharing, high batch occupancy.
• Continuous batching: 2–4× throughput vs static — sequences join/leave the batch per step.
• Speculative decoding: draft k tokens with a small model, verify in one big-model pass, accept longest correct prefix ⇒ 2–3× faster, identical output distribution.
• SLOs: TTFT (prefill) for responsiveness, TPOT (decode) for streaming feel.
• Throughput↑ ⇒ batch↑ ⇒ latency↑. Pick the SLO first, then the max batch that meets it.

Retrieval / RAG (P11)

• Similarity: cosine = (a·b)/(‖a‖‖b‖); normalize once, then it's a dot product.
• HNSW: greedy descent through a multi-layer small-world graph; M (degree) and ef (search width) trade recall vs latency. IVF: cluster, search nprobe nearest clusters.
• Eval the retriever (recall@k, nDCG) separately from the generator (faithfulness) — most "RAG is bad" bugs are retrieval bugs.
• MMR rerank: argmax λ·rel(d) − (1−λ)·max_sim(d, selected) — relevance and diversity.

Agents (P12–P14)

• ReAct step: Thought → Action(tool, args) → Observation → …→ Final. Cap the loop (max steps) or it runs forever / burns tokens.
• Tools: typed JSON-schema args, validate before executing, return errors as observations so the agent can recover.
• Memory tiers: buffer (recent turns) → summary (compressed) → semantic (vector recall).
• Multi-agent: planner decomposes, executors act, critic verifies; a blackboard holds shared state. Termination detection or it deadlocks.
• Neurosymbolic: LLM proposes, a deterministic verifier (rules/solver/types) accepts or rejects → trustworthy answers.

Evaluation & safety (P16)

• Eval the retriever, the generator, and the system separately; one number hides regressions.
• LLM-as-judge: cheap and scalable but biased (position, verbosity, self-preference) — randomize order, use a rubric, spot-check against humans.
• ECE: bin by confidence, measure |accuracy − confidence| per bin, weight by bin size.
• Guardrails are layered: input filter → constrained decode → output validator → human review for high-stakes.

GPU performance (P17)

• Roofline: attainable = min(peak_FLOPs, bandwidth × arithmetic_intensity); left of the ridge = bandwidth-bound, right = compute-bound.
• Coalesced access (adjacent threads → adjacent addresses) can be 10× uncoalesced.
• Occupancy = active warps / max warps per SM; raise it by cutting registers/shared-mem per thread — but past "enough to hide latency" it stops helping.
• Fuse to cut HBM round-trips (FlashAttention, torch.compile); the bottleneck is usually memory traffic, not math.

The senior tradeoff one-liners

• "Make it bigger" trades inference cost forever for one-time training gains — quantize and distill before you scale.
• "Add an agent step" trades latency and a failure mode for autonomy — verify every tool call.
• "Multi-region / multi-GPU" trades complexity and $ for availability/throughput — show the SLA and MFU math first.
• "Lower the temperature" trades diversity for reliability — and constrained decoding beats both for structured output.
• Quality is a distribution: report a metric with a confidence interval, never a single cherry-picked generation.

Interview Prep — Senior AI Engineer (LLMs, VLMs, Agents)

A running reference for the questions a strong AI/research panel actually asks — and the level of answer that separates "uses the API" from "owns the model stack." Pair each section with the phase that builds the mechanism, so your answer is "I implemented this," not "I read about it."

How senior AI interviews differ

A junior candidate explains what a model/framework does. A senior explains how it works under the hood, when it breaks, what it costs in FLOPs/bytes/latency, and what they'd trade. The loop in a strong panel is usually: (1) a from-scratch coding round (implement attention / a sampler / LoRA / the DPO loss), (2) an ML-depth round (derive a gradient, reason about a training failure), (3) a serving/systems round (size a deployment, debug latency), (4) an agent/system-design round, and (5) research taste / "what did you build and why." Every answer should end with a tradeoff and, where possible, a number.

1. Foundations & scaling (P00)

• "How many FLOPs to train a 7B model on 1T tokens? How much GPU memory to serve it?" — 6ND ≈ 4.2e22 to train; weights 14 GB fp16 (7 GB int8) plus the KV-cache, which grows with batch×context — show that the KV-cache, not the weights, is what OOMs you.
• "Is decoding compute- or memory-bound? Why does that decide your whole serving strategy?" — memory-bandwidth bound; you re-read all weights per token, so batching (amortize the read) and quantization (fewer bytes) are the levers, not more FLOPs.
• "Chinchilla — what did it change?" — compute-optimal is ~20 tokens/param; most early large models were under-trained, so a smaller model on more data beats a bigger under-trained one at equal compute and is cheaper to serve.
• "Train a bigger model or fine-tune a smaller one? Defend with numbers." — inference cost is forever (2N/token × every request); quantize/distill/PEFT a small model unless the quality gap is measured and the volume justifies the perpetual bill.

2. Transformer internals (P02)

• "Implement attention on a whiteboard. Why √d_k? Why subtract the max in softmax?" — scale stops dot products from saturating softmax; max-subtraction is overflow safety (mathematically identical). Be ready to write the causal mask.
• "MHA vs MQA vs GQA — what's the tradeoff?" — fewer K/V heads shrink the KV-cache (the serving memory wall) at a small quality cost; GQA is the modern compromise.
• "Why RoPE over learned absolute positions?" — relative, parameter-free, extrapolates to longer context; explain the rotation-of-(q,k) intuition.
• "Where do the parameters live in a transformer block?" — ~2/3 in the FFN, ~1/3 in attention projections; matters for what to shard (TP) and what to quantize.

3. Training & distributed (P03, P10)

• "Derive the gradient of softmax+cross-entropy." — it collapses to (p − y); if you can't show this you haven't earned the training round.
• "Your loss spikes to NaN at step 4000. Debug it." — LR too high / no warmup, missing grad clip, fp16 overflow (need loss scaling/bf16), bad data batch; bisect by reverting the LR schedule and checking grad norms.
• "A 70B model won't fit for training. Walk me through the parallelism." — DP for throughput, TP to split the matmuls within a node (high comms ⇒ NVLink), PP across nodes, ZeRO/FSDP to shard optimizer+grads+params; name where the bubble and the all-reduce live.
• "Explain ring all-reduce and why it's bandwidth-optimal." — 2(N−1) steps each moving 1/N of the tensor; per-GPU traffic is ~2× the tensor regardless of N.

4. Multimodal / VLMs (P04)

• "How does an LLM 'see' an image?" — a vision encoder (ViT) produces patch features; a projector maps them into the token-embedding space; they're prepended as "tokens" the LLM attends over (LLaVA), or fused via cross-attention (Flamingo).
• "Explain CLIP's training objective." — contrastive InfoNCE: a batch's similarity matrix, diagonal = matched pairs = the target; temperature sharpens it. Enables zero-shot + retrieval.
• "Projector tuning vs full fine-tune for a new VLM task?" — freeze encoders, train the projector (+ LoRA on the LLM) first; cheap, preserves pretraining, usually enough.

5. Fine-tuning & compression (P05, P06)

• "How does LoRA work and why is inference free after merging?" — ΔW = (α/r)BA, train r·(d+k) params; merge into W so the served model is identical-shape, zero added latency.
• "QLoRA — how do you fine-tune 65B on one GPU?" — NF4 4-bit frozen base + double quant + paged optimizer states + LoRA adapters in bf16; only the adapters get gradients.
• "int8 vs int4 — when does quality fall off, and how do you protect it?" — int8 ~lossless with per-channel weights; int4 needs outlier handling (AWQ/GPTQ); activations are harder than weights. Always measure on your eval, not perplexity alone.
• "Distillation — what's the loss and why temperature?" — KL of softened teacher/student logits (×T²); T exposes the teacher's "dark knowledge" (relative wrong-class probs).

6. Alignment (P07)

• "RLHF in three stages." — SFT → reward model (Bradley-Terry on preferences) → PPO with a KL penalty to the SFT policy.
• "DPO vs PPO — why did the field move to DPO?" — DPO has a closed-form loss on preference pairs (no reward model, no rollout, no RL instability) using the frozen reference policy and temperature β; simpler, cheaper, more stable — at some loss of online exploration.
• "What's reward hacking and how do you prevent it?" — the policy exploits reward-model flaws; the KL/reference term and held-out human eval are the guards.

7. Decoding (P08)

• "top-k vs top-p vs min-p — pick one for a JSON-extraction service and defend it." — temperature=0 + constrained decoding (mask to schema-valid tokens), not sampling at all; sampling is for open-ended generation.
• "How do you guarantee valid JSON / a valid function call?" — a grammar/FSM that masks the logits to only valid next tokens each step; prompting alone is best-effort.

8. Serving (P09, P17)

• "Explain PagedAttention and continuous batching like I'm a new hire." — KV-cache in fixed blocks like OS pages (no fragmentation, prefix sharing) + sequences joining/leaving the batch every step (no head-of-line blocking) ⇒ much higher GPU occupancy.
• "Speculative decoding — does it change the output?" — no; the verification step makes the accepted tokens distributionally identical to the target model; you just get them 2–3× faster.
• "Your p99 TTFT is fine but TPOT is bad. Diagnose." — decode is memory-bound; check batch occupancy, KV-cache pressure (evictions), quantize the KV/weights, or shard with TP.
• "DeepSpeed vs vLLM vs TensorRT-LLM vs ONNX Runtime — when each?" — vLLM for throughput serving (PagedAttention), TensorRT-LLM for lowest latency on NVIDIA (compiled kernels), DeepSpeed for training+inference at scale, ONNX Runtime for portability/edge.

9. RAG & retrieval (P11)

• "RAG quality is bad. Is it retrieval or generation?" — measure separately: recall@k for the retriever, faithfulness for the generator. Most bugs are retrieval (chunking, embedding mismatch, missing reranker).
• "HNSW vs IVF vs flat — pick for 100M vectors at p99 < 50ms." — HNSW for low-latency high-recall in memory; IVF/IVF-PQ when memory-bound; flat only for small/exact. Name the recall–latency–memory triangle.

10. Agents (P12–P14)

• "Design an agent that books travel. What can go wrong and how do you contain it?" — typed tools with validation, a step cap, errors-as-observations, human approval for irreversible actions, and a verifier/critic; idempotent tool calls so a retry doesn't double-book.
• "Single agent vs multi-agent — when is multi-agent actually worth it?" — when roles are genuinely distinct (planner/executor/critic) or tasks parallelize; otherwise it adds latency, cost, and coordination bugs for little gain. Have a number.
• "Where does neurosymbolic help?" — when answers must be verifiable (math, code, config, constraints): LLM proposes, a deterministic checker accepts/rejects; you get LLM flexibility with symbolic correctness.

11. Evaluation & safety (P16)

• "How do you know a model change is actually better?" — a fixed eval set, the right metric per task, a confidence interval, and a regression gate in CI — never a vibe check on one prompt.
• "LLM-as-judge — what are its biases and how do you control them?" — position, verbosity, self-preference; randomize order, use a rubric, calibrate against human labels.

The system-design loop (use in every round)

1. Clarify the workload: task, modality (text/image/both), QPS, prompt/response token sizes, latency SLO (TTFT/TPOT), quality bar, safety/compliance, budget.
2. Pick the model strategy: API vs open-weights; size; fine-tune (PEFT) vs prompt vs RAG; quantization target. Justify with the FLOPs/memory/cost math.
3. Design the serving path: tokenizer → model (sharding?) → sampler/constraints → cache → batching → autoscaling; and the offline path (training/fine-tune/eval pipeline).
4. Add retrieval/agents only if needed, and design their failure containment.
5. Quantify: tokens/sec, $/1M tokens, KV-cache memory, recall@k, eval score with CI.
6. Fail it: hallucination, prompt injection, tool misuse, cost blowup, drift, the rollback.

See ../system-design/README.md for worked reference designs.

System Design — Reference AI/LLM Architectures

Worked, senior-level designs that tie the whole track together. Each one states the workload numbers first, picks a model strategy (the FLOPs/memory/cost math), designs the serving path and the offline path, then quantifies latency, cost, and quality and names the tradeoff. Use them as templates; the capstone (P19) builds one end-to-end.

The design method (every time)

1. Workload facts first: task + modality, QPS / peak concurrency, prompt and response token sizes, latency SLO (TTFT and TPOT separately), quality bar + eval, safety/compliance, budget.
2. Two paths, always:
   • Serving (online) — tokenizer → model (quantized? sharded?) → sampler/constraints → KV-cache/PagedAttention → continuous batching → autoscaling → guardrails → observability.
   • Offline — data → fine-tune/align (PEFT/DPO) → quantize → eval gate → registry → deploy; plus the RAG indexing pipeline if there is retrieval.
3. Quantify: tokens/sec per GPU, KV-cache bytes at peak concurrency, $/1M tokens, p99 TTFT/TPOT, recall@k, eval score with a confidence interval.
4. Tradeoff: quality vs latency vs cost vs memory vs safety — name what you trade and write the one-sentence decision record.
5. Fail it: hallucination, prompt injection, tool misuse, KV-cache OOM, cost blowup, model regression, drift — and the rollback.

Design 1 — High-throughput LLM serving platform (the canonical one)

Requirement: serve a 13B chat model, 2k concurrent users, p99 TTFT < 1s, smooth streaming (TPOT < 50ms), minimize $/1M tokens.

Model strategy

• Open-weights 13B, int8 (or int4-AWQ) weights ⇒ ~13 GB → ~7–13 GB; fits one 24–48 GB GPU with room for KV-cache. Quantize after measuring eval impact, not before.

Serving path

• vLLM-class server: PagedAttention KV-cache (blocked, prefix-shared) + continuous batching to keep the GPU busy; tensor-parallel across 2 GPUs only if a single one can't hold weights+KV at target concurrency.
• Sampler: temperature 0.7 / top_p 0.9 for chat; a separate temperature=0 + constrained path for tool/JSON endpoints.
• Autoscale on queue depth / KV-cache utilization, not CPU. Admission-control + reject with 429 past the SLO-preserving batch size.

Quantify

• KV-cache: 2 · L · T · d_model · bytes per sequence — at 2k concurrency this, not the weights, sets the GPU count. Show the number; it's the whole capacity plan.
• Throughput: batching amortizes the per-token weight read; report tokens/sec/GPU and $/1M.

Tradeoffs

• Quality vs cost: int4 + draft-model speculative decoding cuts $/token but must pass the eval gate. Latency vs throughput: cap the batch at the largest size that still meets p99.

Design 2 — Production RAG over a 10M-document corpus (P11, P16)

Requirement: grounded Q&A with citations, p99 < 2s end-to-end, "I don't know" instead of hallucination, nightly corpus refresh.

• Indexing (offline): chunk (semantic, ~512 tokens, overlap), embed, build an HNSW (or IVF-PQ if memory-bound) index in a vector DB (FAISS/Milvus/Pinecone/Weaviate) with metadata filters; version the index so a bad rebuild can roll back.
• Query (online): embed → ANN top-k (k≈20) → cross-encoder rerank to top-5 → assemble a grounded prompt with citations → constrained answer that must cite or abstain.
• Eval: retriever (recall@k, nDCG) and generator (faithfulness, citation accuracy) measured separately; gate deploys on both.
• Tradeoff: bigger k + reranker raises recall and cost/latency; abstention raises trust but lowers coverage — tune to the business cost of a wrong answer.
• Fail it: stale index, embedding/model version skew, prompt injection inside retrieved docs (treat retrieved text as untrusted), citation hallucination.

Design 3 — Multimodal VLM assistant (P04)

Requirement: answer questions about uploaded images + text, single model, p99 < 3s.

• Architecture: ViT vision encoder → projector → 13B LLM (LLaVA pattern); images become prepended tokens. Image tokens inflate the context — budget the KV-cache for them.
• Adaptation: freeze encoders, train the projector + a LoRA on the LLM for the domain; cheap and preserves pretraining.
• Serving: vision encode is a separate compute-bound stage (cache by image hash); the LLM decode is the memory-bound stage — pipeline them.
• Tradeoff: more image patches = better grounding, more tokens = more latency/cost; pick resolution to the task.

Design 4 — Production agent with tools (P12–P14, P16)

Requirement: an agent that resolves customer tickets using internal tools (lookup, refund, escalate); must be safe, auditable, and bounded in cost.

• Loop: ReAct with a hard step cap and a token budget; typed JSON-schema tools with input validation; errors returned as observations so the agent can recover.
• Safety: refund/irreversible tools require human approval; all tools idempotent so a retry can't double-act; a critic pass (or neurosymbolic verifier) checks the proposed action against business rules before execution.
• Memory: conversation buffer → summary → semantic recall of similar past tickets.
• Observability: log every thought/action/observation with a trace ID; eval on a fixed ticket suite (pass@k, tool-call accuracy, cost/ticket).
• Tradeoff: more autonomy ⇒ more failure surface; constrain with verification + approvals. Multi-agent only if planner/executor/critic roles genuinely separate.

Design 5 — Train/fine-tune & align a domain model (P03, P05, P07, P10)

Requirement: adapt an open 7B model to a domain, align it to a house style, ship it quantized.

• Pipeline: SFT on curated instructions (LoRA) → preference data → DPO (closed-form, stable) → quantize (int4-AWQ) → eval gate (domain benchmark + safety + calibration) → registry → canary deploy.
• Compute: if full fine-tune is needed, ZeRO/FSDP across the data-parallel group; otherwise QLoRA on a single GPU. Track every run in W&B/MLflow with the exact data + config.
• Tradeoff: DPO is simpler/cheaper than PPO-RLHF but less online exploration; LoRA is cheap but caps how far you can move the model — measure the gap.

Design 6 — Edge / on-device deployment (P06, P17)

Requirement: run a small assistant on a device with no cloud round-trip, limited RAM, battery budget.

• Strategy: distill to a small student → int4 quantize → compile to the target runtime (ONNX Runtime / TensorRT / GGUF); KV-cache is the memory budget, so cap context.
• Tradeoff: every bit of quantization and every dropped parameter trades quality for battery/RAM/latency; measure on the device, not the dev box. Latency may be a safety property if it drives an actuator (P15).

Anti-patterns the panel listens for

• Quoting a model size without the KV-cache memory at target concurrency.
• "We'll fine-tune" when RAG or a better prompt solves it cheaper — or RAG when the knowledge is actually reasoning, not retrieval.
• Sampling (temperature>0) for a structured/extraction endpoint instead of constrained decoding.
• One eval number with no confidence interval; testing on the training distribution.
• An agent with no step cap, no idempotency, no human gate on irreversible actions.
• Claiming a latency/throughput number without saying whether it's TTFT or TPOT, prefill or decode, batched or single-stream.
• Multi-agent / multi-GPU / multi-region with no number justifying the added complexity.

Phase 00 — Foundations: Scaling Laws, FLOPs/Memory Math & Tradeoffs

The phase that builds the mental model and the arithmetic every later phase plugs into. Before you implement attention, fine-tune a model, or design a serving stack, you have to be able to answer — on a whiteboard, in two minutes — how much compute does this take, how much memory, how fast can it possibly go, what does it cost, and what am I trading? This is the round-1 filter of every senior AI interview and the first slide of every design review.

Why this phase exists

The difference between an engineer who "uses LLMs" and a Senior AI Engineer is that the senior treats a model as a physical system with a compute budget, a memory budget, and a bandwidth ceiling, not as a magic API. Almost every important decision in this curriculum — quantize or not, fine-tune or RAG, single GPU or sharded, bigger model or distilled smaller one, what batch size, what context length — is downstream of five numbers you can compute on paper:

1. FLOPs — 6ND to train, ~2N per token to run. Compute is the training bill and part of the inference bill.
2. Weight memory — N × bytes. The thing everyone quotes.
3. KV-cache memory — the thing everyone forgets, and the thing that actually OOMs your server at scale.
4. The roofline — whether you are compute-bound or memory-bandwidth-bound, which decides your entire optimization strategy. (Spoiler: decode is bandwidth-bound, and that one fact explains PagedAttention, batching, and quantization.)
5. Cost & the tradeoff — $/1M tokens, and a weighted resolver that proves "best model" is a category error — there is only "best for these priorities."

Get fluent here and the rest of the track is filling in mechanisms behind numbers you already respect.

Concept map

                        ┌─────────────────────────────────────┐
                        │   A model is a physical system       │
                        └─────────────────────────────────────┘
                          │            │             │
              ┌───────────┘     ┌──────┘      ┌──────┘
              ▼                  ▼             ▼
        COMPUTE (FLOPs)    MEMORY (bytes)   BANDWIDTH (bytes/s)
        6ND train          weights N·b       roofline ridge
        2N / token         KV-cache 2·L·T·d  decode is BW-bound
              │                  │             │
              └────────┬─────────┴──────┬──────┘
                       ▼                ▼
                  $ / 1M tokens    Chinchilla (D≈20N)
                       │                │
                       └───────┬────────┘
                               ▼
                    TRADEOFF RESOLVER
            quality · latency · cost · memory · safety
        (same two designs flip winners when weights change)

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py.

Integrated scenario ideas

• Size a serving deployment: given a 13B model, 1k concurrent users, 8k context — compute weights, KV-cache at peak, and the GPU count. Show that the KV-cache, not the weights, sets the number.
• Defend "distill, don't scale": compute the perpetual inference bill of a 70B vs a distilled 13B at your QPS; show the breakeven volume.
• Pick compute-optimal: given a fixed FLOP budget, use Chinchilla to choose N and D; show why the bigger-but-under-trained option is both worse and more expensive to serve.
• Explain the roofline to a new hire: why throwing more FLOPs at decode does nothing, and why batching + quantization do.

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive 6ND and 2N/token from first principles (MACs × passes).
• You can compute a KV-cache size in your head for a given (L, T, d_model) and explain why GQA shrinks it.
• You can state whether prefill and decode are compute- or memory-bound, and why.
• You can turn any "deploy model X" prompt into FLOPs + memory + $/1M + a named tradeoff.

Key takeaways

• Training is a one-time 6ND; inference is 2N per token, forever. At scale the bill, the latency, and the carbon live on the inference side — which is why this curriculum spends whole phases on quantization, serving, and the KV-cache.
• The KV-cache is the inference memory wall. Weights are fixed; the cache grows with batch × context and is what actually OOMs you. GQA/MQA and quantized KV exist to fight it.
• Decode is memory-bandwidth bound. You re-read every weight per token, so the lever is bytes moved (quantize, batch to amortize), not more FLOPs. This single fact explains half of inference engineering.
• "Best model" is meaningless; "best for these priorities" is engineering. The tradeoff resolver makes that concrete — and the deciding dial is the one sentence you put in the ADR.

Next: Phase 01 — Tokenization & the Multimodal Input Pipeline.

Warmup Guide — Foundations: Scaling Laws, FLOPs/Memory Math & Tradeoffs

Zero-to-senior primer for Phase 00. We start from "what is a neural network actually doing in terms of arithmetic and memory" and end with the five numbers and the one artifact that turn architectural opinion into engineering: the compute budget (6ND / 2N), the memory budget (weights + the KV-cache), the bandwidth ceiling (the roofline, and why decode is memory-bound), the scaling laws (Chinchilla), the serving cost ($/1M tokens), and the weighted tradeoff decision. Every later phase plugs a real mechanism into a slot in the model built here.

Table of Contents

• Chapter 1: A Model Is a Physical System
• Chapter 2: FLOPs — The Compute Budget (6ND and 2N)
• Chapter 3: Memory I — The Weights
• Chapter 4: Memory II — The KV-Cache, the Real Inference Wall
• Chapter 5: The Roofline — Compute-Bound vs Memory-Bandwidth-Bound
• Chapter 6: Scaling Laws & Chinchilla — How Big, How Much Data
• Chapter 7: Throughput & Cost — $/1M Tokens
• Chapter 8: The Tradeoff — Quantifying "Best for These Priorities"
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: A Model Is a Physical System

From zero. A large language model, stripped of all the framing, is a function: it takes a sequence of token IDs and returns a probability distribution over the next token. That function is implemented as a long chain of matrix multiplications interleaved with a few cheap element-wise operations (normalization, activation, residual addition). That is almost the entire computational story — a transformer is, to first order, a stack of matmuls.

Why this framing matters. Once you see the model as "a pile of matmuls over some weights," three physical resources fall out immediately, and every engineering decision you will ever make is about one of them:

The senior's habit. When someone says "let's deploy a 34B model," the senior does not nod. They silently compute: 34e9 × 2 = 68 GB of weights (won't fit one 80 GB GPU once you add the KV-cache and activations), decode is bandwidth-bound so single-stream throughput is roughly HBM_bandwidth / 68 GB ≈ 30 tok/s (too slow without batching), and the $/1M tokens at the target QPS is X. Then they have an actual opinion. This chapter's goal is to make that reflex automatic.

Common misconception. "Bigger model = better, end of story." Bigger raises quality with diminishing returns (Chapter 6) and raises every inference cost linearly forever. The senior move is almost always to find the smallest model that clears the quality bar, then quantize and distill it — because you pay the inference cost on every single request for the life of the product.

Chapter 2: FLOPs — The Compute Budget (6ND and 2N)

What a FLOP is. A floating-point operation: one multiply or one add. A matrix multiply of an m×k by a k×n matrix does m·n·k multiplies and m·n·k adds — about 2·m·n·k FLOPs. The "2" (one multiply + one add per element, a multiply-accumulate or MAC) shows up everywhere.

The forward pass: ~2N per token. Every parameter in the model participates in roughly one MAC per token that flows through it. With N parameters and the factor-of-2 for the MAC, a forward pass over one token costs about

$$ \text{FLOPs}_{\text{fwd}} \approx 2N \quad \text{per token.} $$

This is astonishingly clean and it is the single most useful number in the field. A 7B model costs ~14 GFLOP to advance one token. Generating 500 tokens costs ~7 TFLOP. Memorize 2N.

Training: ~6ND. Training does a forward pass and a backward pass. The backward pass computes gradients with respect to both the inputs and the weights, which is about twice the work of the forward pass. So one training step over a token costs ~2N (fwd) + ~4N (bwd) = ~6N. Over a whole training run of D total tokens:

$$ \boxed{;\text{FLOPs}_{\text{train}} \approx 6 \cdot N \cdot D;} $$

Why you must know this cold. It is the back-of-envelope that sizes a training run, a cluster, and a budget. A 7B model on 1T tokens: 6 × 7e9 × 1e12 = 4.2e22 FLOPs. On a cluster doing 1e18 FLOP/s (an exaflop) at 50% utilization, that's 4.2e22 / 5e17 ≈ 84,000 seconds ≈ a day — and you can now reason about whether your idea is a weekend or a quarter.

Prefill vs decode (same FLOPs, different physics). Processing a 1000-token prompt ("prefill") costs ~2N × 1000 FLOPs, but you do it in one parallel pass over all positions — it's a big, efficient matmul (compute-bound). Generating the next 1000 tokens ("decode") costs the same ~2N × 1000 FLOPs, but you must do it one token at a time, each step a skinny matrix-times-vector that barely uses the GPU's math units (memory-bound — Chapter 5). Same arithmetic, completely different engineering. This distinction drives all of P09.

Common misconception. "Attention dominates the FLOPs." For typical sequence lengths the feed-forward (MLP) layers dominate — they hold ~2/3 of the parameters. Attention's O(T²) cost only overtakes the MLP at long context (the reason FlashAttention and long-context tricks exist). For the 2N estimate, attention's quadratic term is a correction, not the main term, until T gets large.

Chapter 3: Memory I — The Weights

What it is. The weights are N numbers, each stored in some precision. Memory for weights is simply

$$ \text{bytes}_{\text{weights}} = N \times \text{bytes_per_param}. $$

Why precision is a first-class lever. A 70B model in fp16 (140 GB) needs two 80 GB GPUs just for the weights; in int4 (~35 GB) it fits on one with room to spare. Quantization is not a micro-optimization — it changes how many GPUs you need, which changes the cost and the architecture. That is why P06 is a whole phase.

Training memory is much larger than the weights. During training you also hold:

• Gradients: one per parameter (N values).
• Optimizer state: Adam keeps a 1st- and 2nd-moment estimate — 2N values — usually in fp32, plus an fp32 master copy of the weights.
• Activations: the forward-pass intermediates saved for the backward pass (scales with batch and sequence length; activation checkpointing trades compute to shrink this).

A rough fp16-training rule: ~16 bytes/param (2 weights + 2 grad + 4+4 Adam moments + 4 master), so training a 7B model needs ~112 GB before activations. This 2N optimizer-state fact is exactly why ZeRO/FSDP shard it across GPUs (P10).

Common misconception. "We quantized to int4 so training will fit." Quantization is an inference/serving technique for the frozen weights; training still needs full-precision gradients and optimizer state (QLoRA, P05, is the clever exception — it keeps the base frozen and quantized and only trains small adapters in higher precision).

Chapter 4: Memory II — The KV-Cache, the Real Inference Wall

The problem it solves. Autoregressive decoding generates token t+1 by attending over tokens 1..t. Naively, every new token recomputes the keys and values for all previous tokens — O(T²) work to produce a sequence. The KV-cache stores the key and value vectors of every past token so each new step only computes K/V for the one new token and reads the rest from cache. It turns decode from O(T²) into O(T).

The cost. For every layer and every past token you store a key vector and a value vector, each of width d_model (for full multi-head attention):

$$ \text{bytes}{\text{KV}} = 2 \times n{\text{layers}} \times T \times d_{\text{model}} \times \text{bytes_per_elem} \times \text{batch}. $$

The leading 2 is "K and V." Worked example — Llama-2-7B (32 layers, d_model 4096), 4096 context, fp16, one sequence:

$$ 2 \times 32 \times 4096 \times 4096 \times 2 = 2.15\ \text{GB}. $$

Per sequence. Now serve 50 concurrent users and that's ~107 GB of KV-cache — more than the 14 GB of weights, and the actual reason you OOM. At scale, you provision GPUs for the KV-cache, not the weights. Internalize this; it is the most common thing junior candidates miss in a serving design.

The fixes (previewed; built in P02/P09):

• GQA / MQA — share K/V across query heads so the stored width is d_model · n_kv/n_head instead of d_model. With 8 KV heads out of 32, the cache is 4× smaller (0.54 GB above). This is why every modern model uses GQA.
• PagedAttention (P09) — store the cache in fixed blocks like OS pages so you don't waste memory on fragmentation and over-allocation.
• KV quantization — store the cache in int8/int4.
• Sliding-window / shorter context — cap T.

Common misconception. "Long context is just an attention-cost problem." The O(T²) attention compute matters, but the linear-in-T KV-cache memory is usually the binding constraint in production, and it grows with every concurrent request.

Chapter 5: The Roofline — Compute-Bound vs Memory-Bandwidth-Bound

The question. A GPU has two peak rates: how fast it can do math (peak FLOP/s) and how fast it can move data between HBM and the compute units (memory bandwidth, bytes/s). For any operation, which one is the bottleneck? The roofline model answers it with one number.

Arithmetic intensity. Define

$$ I = \frac{\text{FLOPs performed}}{\text{bytes moved}} \quad (\text{FLOP/byte}). $$

The roofline. The attainable performance is

$$ \text{attainable FLOP/s} = \min(\text{peak FLOP/s},; \text{bandwidth} \times I). $$

Plot it and you get a "roof": a slanted line (bandwidth × I) that rises until it hits a flat ceiling (peak FLOP/s). The corner — the ridge point — is at

$$ I_{\text{ridge}} = \frac{\text{peak FLOP/s}}{\text{bandwidth}}. $$

 FLOP/s
   peak ┤        ________________   ← compute-bound (flat roof)
        │       /
        │      /  slope = bandwidth
        │     /
        │    /          ridge point = peak/bandwidth
        └───┴──────────────────────► arithmetic intensity (FLOP/byte)
          memory-bound      compute-bound

• Left of the ridge (I < I_ridge): you finish the math before the data arrives — you are memory-bandwidth-bound. More FLOP/s does nothing; move fewer bytes.
• Right of the ridge (I > I_ridge): you are saturating the math units — compute-bound. Faster memory does nothing; do less math or use better units (tensor cores).

The punchline for LLMs. An A100 has ~312 TFLOP/s (bf16) and ~2 TB/s HBM, so I_ridge ≈ 156 FLOP/byte. Now consider decode: to produce one token you read every weight (~2 bytes) and do ~2 FLOPs per weight → arithmetic intensity ≈ 1 FLOP/byte. That is deeply to the left of 156 — autoregressive decode is hopelessly memory-bandwidth-bound. This is the most important single fact in inference engineering, and it explains:

• why batching helps (you read the weights once and apply them to many sequences, raising I),
• why quantization helps (fewer bytes per weight → more tokens per second),
• why PagedAttention/continuous batching exist (keep I high by keeping the batch full),
• why throwing a faster-FLOPs GPU at a decode workload is often money wasted.

Prefill and training, by contrast, are big dense matmuls with high I → compute-bound → they do benefit from more FLOP/s and tensor cores.

Common misconception. "Our GPU is at 100% utilization, so we're optimal." nvidia-smi utilization means "a kernel was running," not "the math units were busy." A memory-bound decode kernel shows 100% utilization while doing <5% of peak FLOPs. The roofline, not the utilization meter, tells you the truth (P17 goes deeper).

Chapter 6: Scaling Laws & Chinchilla — How Big, How Much Data

Scaling laws. Empirically, test loss falls as a power law in model size, dataset size, and compute, over many orders of magnitude:

$$ L(N) \approx L_\infty + \left(\frac{N_c}{N}\right)^{\alpha_N}. $$

In words: doubling the model gives a predictable, diminishing drop in loss. This is why "just make it bigger" worked for years — and why it eventually stops being worth it.

The compute-optimal question. You have a fixed compute budget C ≈ 6ND. Should you spend it on a bigger model (N) or more data (D)? The Chinchilla result (Hoffmann et al., 2022) found that for compute-optimal training, N and D should scale together, roughly:

$$ D \approx 20 \times N \quad (\text{~20 tokens per parameter}). $$

Why it mattered. Many earlier "large" models were massively under-trained — too big for the data they saw. Chinchilla (70B on 1.4T tokens) beat the 175B GPT-3 (300B tokens) at less compute. The senior lesson: a smaller model trained on more data is often better and much cheaper to serve (smaller N = less inference cost forever).

The inference caveat (post-Chinchilla practice). Chinchilla optimizes training compute. But if you'll serve the model billions of times, you should happily train a smaller model on more than 20×N tokens (past the compute-optimal point) because the extra training cost is paid once and the smaller N saves inference cost forever. This is why models like Llama are trained far past their Chinchilla point (e.g. 7B on ~2T tokens, ~285 tokens/param). Optimize for total cost of ownership, not training compute alone.

Common misconception. "Chinchilla says 20 tokens/param is optimal, so train to exactly there." 20×N is the training-compute-optimal floor; for a model you'll deploy at scale, over-training a small model is the right move. Know both the rule and when to break it.

Chapter 7: Throughput & Cost — $/1M Tokens

Decode throughput ceiling. From the roofline (Chapter 5), single-stream decode is bandwidth-bound: you re-read all the weights every token, so

$$ \text{tok/s}{\text{single}} \approx \frac{\text{bandwidth}}{\text{bytes}{\text{weights}}}. $$

A 7B fp16 model (14 GB) on 2 TB/s HBM: 2e12 / 1.4e10 ≈ 143 tok/s — a ceiling, real systems hit less. Batching reads the weights once for B sequences, so throughput scales ~linearly with batch (until you become compute-bound or run out of KV-cache room):

$$ \text{tok/s}{\text{batch}} \approx B \times \frac{\text{bandwidth}}{\text{bytes}{\text{weights}}}. $$

This is the whole economic argument for batched serving: at batch 32 our 7B model approaches ~4,500 tok/s on the same GPU.

Cost per million tokens. Given a GPU hourly price and a sustained throughput:

$$ \frac{$}{1\text{M tokens}} = \frac{\text{gpu_$/hour}}{\text{tok/s} \times 3600} \times 10^6. $$

At $2/hour and 4,500 tok/s, that's about $0.12 / 1M tokens — and now you can compare an open model on your own GPU against an API's published price with an apples-to-apples number, which is a real decision seniors make weekly.

MFU (Model FLOPs Utilization). achieved_FLOP/s ÷ peak_FLOP/s. Good training runs hit 40–55%; memory-bound decode is far lower (because it's not FLOP-limited at all). MFU is the honest efficiency metric — it tells you how much of the hardware you actually paid for is doing useful model work.

Common misconception. "The API is more expensive than self-hosting." Sometimes — but only once you account for batching (utilization), engineering time, idle GPUs at low traffic, and the quality gap. Compute the $/1M at your real utilization, not at theoretical peak, before you make the build-vs-buy call.

Chapter 8: The Tradeoff — Quantifying "Best for These Priorities"

The core idea. There is no "best model" or "best architecture." There is only best for a given set of priorities. A regulated medical assistant weights quality and safety far above cost; an internal batch summarizer weights cost and throughput above latency. The same two candidate designs can have opposite winners under those two weightings — and proving that to yourself is the entire point of the lab's tradeoff resolver.

The mechanism. Score each candidate 0..1 on each dial (here: quality, latency, cost, memory, safety, where 1 = best). Pick weights for the workload. The candidate score is the weight-normalized sum:

$$ \text{score}(c) = \sum_{d \in \text{dials}} \text{score}_c[d] \cdot \frac{w_d}{\sum_j w_j}. $$

The deciding dial is the one whose weighted (winner − loser) gap is largest — that single sentence is what you write in the decision record: "We chose the local int4-7B because the workload weights cost 4× and it wins decisively on cost, accepting a measured 0.25 quality gap."

Why "deciding dial" beats a raw score. A score of 0.71 vs 0.68 tells a reviewer nothing. "We traded ~0.25 of quality to win 4× on cost, and here's the eval showing the quality gap is acceptable for this use case" is an argument. Seniors win design reviews with the deciding dial and the number, not the aggregate.

The decision record (ADR-lite). For any model/serving decision, write three lines:

• Context — the workload facts and the weights (what we optimized for).
• Decision — the candidate chosen and the deciding dial.
• Consequence — the dial(s) we traded away, and the number that makes the trade acceptable.

Common misconception. "We'll just pick the model with the best benchmark score." Benchmarks measure quality on their distribution at unbounded cost/latency. Production is a constrained optimization; the resolver makes the constraints explicit so the decision is defensible six months later when someone asks "why didn't we use the bigger model?"

Lab Walkthrough Guidance

The lab (lab-01-transformer-cost-modeler) turns this chapter into code. Suggested order (matches the file):

1. FLOPs (training_flops, forward_flops_per_token, inference_flops) — Chapter 2. The only trick is validation; the formulas are 6ND and 2N.
2. Memory (weight_memory_bytes, kv_cache_bytes) — Chapters 3–4. The KV-cache GQA branch is the interesting one: scale d_kv by n_kv/n_head, and require both head counts or neither.
3. Scaling (chinchilla_optimal_tokens/params) — Chapter 6. A round-trip: tokens↔params at ratio 20.
4. Roofline (arithmetic_intensity, roofline_ridge_intensity, roofline_attainable_flops, is_memory_bound) — Chapter 5. The ridge is peak/bandwidth; attainable is the min.
5. Throughput & cost (decode_tokens_per_sec, cost_per_million_tokens, mfu) — Chapter 7. Note batch multiplies throughput.
6. Tradeoff (Candidate.__post_init__, resolve_tradeoff) — Chapter 8. Validate the dials, weight-normalize, find the deciding dial. The test_weights_flip_the_winner test is the soul of the lab.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked numbers.

Success Criteria

• All tests pass against your lab.py and against LAB_MODULE=solution.
• You can derive 6ND (MAC × 3 passes × N × D) and 2N/token without notes.
• You can compute the KV-cache size for (L=32, T=4096, d=4096, fp16) ≈ 2 GB in your head and explain why GQA with 8/32 KV heads makes it 0.54 GB.
• You can state the roofline ridge for an A100 (~156 FLOP/byte) and explain why decode (I≈1) is far left of it and therefore memory-bound.
• You can take any "deploy model X" prompt and produce, in two minutes on paper: training FLOPs, weight + KV-cache memory, a decode $/1M, and a named tradeoff with the deciding dial.

Interview Q&A

• "How many FLOPs to train a 13B model on 2T tokens?" — 6 × 13e9 × 2e12 = 1.56e23. Then reason about cluster-days at a given FLOP/s and MFU.
• "Serve a 70B model — how many 80 GB GPUs, and what limits it?" — fp16 weights = 140 GB → 2 GPUs minimum for weights alone; then add the KV-cache at your concurrency/context (often the binding constraint) and activations; int4 (~35 GB) changes the answer to one GPU. Lead with the KV-cache.
• "Is decoding compute- or memory-bound? Prove it." — memory-bound: I≈1 FLOP/byte vs an A100 ridge of ~156; you re-read all weights per token. Therefore batch and quantize, don't buy more FLOPs.
• "Bigger model or fine-tune a smaller one?" — inference cost is 2N/token forever; prefer the smallest model that clears the eval bar, then quantize/distill — unless measured quality and volume justify the perpetual bill. Cite Chinchilla and TCO, not training compute alone.
• "What did Chinchilla change and when do you ignore it?" — D≈20N is training-optimal; for a model you'll serve at scale, over-train a smaller model past that point because inference cost dominates TCO.
• "Why is nvidia-smi at 100% not the same as efficient?" — utilization = "a kernel ran," not "FLOPs were busy"; memory-bound decode shows 100% util at low MFU. Use the roofline.

References

• Kaplan et al., Scaling Laws for Neural Language Models (2020) — the original power laws.
• Hoffmann et al., Training Compute-Optimal Large Language Models ("Chinchilla", 2022).
• Williams, Waterman, Patterson, Roofline: An Insightful Visual Performance Model (2009).
• Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention (vLLM, 2023).
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models (2023).
• "Transformer Inference Arithmetic" (kipp.ly) and "LLM inference speed of light" (the cleanest write-ups of the 2N / KV-cache / bandwidth-bound math).
• NVIDIA A100/H100 datasheets — peak FLOP/s and HBM bandwidth numbers for your roofline.

Hitchhiker's Guide — Foundations: Scaling Laws & the Cost Math

The compressed practitioner tour. If WARMUP.md is the professor, this is the senior who leans over and says "here's what you actually need to remember."

The 30-second mental model

A model is a pile of matmuls over N weights. Three resources govern everything: compute (FLOPs), memory (bytes), bandwidth (bytes/s). Training costs 6ND once; inference costs 2N per token forever. The weights are the number everyone quotes; the KV-cache is the number that OOMs you. Decode is memory-bandwidth-bound, so the levers are batch and quantize, not more FLOPs. There's no "best model," only "best for these priorities."

The numbers to tattoo on your arm

Back-of-envelope one-liners

# "Can a 70B fit on one 80GB GPU?"
70e9 * 2 = 140 GB fp16  → no (need 2);  70e9 * 0.5 = 35 GB int4 → yes (+ KV headroom)

# "How fast can a 7B decode, single stream?"
2e12 (HBM B/s) / 14e9 (fp16 weights) ≈ 143 tok/s ceiling  → batch to go faster

# "KV-cache for 13B, 8k ctx, 40 layers, d=5120, 100 users, fp16?"
2 * 40 * 8192 * 5120 * 2 * 100 ≈ 670 GB  → KV-cache, not weights, sets the GPU count

# "Train 7B on 2T tokens — how big a job?"
6 * 7e9 * 2e12 = 8.4e22 FLOPs;  / (5e17 FLOP/s @ 50% MFU) ≈ 1.9 days

The framework one-liners (where these numbers live in real tools)

# Count params (the N in every formula)
sum(p.numel() for p in model.parameters())            # PyTorch
# Estimate FLOPs
from deepspeed.profiling.flops_profiler import FlopsProfiler
# Measured throughput / MFU
# vLLM logs tokens/s; nvidia-smi dmon for bandwidth; torch.profiler for kernels
# KV-cache config knobs that change the memory math:
#   num_key_value_heads (GQA), max_model_len (T), gpu_memory_utilization (vLLM)

War stories

• The KV-cache OOM at 3pm. Load tests passed at batch 8; production hit batch 60 at peak and OOM'd — because nobody multiplied the KV-cache by concurrency. The weights fit fine; the cache didn't. Capacity = weights + KV-cache at peak concurrency × context + activations. Always.
• The "100% GPU but slow" ticket. nvidia-smi showed 100%; latency was terrible. It was a memory-bound decode at ~4% MFU — the GPU was busy waiting on HBM, not computing. We fixed it with bigger batches (continuous batching) and KV quantization, not a faster GPU.
• The bigger-model regret. Team shipped a 70B because it scored 2 points higher on a benchmark. The inference bill was 5× and p99 latency missed SLO. A distilled 13B + RAG would have cleared the actual quality bar at a fifth of the cost. They optimized the benchmark, not the TCO.
• The under-trained giant. Pre-Chinchilla, a team scaled params and starved the model of data; a half-size model on 3× the tokens beat it at less compute and served cheaper.

Vocabulary (rapid-fire)

• FLOP — one multiply or add. MAC = multiply-accumulate ≈ 2 FLOPs.
• KV-cache — stored keys/values so decode is O(T) not O(T²); the inference memory wall.
• Roofline / ridge — min(peak, bw·I); ridge = peak/bw, the compute/memory crossover.
• Arithmetic intensity — FLOPs per byte moved.
• MFU — fraction of peak FLOPs actually used; honest efficiency metric.
• TTFT / TPOT — time to first token (prefill) / time per output token (decode).
• Chinchilla — compute-optimal D≈20N.
• TCO — total cost of ownership; inference cost × lifetime, the real objective.

Beginner mistakes

• Quoting weight memory and forgetting the KV-cache (and multiplying by concurrency).
• Trying to fix memory-bound decode with a faster-FLOPs GPU instead of batching/quantizing.
• Treating nvidia-smi utilization as efficiency.
• Picking the model by benchmark score with no cost/latency constraint.
• Confusing training memory (huge: grads + Adam) with inference memory (weights + KV).
• Forgetting that inference cost is paid per request, forever — the asymmetry that makes distillation/quantization/PEFT the default senior move.

The one thing to take away

Before you have an opinion about any model decision, compute FLOPs, weight memory, KV-cache, and $/1M — and name the deciding dial. If you can't, you're guessing; if you can, you're engineering.

Brother Talk — Foundations

Off the record. The stuff I'd tell you over a beer, not in a design review.

This phase is the one that makes people respect you, and almost nobody does it. Everyone can prompt a model. Very few engineers — even ones with "AI" in their title — can tell you, on the spot, how many GPUs it takes to serve a 34B model at 1k QPS and why. The moment you can do the FLOPs/memory/KV-cache/cost math out loud, you stop being "the person who calls the API" and become "the person we ask before we commit budget." That reputation is worth more than any framework you'll learn later.

The KV-cache is the thing that will save you at 2 a.m. I promise you that at some point a service will OOM in production and a roomful of smart people will be staring at it confused because "the model fits, we checked." You'll be the one who says "what's the concurrency × context? That's the KV-cache, and it's bigger than the weights." You'll fix it in ten minutes. Memorize the KV-cache formula like a phone number.

"Memory-bound decode" is a cheat code for sounding senior. Half the inference-optimization conversations in this industry are people reinventing the same insight: decode re-reads the weights every token, so it's bandwidth-bound, so batch and quantize. Once you really get the roofline, you'll watch teams waste months and money throwing FLOPs at a bandwidth problem, and you'll be able to redirect them in one sentence. That one sentence is a promotion.

Resist the bigger-model dopamine. There is a real, almost emotional pull toward "use the biggest, smartest model." It feels safe and impressive. But you pay for inference on every single request for the life of the product, and that cost compounds in a way training cost doesn't. The senior move — the one that's boring in the demo and heroic on the P&L — is the smallest model that clears the bar, quantized and distilled. Train yourself to feel good about shipping small.

Numbers end arguments; opinions start them. When you walk into a design review with "here's the FLOPs, here's the KV-cache at peak, here's the $/1M, and the deciding dial is cost," the conversation is over in your favor. When you walk in with "I think we should use the smaller model," you're in for an hour of vibes. Do the arithmetic before the meeting. Every time.

You don't need a PhD to do this, but you do need to not be scared of the math. The math in this phase is multiplication and one min(). That's it. The reason it's a differentiator isn't that it's hard — it's that most people never bother to make it a reflex. Be the person who bothered.

What's actually worth caring about: the four numbers (FLOPs, weights, KV-cache, $/1M) and the deciding dial. What's not worth losing sleep over: the exact constant in the scaling-law fit, or whether it's 6ND or 6.02ND. Estimates that are right to within 2× win design reviews; false-precision to 5 decimals just slows you down.

The career framing. This phase is the foundation of the "I own the model stack" identity. Every later phase — attention, quantization, serving, agents — is more impressive, but this is the one that makes the rest credible, because it proves you think about models as systems with budgets, not as magic. Get this cold and the whole track compounds. Now go make pytest green.

Lab 01 — Transformer FLOPs / Memory / Cost & Tradeoff Modeler

Phase: 00 — Foundations: Scaling Laws, FLOPs/Memory Math & Tradeoffs Difficulty: ⭐⭐☆☆☆ (the arithmetic is easy; the judgment is ⭐⭐⭐⭐⭐) Time: 3–4 hours

Round-1 of an AI engineering interview — and the first slide of every design review — is back-of-envelope arithmetic worn as architecture. This lab builds the calculator: training and inference FLOPs (6ND, 2N), weight memory across precisions, the KV-cache (with GQA), Chinchilla-optimal data, the roofline (compute- vs memory-bound), the memory-bandwidth ceiling on decode throughput, the serving cost per million tokens, and a weighted tradeoff resolver whose entire point is that the same two designs flip winners when the workload's priorities change — which is why "best model" is a meaningless phrase.

What you build

• training_flops / forward_flops_per_token / inference_flops — the compute budget: 6·N·D to train, 2·N per token to run. The numbers that size a cluster and a bill.
• weight_memory_bytes — N × bytes_per_param across fp32/fp16/int8/int4 — the lever that decides how many GPUs you need.
• kv_cache_bytes — 2 · L · T · d_model · bytes · batch, with the GQA branch (d_kv = d_model · n_kv/n_head) — the memory wall everyone forgets and the thing that actually OOMs a server.
• chinchilla_optimal_tokens / chinchilla_optimal_params — the compute-optimal D≈20N allocation (and its inverse).
• arithmetic_intensity / roofline_ridge_intensity / roofline_attainable_flops / is_memory_bound — the roofline: min(peak, bw·I), ridge = peak/bw, and the check that proves decode is memory-bandwidth-bound.
• decode_tokens_per_sec / cost_per_million_tokens / mfu — throughput ceiling (and how batching amortizes the weight read), $/1M tokens, and the honest efficiency metric.
• Candidate (five dials) + resolve_tradeoff — score two designs under the workload's weights and return the winner and the deciding dial (the one sentence for your decision record).

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why training is 6ND and inference is 2N/token (MACs × passes).
• You can compute the KV-cache for (L=32, T=4096, d=4096, fp16) ≈ 2 GB in your head and why GQA 8/32 makes it 0.54 GB.
• You can explain why test_decode_tokens_per_sec_and_batching shows ~linear scaling with batch, and why that is the entire economic argument for batched serving.
• You can explain why is_memory_bound(1.0, …) is True for decode and why that one fact governs inference optimization.
• You can explain why test_weights_flip_the_winner is the whole point of the lab.
• Given any "deploy model X" prompt you can, in two minutes, write FLOPs + weight/KV memory + a $/1M estimate + the deciding dial on paper.

How this maps to the real stack

Limits of the miniature (be honest in the interview): the FLOP rules ignore attention's O(T²) term (matters at long context), embeddings, and the LM head; real KV-cache adds block fragmentation and over-allocation (which is why PagedAttention exists); decode throughput is a ceiling — real systems lose to kernel launch overhead, sampling, and Python; and a tradeoff score is a structured conversation, not a number you can sum mechanically across every workload.

Extensions (build these on real hardware)

• Add the attention O(T²) FLOPs term and find the context length where it overtakes the MLP.
• Add training_memory_bytes(N, optimizer="adamw", dtype="fp16", activation_checkpointing=bool) using the ~16 bytes/param rule, and show why a 7B model needs ZeRO to train on one GPU.
• Add paged_kv_waste(seq_lens, block_size) modeling fragmentation, and show how PagedAttention's block allocator reclaims it (preview of P09).
• Plot a real roofline for your GPU (peak_flops, bandwidth from the datasheet) and mark where prefill, decode, and a big GEMM land on it.
• Pull params with sum(p.numel() for p in model.parameters()) on a real HF model and check your weight_memory_bytes against model.get_memory_footprint().

Interview / resume

• Talking points: "How many GPUs to serve a 70B at 1k QPS, and what limits it?" "Is decode compute- or memory-bound — prove it." "Bigger model or distill a smaller one — defend with numbers." "Compose the $/1M for self-hosting vs an API."
• Resume bullet: Built a transformer cost-and-tradeoff model that computes training/inference FLOPs, weight and KV-cache memory (with GQA), the compute-vs-bandwidth roofline, decode throughput and $/1M-token serving cost, and produces explainable, priority-weighted model-selection decisions for design reviews.

Phase 01 — Tokenization & the Multimodal Input Pipeline

The phase that builds the boundary between the human world and the tensor world. A model never sees your text, your image, or your audio — it sees a sequence of integer token IDs, and it emits integer IDs back. Everything upstream of the first matmul is tokenization, and it is the most under-respected, highest-leverage component in the stack: get it wrong and your round-trips break, your chat quality silently collapses, your multilingual users pay 3× the cost, and a "harmless" string of text reserves a control token it was never supposed to. Phase 00 taught you the model is a physical system; this phase teaches you what you actually feed it.

Why this phase exists

Every later phase — attention, fine-tuning, RAG, serving, agents — assumes a sequence of token IDs already exists. Where do they come from, and why are they the right ones? The answer is a learned subword vocabulary (BPE), a byte-level base alphabet that makes the codec closed (no <unk>, a true round-trip), a set of reserved special/control tokens the model attends to for structure, and a chat template that must match the training format byte-for-byte. Get these wrong and the failures are silent — no exception, no crash, just degraded quality that takes a week to diagnose. That makes tokenization a senior-level topic disguised as a beginner one.

The five things you will own after this phase:

1. Why subword. Char-level makes sequences too long (and 2N/token cost too high); word-level makes the vocab huge and chokes on out-of-vocabulary words with <unk>. BPE is the data-driven compromise: learn a merge table from a corpus.
2. Byte-level + byte fallback. Start the alphabet from the 256 byte values so every possible input — any language, any emoji, any control byte — decomposes into base tokens. The codec is closed: no <unk> for normal text, and decode(encode(x)) == x exactly.
3. Special/control tokens. <bos>, <eos>, <pad>, and the ChatML <|im_start|> / <|im_end|> are reserved IDs added structurally — never produced by tokenizing user text that happens to contain those strings. Mishandling them is the #1 chat bug.
4. The chat template. Models are trained on an exact string format. Serving must reproduce it byte-for-byte; a stray space or missing newline degrades quality with no error.
5. The bridge to multimodal. An image is just more tokens: cut it into patches, embed each patch, and splice the patch embeddings into the same sequence behind an <image> placeholder. Tokens are the universal interface (full mechanism in Phase 04).

Concept map

                ┌──────────────────────────────────────────────┐
                │  text / image / audio  →  INTEGER TOKEN IDs   │
                └──────────────────────────────────────────────┘
                                     │
          ┌──────────────────────────┼──────────────────────────┐
          ▼                          ▼                          ▼
   GRANULARITY               BASE ALPHABET               STRUCTURE
   char  → seqs too long     bytes (256)  ──────────►    special tokens
   word  → huge vocab, <unk> closed vocab               <bos> <eos> <pad>
   SUBWORD (BPE) ◄───────────no <unk>, round-trips      <|im_start|>/<|im_end|>
          │                          │                          │
          ▼                          ▼                          ▼
   merge table (learned)      byte fallback              CHAT TEMPLATE (ChatML)
   freq-greedy + tie-break    decode(encode(x))==x       byte-for-byte = training
          │                          │                          │
          └──────────────┬───────────┴──────────────┬───────────┘
                         ▼                          ▼
                  VOCAB-SIZE TRADEOFF        PITFALLS
                  compression ↑ vs           numbers/arithmetic, whitespace,
                  embedding table ↑          multilingual fairness, glitch tokens
                         │
                         ▼
                  MULTIMODAL: image patches as tokens  (→ Phase 04)

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py. The soul test is the round-trip: if decode(encode(x)) ever differs from x, the codec is broken — and you will be able to say exactly which invariant failed.

Integrated scenario ideas

• Diagnose a "the model got dumber" report. A chat app's quality dropped after a tokenizer upgrade. Show that the served chat template no longer matches the training format (an extra space after the role), and that this is a byte-for-byte mismatch the model never saw — no exception, just degraded completions.
• Defend a vocab size in a design review. Given a target language mix and an embedding budget, argue 32k vs 128k: bigger vocab compresses sequences (fewer tokens → cheaper 2N/token and longer effective context) but inflates the embedding + LM-head matrix (vocab × d_model). Name the deciding dial (carry the Phase 00 tradeoff habit forward).
• Explain a multilingual cost complaint. A user writing in a low-resource language is billed 3× for the "same" text. Show the fertility gap (tokens per word) and tie it to who the BPE corpus was trained on — a fairness and a cost problem.
• Hunt a glitch token. Reproduce the SolidGoldMagikarp failure mode: a token that exists in the vocab (it appeared in the BPE corpus) but was rarely or never seen in training, so its embedding is essentially random and the model behaves bizarrely when it appears.
• Sketch the multimodal splice. Reserve an <image> token, expand it into k patch slots, and describe how patch embeddings replace those slots in the sequence — the same interface, more tokens (preview of Phase 04).

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can explain char vs word vs subword tradeoffs and why BPE won.
• You can run one BPE merge step by hand: count pairs, pick the most frequent, apply the lexicographic tie-break, rewrite the words.
• You can explain why byte-level BPE has no <unk> and why decode(encode(x)) == x for any input including emoji and CJK.
• You can explain why encoding the literal text <bos> must not yield the reserved bos_id.
• You can state, byte-for-byte, what a ChatML turn looks like and what breaks on a mismatch.
• You can articulate the vocab-size tradeoff (compression vs embedding-table size) with the deciding dial.
• You can name at least three tokenization pitfalls (numbers/arithmetic, multilingual fairness, glitch tokens) and their root cause.

Key takeaways

• The model never sees text — it sees integer IDs. The tokenizer is the entire interface, and every silent quality bug at the edges of an LLM app starts here. Respect it accordingly.
• Byte-level BPE makes the codec closed. Starting from the 256 byte values means every possible input decomposes into base tokens, so there is no <unk> for normal text and the round-trip is exact. This one design choice removes a whole class of data-loss bugs.
• Determinism comes from the tie-break, not from luck. When two pairs tie on frequency, the lexicographically smallest wins — so the merge table is identical regardless of corpus ordering or dict iteration. Determinism is a feature you engineer, not one you hope for.
• Special tokens are structural, not textual. They are reserved IDs added by the framework, and the literal angle-bracket strings in user text must never collide with them — that boundary is both a correctness and a security property.
• The chat template must match training byte-for-byte. This is the most common and most expensive silent failure in chat systems: no error, just worse answers. Always serve the exact format the model was trained on.
• Tokens are the universal interface. Text subwords, image patches, audio frames — all become entries in the same sequence. Once you internalize "everything is tokens," multimodal stops being mysterious (Phase 04 fills in the mechanism).

Next: Phase 02 — The Transformer From Scratch.

Warmup Guide — Tokenization & the Multimodal Input Pipeline

Zero-to-senior primer for Phase 01. The model is a function from a sequence of integer IDs to a probability distribution over the next ID. Everything between the human-readable input and that first integer is tokenization — and it is where the most expensive, most silent bugs in the entire LLM stack live. We start from "why can't the model just read characters?" and end at "an image is just more tokens." Along the way: the BPE training algorithm step by step, byte-level encoding and why it kills <unk>, the special/control tokens that make or break chat, the chat template you must reproduce byte-for-byte, the vocabulary-size tradeoff, the famous pitfalls (arithmetic, whitespace, multilingual fairness, glitch tokens), and the bridge to multimodal input.

Table of Contents

• Chapter 1: The Model Speaks Integers, Not Text
• Chapter 2: Char vs Word vs Subword — Why Tokenization Exists
• Chapter 3: The BPE Training Algorithm, Step by Step
• Chapter 4: Merge Priority & the Deterministic Tie-Break
• Chapter 5: Byte-Level BPE & the Byte Fallback (No <unk>)
• Chapter 6: Special & Control Tokens — the #1 Chat Bug
• Chapter 7: Chat Templates (ChatML) & Byte-for-Byte Matching
• Chapter 8: Vocabulary-Size Tradeoffs — Compression vs the Embedding Table
• Chapter 9: Tokenization Pitfalls — Numbers, Whitespace, Fairness, Glitch Tokens
• Chapter 10: The Bridge to Multimodal — Image Patches as Tokens
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: The Model Speaks Integers, Not Text

From zero. A transformer's first layer is an embedding lookup: a table with one row per vocabulary entry, where row i is a learned d_model-dimensional vector. To use that table, the input must already be a list of integer IDs, one per position. The model has no notion of "letters" or "words" — it has a fixed set of V symbols numbered 0..V-1, and it manipulates the vectors those numbers point to. The entire job of the tokenizer is the two functions that bracket the model:

   "the cat"  ──encode──►  [1037, 4937]  ──►  [embedding lookup] ──► transformer ──► logits
                                                                                        │
   "the cat"  ◄──decode──  [1037, 4937]  ◄──────────────── argmax / sampling ◄──────────┘

Why this framing matters. Because the model only ever sees IDs, three properties of the tokenizer become correctness-critical, and none of them are about the model:

• It must be reversible. decode(encode(x)) should return x exactly, or you have silently corrupted the user's input before the model ever ran.
• It must be closed. Every possible input must map to some sequence of IDs — there can be no "I don't have a symbol for this" hole that loses data.
• It must match training. The model learned statistics over a specific tokenization. Feed it a different one at serving time and you are off-distribution, with no error to tell you so.

Production significance. Every "the model got dumber overnight" incident, every "it can't count the letters in a word" complaint, every "non-English users are billed triple" ticket traces back to a tokenizer property, not a model weight. The tokenizer is the cheapest component to get wrong and the most expensive to debug, because it fails silently.

Common misconception. "Tokenization is preprocessing — boring plumbing I can ignore." It is the plumbing, and the plumbing is where the leaks are. Senior engineers treat the tokenizer as a first-class part of the model contract, versioned and tested like one.

Chapter 2: Char vs Word vs Subword — Why Tokenization Exists

The design question. We need to chop text into a sequence of symbols drawn from a fixed vocabulary. What should the symbols be? There are three natural choices, and tokenization exists because the obvious two are both bad.

Option A — characters (or bytes). Vocabulary is tiny (~100 characters, or exactly 256 bytes). Closed by construction (every string is a sequence of characters). But sequences are long: "internationalization" is 20 character-tokens. Since inference cost is ~2N FLOPs per token (Phase 00) and attention is O(T²), long sequences are directly expensive, and the model must learn to compose meaning from tiny pieces over a long range — harder to learn.

Option B — words. Sequences are short (one token per word) and each token is meaning-rich. But the vocabulary is enormous and open: natural language has a long tail of rare words, names, typos, and morphology, so you either keep a million-entry vocab (a giant embedding table) or you cap it and emit <unk> for everything outside — which is lossy (you can't decode <unk> back) and brittle (the model is blind to exactly the rare, information-dense words you most need it to read).

Option C — subwords (the winner). Keep whole common words as single tokens, but break rare words into reusable sub-word pieces. "tokenization" might be token + ization; an unseen name breaks into smaller fragments that still exist in the vocab. You get short-ish sequences and a bounded vocabulary and, with a byte-level base alphabet (Chapter 5), no <unk> at all. The data decides where the boundaries are, via a learned algorithm — Byte-Pair Encoding (BPE).

How it works under the hood (the intuition). BPE starts from the smallest pieces (bytes) and greedily merges the most frequent adjacent pair over and over. Frequent sequences — common words, common suffixes — get merged into single tokens because they show up a lot; rare sequences stay broken into smaller pieces. The vocabulary self-organizes around the statistics of the corpus.

Production significance. Compression (characters per token) directly sets your effective context length and your per-request cost. A good subword vocab gets English to ~4 characters per token (≈ 0.75 words/token) — that's the number to remember.

Common misconception. "Tokens are words." They are frequency-driven fragments. " the" (with a leading space) is often one token; "tokenization" might be two or three; a rare name might be five. The model reasons over these fragments, which is exactly why it struggles to count letters or do digit arithmetic (Chapter 9).

Chapter 3: The BPE Training Algorithm, Step by Step

What BPE training produces. Two artifacts: an ordered list of merges (each a pair of symbols that fuse into one) and a vocabulary (every token string → integer ID). The ordered merge list is the heart of it — its order is the encoding priority (Chapter 4).

The algorithm. Given a corpus and a target number of merges:

1. Pre-tokenize the corpus into "words." (Our lab splits on \S+|\s+, so whitespace runs are their own words — this is what makes spaces survive the round-trip; Chapter 9. Real GPT-2/GPT-4 use a richer regex.) Represent each word as a tuple of its base symbols (in our lab, its UTF-8 bytes). Count how many times each distinct word-form occurs.
2. Repeat num_merges times: a. Count every adjacent symbol pair across all words, weighted by word frequency. b. Pick the most frequent pair (tie-break: lexicographically smallest — Chapter 4). c. Append it to the merge list and rewrite every word, fusing that pair into one symbol. d. Stop early if there are no pairs left.
3. Assemble the vocab: reserved specials first (stable IDs), then the base alphabet, then each merged token in merge order.

A fully worked example. Take the toy corpus and pretend the base symbols are the characters themselves (the lab uses bytes, but the counting is identical):

corpus words (with counts):   low ×5   lower ×2   newest ×6   widest ×3

Represent each as characters and count adjacent pairs across all words, weighted by the word counts. The pair e+s appears in newest (×6) and widest (×3) → count 9. The pair s+t appears in the same two words → count 9 too. (Here e+s wins on the tie-break since it's lexicographically smaller.) So:

merge 1:  e + s  ->  es        # newest -> n e w es t,  widest -> w i d es t

Recount. Now es+t appears in both → count 9, the new winner:

merge 2:  es + t -> est        # newest -> n e w est,   widest -> w i d est

Recount. Now l+o appears in low (×5) and lower (×2) → count 7:

merge 3:  l + o  -> lo
merge 4:  lo + w -> low        # "low" is now a single token; "lower" -> low e r

The merge count is the literal answer to "how many times did we run the loop" — four merges produced the tokens es, est, lo, low. After k merges the vocabulary is (base alphabet) + k (minus any key collisions). This is exactly what train_bpe(corpus, num_merges) does in the lab; its first merge on the lab corpus is ('a','t') (frequency 8: cat×3, sat×2, mat, rat, ate), which beats ('t','h') at frequency 6 (the×6) — count the actual most-frequent pair, not the eye-catching one.

Why frequency-greedy. The original Sennrich et al. (2016) insight was that you don't need a clever objective: greedily merging the most common pair already discovers morphologically sensible units (stems, suffixes, common words) because those are the frequent sequences. It's simple, fast, and good enough — which is why it's everywhere.

Production significance. This is the exact loop inside HuggingFace tokenizers BpeTrainer and SentencePiece's BPE mode. The only production differences are scale (millions of merges, parallelized counting), a richer pre-tokenizer, and stopping at a target vocab size rather than a fixed merge count (Chapter 8).

Common misconception. "BPE finds the linguistically correct morphemes." It finds frequency units. They often look morphological (because frequent things often are), but BPE will happily merge " the" or "ization" or an artifact of your corpus's quirks — it has no grammar, only counts.

Chapter 4: Merge Priority & the Deterministic Tie-Break

The problem. Two facts make BPE subtle. First, the merge list isn't just a set — it is an ordered list, and that order is reused at encoding time. Second, ties are common (many pairs share the same frequency), and if you resolve them by "whatever the dict happened to yield first," your tokenizer becomes non-deterministic across Python versions, machines, and corpus orderings — a nightmare, because two runs produce different vocabularies and the IDs no longer mean the same thing.

How merge order becomes encoding priority. When you encode a new word, you don't recount frequencies — you apply the learned merges in rank order. The merge learned first has rank 0 (highest priority); to tokenize a word you repeatedly find the present adjacent pair with the lowest rank and fuse it, until no learned pair remains. In the lab this is merge_ranks (a pair -> rank dict) and the _bpe method:

candidate = min(
    (p for p in pairs if p in self.merge_ranks),  # only learned pairs
    key=lambda p: self.merge_ranks[p],            # lowest rank = applied first
    default=None,
)

This is why the order matters: applying es+t before e+s would give a different segmentation. Encoding must replay training's priority exactly.

The deterministic tie-break. During training, when two pairs tie on frequency, pick the lexicographically smallest pair. This single rule makes the whole algorithm a pure function of (corpus contents, num_merges) — independent of dict iteration order or the order lines appear in the corpus. The lab proves this with test_training_is_order_independent_of_dict_iteration: reversing the corpus lines yields byte-identical merges. (The lab implements "smallest wins" by inverting the key inside a max; the effect is the same.)

Why tattoo this on your arm. A tokenizer is a contract. Its IDs are baked into a trained model. If the same corpus can produce two different vocabularies, you cannot reproduce a model, cannot share a tokenizer across a team, and cannot debug a mismatch. Determinism is engineered, via the tie-break — never assumed.

Production significance. Real trainers (HF, SentencePiece) pin determinism the same way (a defined tie-break, fixed iteration). It's why a merges.txt is a versioned, frozen artifact shipped with the model.

Common misconception. "Ties are rare, so the tie-break doesn't matter." On real corpora, ties are everywhere (huge numbers of pairs share low frequencies), and an unstable tie-break is exactly the kind of bug that passes your tests on one machine and corrupts a model on another.

Chapter 5: Byte-Level BPE & the Byte Fallback (No <unk>)

The problem it solves. If your base alphabet is "characters seen in the training corpus," then any unseen character — an emoji, a CJK glyph, a rare diacritic, a raw control byte — has no base token and must become <unk>. That is lossy (you can't decode it back) and the model is blind to it. The fix is to make the base alphabet the 256 byte values themselves. Since every string, in any language, is ultimately a sequence of UTF-8 bytes, every possible input decomposes into base byte-tokens. There is nothing left to be out-of-vocabulary.

How it works under the hood. Encode the text to UTF-8 bytes; each byte is a base symbol; BPE merges sit on top of the byte alphabet exactly as before. A 4-byte emoji with no learned merges simply becomes its 4 byte-tokens — test_byte_fallback_for_unseen_unicode asserts exactly len(ids) == len("🦩".encode("utf-8")). Decoding maps each byte-token back to its byte and UTF-8-decodes the byte stream. The codec is closed, and decode(encode(x)) == x holds for any round-trippable input — the "soul test."

The GPT-2 byte↔unicode trick (why the lab has BYTE_ENCODER). Some byte values are control or whitespace bytes that are ugly or ambiguous as dictionary keys. GPT-2 (and the lab) map each of the 256 bytes to a distinct printable unicode character via a fixed bijection, so every base token is a clean single-character string with no raw control bytes inside vocab keys. It is perfectly invertible (BYTE_DECODER is the inverse), so it changes nothing semantically — it just makes the implementation tidy. This is why you'll see funny characters like Ġ (a space) in a GPT-2 vocab dump.

Production significance. GPT-2, GPT-3, GPT-4 (cl100k_base/o200k_base via tiktoken), Llama, and friends are all byte-level. The practical payoff: no <unk> token in normal operation, exact round-trips, and graceful handling of any language or binary-ish input — a whole category of data-loss bugs simply cannot occur.

Common misconception. "Byte-level means it tokenizes every character as a byte, so it's slow and long." No — the base is bytes, but the merges fuse common byte sequences back into whole-word tokens. Byte-level BPE on English still hits ~4 chars/token; bytes only show up un-merged for the rare, unseen stuff. You get closedness for free without sacrificing compression on common text.

Chapter 6: Special & Control Tokens — the #1 Chat Bug

What they are. Special tokens are reserved vocabulary entries that carry structure, not text: <bos> (beginning of sequence), <eos> (end of sequence — the stop signal), <pad> (filler to make a batch rectangular), <unk> (the unknown token we engineered away in Chapter 5 but still reserve an ID for), and the chat control tokens <|im_start|> / <|im_end|> that mark role boundaries (Chapter 7). In the lab these are SPECIAL_TOKENS = ("<pad>", "<unk>", "<bos>", "<eos>") plus IM_START / IM_END.

Why they exist. The model needs to know where a sequence starts and stops, where one turn ends and the next begins, and which positions are real vs padding (so attention can mask the pad). These are structural signals the model is trained to recognize as dedicated symbols — they get their own embedding rows and the model attends to them as hard boundaries.

How they work under the hood — and the trap. The critical property: a special token is added structurally (e.g., encode(text, add_bos=True) prepends bos_id), and it must never be produced by tokenizing user text that merely contains the same characters. If a user types the literal string <bos> and your tokenizer maps it to the reserved bos_id, you have let user input forge a control signal — a prompt-injection-flavored bug. The byte path makes this safe: the literal <bos> decodes into ordinary byte-tokens, so test_special_token_string_is_not_split_into_pieces asserts bos_id not in encode("<bos>") and that it still round-trips. Reserved IDs are added by the framework, never minted from the content.

  encode("<bos>")            -> byte-tokens for '<','b','o','s','>'   (NOT bos_id)  ✅ safe
  encode("hi", add_bos=True) -> [bos_id, ...byte-tokens for 'hi']     (structural)  ✅ correct

Production significance. This is the #1 chat bug, in two flavors. (1) Stop-token mishandling: if you don't recognize <eos> / <|im_end|> as a stop, generation runs on past the turn into hallucinated "user" messages; if you accidentally split it, the model never stops. (2) Injection: if content can mint role/control tokens, a user can fake a system turn. Decode must skip specials (they carry no source text — the lab does if tok in specials: continue), and encode must add them structurally.

Common misconception. "<eos> is just text the model emits." <eos> is a reserved ID and the serving loop's stop condition. Treating it as ordinary text — or forgetting to register it as a stop — is why a model "won't shut up" or bleeds into a fake next turn.

Chapter 7: Chat Templates (ChatML) & Byte-for-Byte Matching

What a chat template is. A base language model just continues text. To make it a chat model, you fine-tune it on conversations rendered into a specific string format with role markers. The template is the function that turns a list of {"role", "content"} messages into that exact string. The dominant format is ChatML (OpenAI's "Chat Markup Language"), which the lab implements:

<|im_start|>system
You are concise.<|im_end|>
<|im_start|>user
Hi!<|im_end|>
<|im_start|>assistant

Each message becomes <|im_start|>{role}\n{content}<|im_end|>\n, and when you want the model to respond, you append an open assistant turn — <|im_start|>assistant\n with no closing tag — so the model's job is literally "continue this string." That trailing open turn is the add_generation_prompt=True flag (test_chat_template_exact_string pins the exact bytes).

How it works under the hood. The control tokens (<|im_start|>, <|im_end|>) are special tokens (Chapter 6) the model learned as role boundaries. The model was trained to emit <|im_end|> to end its turn — that's the stop signal the serving loop watches for. Generation is "complete the open assistant turn until you produce <|im_end|>."

Why byte-for-byte matching is non-negotiable. The model learned the exact boundary bytes, including the newline after the role and the newline after <|im_end|>. If your serving renderer adds a space (<|im_start|>user instead of <|im_start|>user\n), drops the trailing newline, or uses a different role word than training did, you are now feeding the model a string it never saw during training. There is no error — the model just produces measurably worse completions. This is the single most common, most maddening silent quality regression in chat systems.

  TRAINING:  ...<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n
  SERVING :  ...<|im_start|>user Hi!<|im_end|> <|im_start|>assistant       ← extra spaces, no \n
                              ▲                ▲                 ▲
                       off-distribution: no exception, just worse answers

Production significance. Every model ships its own template (tokenizer_config.json carries a Jinja chat_template); tokenizer.apply_chat_template(messages, add_generation_prompt=True) renders it. Using the wrong model's template, hand-rolling the format, or a framework that "helpfully" normalizes whitespace are classic causes of a model that "got dumber" with no code change to the model.

Common misconception. "Any reasonable chat format works; the model is smart enough." It is not a robustness question — the model has memorized the exact format. Match it byte-for-byte or eat a silent quality hit. Use the model's own template; never improvise.

Chapter 8: Vocabulary-Size Tradeoffs — Compression vs the Embedding Table

The lever. Vocabulary size V is a hyperparameter you choose (via the number of merges / a target size). Typical modern values are 32k–128k (GPT-2: 50k; Llama-2: 32k; GPT-4 cl100k: ~100k; o200k: ~200k). Bigger isn't free in either direction — it's a genuine tradeoff.

The benefit of a bigger vocab — compression. More merges means more whole words/phrases are single tokens, so a given text becomes fewer tokens. Fewer tokens means: lower 2N/token compute cost (Phase 00), more text fitting in a fixed context window, and (often) better quality because the model reasons over meaning-rich units. Roughly, more vocab → fewer tokens per character (higher compression).

The cost of a bigger vocab — the embedding table. The embedding matrix and the output (LM-head) matrix are each V × d_model. Doubling V doubles those two matrices' parameters and FLOPs (and the final softmax is over V logits). For a model with d_model = 4096, going from 32k → 128k vocab adds (128k - 32k) × 4096 × 2 (in + out) ≈ 0.8B parameters — pure tax, plus the softmax cost on every step. There are also diminishing compression returns and a rare-token problem: a huge vocab has many tokens seen too rarely in training to learn good embeddings (the glitch-token risk, Chapter 9).

$$ \text{embedding+head params} = 2 \times V \times d_{\text{model}}. $$

How to decide (carry the Phase 00 habit). Score the dials: compression (→ context + $/token), embedding/softmax cost, multilingual coverage, and rare-token risk; weight them for your workload; name the deciding dial. A multilingual product weights coverage and compression heavily (favoring a bigger, balanced vocab); a single-language, latency-tight product may favor a smaller one.

Production significance. The trend is upward (GPT-4 → o200k) because larger, well-balanced vocabs improve multilingual fairness and compression, and the V·d tax is small relative to a big model. But for a small model, the embedding table can be a large fraction of total params — there the tradeoff bites hard.

Common misconception. "Bigger vocab is strictly better compression, so maximize it." Compression returns diminish, the V·d cost is real, and an oversized vocab manufactures rarely-trained tokens. There is a knee; find it for your corpus and model size.

Chapter 9: Tokenization Pitfalls — Numbers, Whitespace, Fairness, Glitch Tokens

Tokenization causes a cluster of famous, non-obvious model failures. Knowing their root cause is a senior signal.

Numbers and arithmetic. BPE merges digits by frequency, not by place value, so "327" might be one token while "328" is two, and "1234567" fragments unpredictably. The model sees inconsistent, non-positional chunks, which is a major reason LLMs are shaky at multi-digit arithmetic and digit manipulation. The fix some models adopt: split every digit into its own token (a pre-tokenizer rule), giving a consistent positional representation. When someone asks "why is GPT bad at math," the first-order answer is tokenization, not reasoning.

Whitespace. Spaces are part of tokens (often a leading-space token like " the"), and exact whitespace must survive the round-trip — otherwise the model can't reproduce code indentation, can't count spaces, and decode(encode(x)) != x. The lab guarantees this by splitting on \S+|\s+ so whitespace runs are their own words (test_whitespace_only, the multi-space and \n/\t round-trip cases). Whitespace bugs are why a model mangles Python indentation or "loses" blank lines.

Multilingual fairness. If the BPE corpus is English-heavy, English compresses to ~4 chars/token but a low-resource language may compress to ~1–2 chars/token — its text becomes 2–3× more tokens. Consequences: those users pay 2–3× more (cost is per token), hit the context limit sooner, and the model has less effective capacity for their text. This fertility gap is a real fairness and economics problem, and it's a direct artifact of whose data trained the tokenizer. Balanced multilingual corpora and bigger vocabs (Chapter 8) narrow it.

Glitch / "SolidGoldMagikarp" tokens. A token can land in the vocabulary (it appeared in the BPE corpus — e.g., a Reddit username or a scraping artifact) but be rare or absent in the model's training data. Its embedding row is then barely trained — essentially random — so when that token appears at inference the model behaves bizarrely: refuses, hallucinates, repeats, or emits gibberish. The infamous SolidGoldMagikarp was one such token. Root cause: a mismatch between the tokenizer's corpus and the model's training corpus. Mitigations: train them on aligned data, filter pathological tokens, and monitor rare-token activations.

Production significance. Each of these is a recurring support ticket: "bad at math," "broke my code's indentation," "non-English costs more," "weird output on this one string." A senior names the tokenizer cause immediately instead of blaming the model.

Common misconception. "These are model intelligence limits." They are representation artifacts of how text was chopped before the model ever saw it. Change the tokenization and several of them improve — which is why digit-splitting and balanced multilingual vocabs exist.

Chapter 10: The Bridge to Multimodal — Image Patches as Tokens

The unifying idea. A transformer doesn't care that its tokens came from text. It consumes a sequence of d_model vectors and does attention over them. So to make a model see, you only need to turn an image into a sequence of d_model vectors — i.e., into tokens. Everything you learned about sequences, special tokens, and templates carries over unchanged.

How an image becomes tokens (the ViT recipe). Cut the image into a grid of fixed-size patches (e.g., 16×16 pixels). Flatten each patch and pass it through a small linear projection to get one d_model vector — a patch embedding, the image's analog of a token embedding. A 224×224 image at 16×16 patches yields a 14×14 grid = 196 patch tokens. Add positional information (so the model knows the 2D layout), and you now have a sequence the transformer can attend over, exactly like text.

   image ──► [16×16 patches] ──► linear proj ──► 196 patch vectors  (each d_model)
                                                        │
   text  ──► [BPE tokens] ───────► embed ────► text vectors (each d_model)
                                                        │
                         ┌──────────────────────────────┘
                         ▼
   "What is in <image> ?"  →  [..., <image-placeholder expands to 196 patch vectors>, ...]
                                                        │
                                                  one mixed sequence → transformer

The interface trick (why this connects to Chapter 6). A multimodal model reserves an <image> special token. At input assembly, that one placeholder is expanded into k patch slots, and the patch embeddings are spliced into those positions in the sequence. The text tokens and image tokens live in the same sequence, the same attention, the same template — tokens are the universal interface. Audio is the same story (frames/spectrogram patches → tokens).

Production significance. This is how GPT-4V, Llava, Qwen-VL, and friends work: a vision encoder produces patch tokens, a projector maps them into the language model's embedding space, and they're inserted at the <image> placeholder. The full mechanism — patch embedding, the vision encoder, cross-modal alignment — is Phase 04. For now, internalize the one sentence: an image is just more tokens.

Common misconception. "Multimodal models have a totally separate 'vision brain.'" There's a vision encoder up front, but the heavy lifting happens in the same transformer over a single token sequence. Once you see "everything is tokens," multimodality stops being magic and becomes plumbing — the same plumbing this phase is about.

Lab Walkthrough Guidance

The lab (lab-01-bpe-tokenizer) turns these chapters into code. Suggested order (matches the file):

1. The byte bijection (BYTE_ENCODER / BYTE_DECODER) — Chapter 5. Given; understand that it maps each of the 256 bytes to a clean printable char and is invertible.
2. train_bpe — Chapters 3–4. Pre-tokenize on \S+|\s+, count frequency-weighted pairs, pick the most frequent with the lexicographic tie-break, rewrite the words, repeat. Assemble the vocab: specials, then 256 bytes, then merges. test_most_frequent_pair_merged_first and test_tie_break_is_lexicographically_smallest are the correctness anchors.
3. BPETokenizer.__post_init__ — Chapters 5–6. Validate the codec is closed (all specials + all 256 bytes present, unique IDs); build id_to_token and merge_ranks.
4. _bpe + encode — Chapter 4. Apply merges in rank order (lowest rank first). Add <bos> / <eos> structurally.
5. decode — Chapters 5–6. Skip specials, map byte-tokens back to bytes, UTF-8 decode. This is the soul test: test_decode_encode_roundtrip across ASCII/spaces/newlines/accents/CJK/emoji.
6. render_chat — Chapter 7. Emit <|im_start|>{role}\n{content}<|im_end|>\n per message and the open assistant turn; test_chat_template_exact_string pins it byte-for-byte.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked example.

Success Criteria

• All tests pass against your lab.py and against LAB_MODULE=solution.
• You can run one BPE merge step by hand (count pairs, pick the most frequent, apply the lexicographic tie-break, rewrite) and explain why the lab's first merge is ('a','t'), not ('t','h').
• You can explain why the merge order is the encoding priority and why the tie-break makes training a pure function of (corpus, num_merges).
• You can explain why byte-level BPE has no <unk> and why decode(encode(x)) == x holds for emoji and CJK — and demonstrate it.
• You can explain why encode("<bos>") must not contain bos_id, and why that is a security property.
• You can write a ChatML turn byte-for-byte and name the failure mode of a whitespace mismatch.
• You can state the vocab-size tradeoff (compression vs 2 · V · d_model) and the deciding dial for a given workload.
• You can name three tokenization pitfalls and their root cause, and explain "an image is just more tokens."

Interview Q&A

• "Char, word, or subword — and why did subword win?" — Char: tiny closed vocab but very long sequences (expensive 2N/token, O(T²) attention). Word: short sequences but unbounded vocab and lossy <unk>. Subword (BPE) balances both, and byte-level makes it closed. ~4 chars/token for English is the number.
• "Walk me through one BPE merge step." — Count all adjacent symbol pairs weighted by word frequency; merge the most frequent (tie-break: lexicographically smallest pair); rewrite every word fusing that pair; repeat. The ordered merge list is the encoding priority.
• "Why does byte-level BPE have no <unk>?" — The base alphabet is the 256 byte values, so every possible input decomposes into base byte-tokens. Nothing is out-of-vocabulary; the codec is closed and decode(encode(x)) == x exactly.
• "How do you guarantee a deterministic tokenizer?" — A defined tie-break (lexicographically smallest pair on a frequency tie) makes the merge list a pure function of (corpus, num_merges), independent of dict order or corpus line order. Determinism is engineered, not assumed.
• "What's the #1 chat bug you've seen?" — A chat-template mismatch: the served format differs from training (a stray space, a missing newline, the wrong role word, or the wrong model's template). No error, just silently worse answers. Always use the model's own template, byte-for-byte.
• "Why must encoding the literal text <bos> not yield the <bos> id?" — Special tokens are structural IDs added by the framework, not minted from content. If user text could forge a control token, that's a prompt-injection vector. The byte path makes <bos> decode into ordinary byte-tokens.
• "Why are LLMs bad at arithmetic?" — Largely tokenization: BPE chunks digits by frequency, not place value, so numbers fragment inconsistently. Digit-splitting pre-tokenizers fix much of it. It's a representation artifact, not (only) a reasoning limit.
• "What is a glitch token like SolidGoldMagikarp?" — A token that's in the vocabulary (it appeared in the BPE corpus) but barely or never in the model's training data, so its embedding is essentially random and the model behaves bizarrely when it appears. Root cause: tokenizer-corpus vs training-corpus mismatch.
• "How does vocab size trade off?" — Bigger vocab → better compression (fewer tokens → cheaper, more context, often better) but a larger 2 · V · d_model embedding+head cost and more rarely-trained tokens. 32k–128k is the usual range; name the deciding dial for the workload.
• "How does multimodal input reuse this?" — An image is cut into patches; each patch is linearly projected to a d_model vector (a patch token) and spliced into the sequence at a reserved <image> placeholder. Same sequence, same attention, same template — tokens are the universal interface (mechanism in Phase 04).

References

• Sennrich, Haddow, Birch, Neural Machine Translation of Rare Words with Subword Units (2016) — the paper that introduced BPE to NLP.
• Radford et al., Language Models are Unsupervised Multitask Learners (GPT-2, 2019) — byte-level BPE and the bytes_to_unicode trick.
• OpenAI tiktoken — the production byte-level BPE tokenizer and cl100k_base / o200k_base encodings: https://github.com/openai/tiktoken.
• HuggingFace tokenizers docs — BPE, BpeTrainer, pre-tokenizers, and apply_chat_template: https://huggingface.co/docs/tokenizers and the Chat Templates guide.
• Kudo, Richardson, SentencePiece: A simple and language independent subword tokenizer (2018) — and Kudo, Subword Regularization (the unigram-LM alternative to BPE).
• Andrej Karpathy, minBPE and his "Let's build the GPT Tokenizer" walkthrough — the clearest from- scratch BPE build: https://github.com/karpathy/minbpe.
• OpenAI ChatML documentation — the chat markup format and its control tokens.
• Rumbelow & Watkins, SolidGoldMagikarp (LessWrong, 2023) — the glitch-token investigation.
• Dosovitskiy et al., An Image is Worth 16×16 Words (ViT, 2020) — patches as tokens (Phase 04 preview).

Hitchhiker's Guide — Tokenization & the Multimodal Input Pipeline

The compressed practitioner tour. If WARMUP.md is the professor, this is the senior who leans over and says "here's what you actually need to remember about the thing that turns text into numbers."

The 30-second mental model

The model speaks integers, not text. The tokenizer is the whole bridge: encode text → IDs, decode IDs → text. Use subword units (between char-level, which is too long, and word-level, which is too big and <unk>-lossy), learned by BPE = "greedily merge the most frequent adjacent pair, repeat." Make the base alphabet the 256 bytes so the codec is closed — no <unk>, decode(encode(x)) == x for anything (emoji, CJK, binary). Special tokens (<bos>, <eos>, <|im_start|>) are structural IDs, never minted from user text. The chat template must match training byte-for-byte or quality silently dies. And remember: an image is just more tokens (patches → vectors → spliced into the sequence).

The numbers to tattoo on your arm

The framework one-liners (where these live in real tools)

# OpenAI / tiktoken — byte-level BPE, closed vocab, no <unk>
import tiktoken
enc = tiktoken.get_encoding("cl100k_base")
ids = enc.encode("the cat");  enc.decode(ids)            # exact round-trip
len(enc.encode(text))                                    # token count == your bill

# HuggingFace — train, encode, and (critically) the chat template
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
tok.bos_token_id, tok.eos_token_id, tok.pad_token_id     # the structural IDs
prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
# ^ ALWAYS use this. Never hand-roll the chat string.

# SentencePiece — BPE or unigram, language-independent
# spm_train --model_type=bpe --vocab_size=32000 --input=corpus.txt

War stories

• The chat template that ate 8 points of quality. A team hand-built the prompt string with an extra space after the role (<|im_start|>user instead of \n). No crash. Evals quietly dropped; it took a week to find. Fix was one line: apply_chat_template. The model had memorized the exact bytes.
• The <eos> that wouldn't stop. The serving loop didn't register the model's end token as a stop, so generations ran past the turn and hallucinated a fake user reply. Classic special-token mishandling — <eos>/<|im_end|> is a stop condition, not just text.
• The triple-billed customer. A non-English user complained their "same length" messages cost 3× and truncated early. The English-trained tokenizer had high fertility on their language — 2–3 bytes/token vs ~4 chars/token for English. Not a bug in the meter; a bias in the vocab.
• The cursed string. One specific username made the model emit gibberish on sight. A glitch token (SolidGoldMagikarp-style): in the vocab, but unseen in training → random embedding → chaos.
• "Why is it bad at math?" Asked to add 7-digit numbers, the model flailed. Root cause wasn't reasoning — BPE had chunked the digits non-positionally. Digit-splitting tokenizers exist for exactly this.

Vocabulary (rapid-fire)

• BPE — Byte-Pair Encoding; greedily merge the most frequent adjacent pair, repeat.
• Merge / merges.txt — the ordered, frozen list of learned merges; its order is the encoding priority.
• Byte-level / byte fallback — base alphabet is the 256 bytes; makes the codec closed (no <unk>).
• Special / control token — reserved structural ID: <bos>, <eos>, <pad>, <|im_start|>.
• Chat template / ChatML — the exact string format the chat model was trained on.
• Fertility — tokens per word; the multilingual fairness/cost knob.
• Compression — chars per token; sets context length and $/token.
• Glitch token — in the vocab but undertrained → bizarre behavior.
• Patch token — an image patch projected to a d_model vector; multimodal's "token."
• add_generation_prompt — append the open <|im_start|>assistant\n turn for the model to complete.

Beginner mistakes

• Hand-rolling the chat string instead of apply_chat_template (the silent quality killer).
• Treating <eos>/<|im_end|> as text instead of a registered stop.
• Letting user text mint control tokens (encoding <bos> to bos_id) — an injection vector.
• Assuming decode(encode(x)) == x without testing spaces, newlines, CJK, and emoji.
• Quoting "tokens ≈ words" — they're frequency fragments; the, ization, and a 5-piece name all happen.
• Blaming the model for arithmetic / indentation / multilingual-cost issues that are tokenization artifacts.
• Maxing out vocab size for "better compression" while ignoring the 2 · V · d_model tax and rare- token risk.

The one thing to take away

The tokenizer is the contract between text and the model, and its failures are silent. Before you blame the model, check the boundary: is the round-trip exact, are the special tokens structural, and does the chat template match training byte-for-byte? If yes, it's the model. If you didn't check, you're guessing.

Brother Talk — Tokenization & the Multimodal Input Pipeline

Off the record. The stuff I'd tell you over a beer, not in a design review.

This is the phase everyone skips and then quietly suffers for. Tokenization looks like boring preprocessing — a thing you import and forget. So people do, and then they spend a week debugging "the model got dumber" only to discover it was a stray space in a chat string. The engineers who actually understand the tokenizer are the ones who fix that bug in ten minutes while everyone else is staring at the model weights. Boring topic, outsized leverage.

The 2 a.m. truth: it's almost always the chat template. When a chat model's quality tanks with no change to the model, the first thing I check — every single time — is whether the served prompt matches the training format byte-for-byte. A space after the role, a missing newline, the wrong model's template, a framework "helpfully" normalizing whitespace. There's no exception, no error — the model just gets quietly worse because you fed it a string it never saw in training. It memorized those exact bytes. Memorize this failure mode like a phone number, because you will meet it, and the person who names it instantly looks like a wizard.

Use apply_chat_template. Do not be clever. I know it's tempting to hand-build the prompt string — it feels more "in control." Resist. The model's own template ships with the model for a reason. Hand-rolling it is how you reintroduce the exact silent bug above. The boring call is the correct call.

decode(encode(x)) == x is the soul test, and it's also a trap if you only test ASCII. Plenty of tokenizers "work" until someone pastes an emoji, a CJK sentence, or a tab-indented code block, and suddenly bytes go missing. Byte-level BPE gives you a closed codec for free — no <unk>, exact round-trips for anything — and the moment you internalize why (everything is bytes, 256 base symbols, nothing is OOV), a whole class of data-loss bugs just stops existing for you. Test the weird stuff, always.

Half the "the model is dumb" complaints are tokenization, not intelligence. Bad at multi-digit math? BPE chunked the digits non-positionally. Mangled your Python indentation? Whitespace tokens. Non-English costs triple? Fertility — the vocab was trained on English. A glitch token making it spew gibberish? An undertrained embedding from a tokenizer/training-corpus mismatch. When you can route these to the real cause out loud, you stop sounding like a user and start sounding like someone who owns the stack.

Special tokens are a security surface, not just plumbing. If user content can mint a <bos> or a role token, a user can forge a system turn — that's prompt injection wearing a tokenizer costume. The rule is simple and absolute: control tokens are structural, added by your code, never minted from content. The byte path enforces it. Get this reflex early; it'll matter more as you build agents.

What's worth caring about: the round-trip invariant, the byte-for-byte chat template, and the structural-vs-textual line on special tokens. What's not worth losing sleep over: the exact merge count, whether your vocab is 48k or 52k, or memorizing GPT-2's bytes_to_unicode table by heart. Get the three invariants right and you've covered 95% of the real-world pain.

The multimodal punchline is smaller than it sounds. People treat "the model can see images now" as some separate magic. It isn't. An image is cut into patches, each patch becomes a vector, and the vectors are spliced into the same token sequence behind an <image> placeholder. Same attention, same template, more tokens. Once "everything is tokens" clicks, multimodality (Phase 04) stops being intimidating and becomes the plumbing you already understand.

The career framing. This phase doesn't make you look impressive in a demo — nobody claps for a tokenizer. But it's the difference between an engineer who uses an LLM and one who owns the input contract, and that's the person teams trust when the edges break. Get the invariants into your hands by building the lab, and you'll debug in minutes what costs others days. Now go make pytest green.

Lab 01 — BPE Tokenizer + Chat-Template Renderer

Phase: 01 — Tokenization & the Multimodal Input Pipeline Difficulty: ⭐⭐⭐☆☆ (the algorithm is small; the invariants are where seniors are made) Time: 3–4 hours

The model never sees your text. It sees integer IDs, and it speaks integer IDs back. The tokenizer is the entire boundary between the human world and the tensor world — and it is the single most under-respected component in the stack, the place where round-trips break, where chat quality silently dies, and where prompt-injection-flavored bugs hide. This lab builds the real thing: learn a byte-level BPE merge table from a corpus, encode/decode with a UTF-8 byte fallback so the codec is closed (no <unk> for normal text, a true decode(encode(x)) == x round-trip), reserve the special/control tokens chat models live and die by, and render a ChatML chat template byte-for-byte the way the model was trained to expect.

What you build

• train_bpe(corpus, num_merges, *, lowercase=False) — learn the ordered merge list and the vocabulary. Greedy and deterministic: pre-tokenize into words (\S+|\s+), represent each word as a tuple of base byte-tokens, then repeatedly count adjacent pairs, merge the most frequent (tie-break: lexicographically smallest pair), and rebuild. Returns (merges, vocab).
• BPETokenizer — the trained codec (a @dataclass). Validates that the vocab is closed (all 4 specials + all 256 base byte-tokens present, IDs unique) and builds the id_to_token inverse and the merge_ranks priority table.
  • BPETokenizer.from_training(corpus, num_merges, *, lowercase=False) — train + construct.
  • vocab_size, bos_id, eos_id, pad_id, unk_id — properties.
  • encode(text, *, add_bos=False, add_eos=False) — text → list of IDs, applying merges per word in learned-priority order, with the byte fallback so nothing is ever OOV.
  • decode(ids) — IDs → exact original text; specials are skipped, byte-tokens are mapped back to bytes and the byte stream is UTF-8 decoded.
• render_chat(messages, *, add_generation_prompt=True) — render a conversation into ChatML: <|im_start|>{role}\n{content}<|im_end|>\n per message, plus an open <|im_start|>assistant\n turn for generation. Must match training byte-for-byte.
• Module constants — SPECIAL_TOKENS (<pad>, <unk>, <bos>, <eos>), IM_START / IM_END (the ChatML control strings), BYTE_ENCODER / BYTE_DECODER (the GPT-2 byte↔printable-unicode bijection).

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• decode(encode(x)) == x holds for ASCII, multiple spaces, newlines/tabs, accents (café), CJK (你好世界), and emoji — the soul test. You can explain why the byte fallback makes this total.
• You can hand-count the first merge of the lab corpus and explain why ('a','t') (freq 8) beats ('t','h') (freq 6) — counting the actual most-frequent pair, not the eye-catching one.
• You can explain why train_bpe(reversed(CORPUS), …) learns the identical merges (the tie-break, not dict order, is the source of determinism).
• You can explain why encoding the literal string <bos> does not produce the reserved bos_id, and why that property matters for safety.
• You can explain why render_chat must match the training format byte-for-byte, and name the failure mode when it doesn't.

How this maps to the real stack

Limits of the miniature (be honest in the interview): real tokenizers use a more elaborate GPT-2/GPT-4 pre-tokenization regex (splitting digits, contractions, and leading spaces specially) instead of plain \S+|\s+; production trainers prune rare words and cap vocabulary at the target size (32k–128k) rather than a fixed merge count; SentencePiece operates on a unigram language model as an alternative to BPE and treats the input as a raw character stream with a ▁ meta-space; and special tokens in real systems get dedicated embedding rows the model attends to, which this lab models only as reserved string IDs.

Extensions (build these on real hardware / real libs)

• Swap the \S+|\s+ splitter for the actual GPT-2 / cl100k_base pre-tokenization regex and re-measure compression (tokens per character) on a paragraph of English.
• Add a train_to_vocab_size(corpus, target) that stops when len(vocab) == target instead of a fixed num_merges, and plot compression vs vocab size — find the knee.
• Train on a multilingual corpus and measure the fertility (tokens per word) gap between English and a low-resource language; explain the fairness/cost implication.
• Build the <image> placeholder path: reserve an image special token, expand it to k patch IDs, and show how a vision tokenizer slots patch embeddings into the same sequence (forward reference to Phase 04).
• Compare your merges against tiktoken on the same text: tiktoken.get_encoding("gpt2").

Interview / resume

• Talking points: "char vs word vs subword — what does BPE actually buy you?" "Why does byte-level BPE have no <unk>?" "Walk me through one BPE merge step and the tie-break." "Why must a chat template match the training format byte-for-byte, and what breaks if it doesn't?" "What is a glitch token and where do they come from?"
• Resume bullet: Implemented a byte-level BPE tokenizer from scratch — deterministic merge-training with lexicographic tie-break, a closed UTF-8 byte-fallback codec with a verified decode(encode(x)) == x round-trip across ASCII/CJK/emoji, reserved special/control tokens, and a byte-exact ChatML chat-template renderer.

Phase 02 — The Transformer From Scratch

The phase where "I use transformers" becomes "I can build one." Phase 00 taught you to treat a model as a physical system with a compute, memory, and bandwidth budget; Phase 01 turned text and images into token IDs. Now we open the box those tokens flow through and build the actual mechanism — attention, the causal mask, multi-head, RoPE, RMSNorm, the SwiGLU MLP, the residual stream, the full forward pass, and the KV-cache — by hand, in pure Python, with no framework hiding the arithmetic. This is the architecture every later phase trains, quantizes, serves, and attacks.

Why this phase exists

You cannot debug, optimize, or reason about a thing you can only call. The senior who can answer "why √d_k," "what exactly is the mask doing," "where do the parameters live," and "why is the KV-cache correct" is the one trusted to own the model stack. Every one of those answers is a few lines of code you will write here.

The transformer is also the load-bearing abstraction of the entire curriculum. Autograd (P03) differentiates this forward pass. Fine-tuning (P05) injects adapters into these linear layers. Quantization (P06) compresses these weights. Serving (P09) caches these keys and values with PagedAttention. Long-context, FlashAttention, GQA, MoE — all are modifications to the block you build in this lab. Get the forward pass in your hands and the rest of the track is changes to a thing you already understand, not new magic.

Three ideas do the heavy lifting, and you should leave the phase able to defend each from first principles:

1. Attention is content-based, weighted retrieval. softmax(Q·Kᵀ/√d_k)·V scores every query against every key (similarity), normalizes to a distribution, and mixes the value rows by it. The causal mask makes it autoregressive; multiple heads make it multi-relational.
2. Position is injected, not inherent. Attention is permutation-equivariant on its own, so order has to be added — by RoPE (rotate Q/K so the score depends on the relative offset), the modern default over learned absolute embeddings.
3. The residual stream is the highway, and most of the model is the MLP. Pre-norm blocks read from the residual stream, compute a correction, and write it back; ~2/3 of the parameters live in the feed-forward, not in attention.

Concept map — a transformer block

   token ids ──► embedding lookup ──► + position (RoPE) ──► residual stream X  ┐
                                                                               │
   ┌───────────────────── one pre-norm block (× n_layers) ────────────────────┤
   │                                                                          │
   │      X ──► RMSNorm ──► Multi-Head Attention ──► (+) ──► X'               │
   │                          │  split d_model into h heads    ▲              │
   │                          │  softmax(Q·Kᵀ/√d_k + mask)·V   │ residual     │
   │                          └───────────────────────────────┘              │
   │                                                                          │
   │      X' ──► RMSNorm ──► SwiGLU FFN (~2/3 of params) ──► (+) ──► X''      │
   │                          down(silu(x·W_gate) ⊙ (x·W_up))    ▲           │
   │                          └──────────────────────────────────┘ residual  │
   └──────────────────────────────────────────────────────────────────────────┘
                                       │
        final RMSNorm ──► LM head (linear to vocab) ──► logits [T, vocab]
                                       │
                       (the causal mask: token i attends only to j ≤ i)

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py. Pure stdlib, deterministic from a seed, offline.

Integrated scenario ideas

• Prove the KV-cache to a skeptic. Run full causal attention, then run the same input through attention_with_kv_cache step by step, and show the last-position output is identical — then explain why (the last query under a causal mask sees exactly tokens 0..t, which is the cache).
• Find where the future leaks. Remove the causal mask and show a token's prediction changes when you edit a later token — the canonical training-data-leakage bug — then put the mask back.
• Locate the parameters. Count params in attention vs the FFN for a real config (e.g. Llama-style d_model, d_ff ≈ 8/3·d_model) and show the FFN is ~2/3 — connect to Phase 00's 2N/token.
• RoPE is relative. Take two tokens, RoPE them at positions (p, q) and at (p+k, q+k), and show the QK dot product is unchanged — the property that lets RoPE extrapolate context length.
• Norm ablation. Swap RMSNorm for LayerNorm in the block and show the forward pass still runs; discuss why the field moved to RMSNorm (cheaper, no bias, equally good).

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive softmax(Q·Kᵀ/√d_k)·V and explain every term, including why √d_k is there.
• You can explain the causal mask as additive -inf → zero softmax weight → autoregression.
• You can explain why multiple heads exist and what _split_heads/_concat_heads do.
• You can explain RoPE's rotation, why it preserves the norm, and why the score is relative.
• You can state where the parameters live (~2/3 FFN) and contrast RMSNorm with LayerNorm.
• You can prove the KV-cache equals the full causal recompute's last row and say why.

Key takeaways

• Attention is weighted retrieval, the mask makes it causal. softmax(Q·Kᵀ/√d_k + mask)·V is the whole engine; the -inf mask is the one line that makes a transformer a language model instead of an encoder.
• Position is added on purpose. Self-attention is order-blind; RoPE injects relative position by rotation, which is why modern long-context tricks (NTK/YaRN scaling) are RoPE-base tweaks.
• Most of the model is the MLP. ~2/3 of the parameters are in the feed-forward — which is why the 2N/token rule from Phase 00 holds and why MoE puts its experts there.
• The KV-cache is correct because of the causal mask. Decode never recomputes past K/V, and the result is bit-for-(float)-bit the full recompute's last row — the identity that makes serving (P09) possible and PagedAttention worth building.

Next: Phase 03 — Autograd & Training Dynamics.

Warmup Guide — The Transformer From Scratch

Zero-to-senior primer for Phase 02. We start from "a token is an integer" and build the entire decoder-only transformer one mechanism at a time — embeddings, self-attention, the causal mask, multi-head, the KV-cache, RoPE, RMSNorm, the SwiGLU MLP, the residual stream, the full forward pass — and finish with FlashAttention as the IO-aware version of the attention you just wrote. Each chapter is the same shape: what it is → why it exists → how it works under the hood (mechanism, a diagram, a little math) → why it matters in production → the misconception that trips people up. By the end you can build a GPT forward pass from nothing and defend every line.

Table of Contents

• Chapter 1: From Token IDs to Vectors — The Embedding Table
• Chapter 2: Self-Attention — Q, K, V and the Dot Product as Similarity
• Chapter 3: Why /√d_k and Why softmax
• Chapter 4: The Causal Mask & Autoregression
• Chapter 5: Multi-Head Attention — Many Subspaces at Once
• Chapter 6: MQA / GQA and the KV-Cache
• Chapter 7: Position — RoPE vs Learned Absolute
• Chapter 8: RMSNorm vs LayerNorm
• Chapter 9: The Feed-Forward / SwiGLU — Where the Parameters Live
• Chapter 10: The Residual Stream & the Pre-Norm Block
• Chapter 11: The Full Forward Pass — Embedding to Logits
• Chapter 12: FlashAttention — The IO-Aware Fused Version
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: From Token IDs to Vectors — The Embedding Table

From zero. Phase 01 turned text into a list of integers — token IDs like [1, 5, 9, 2]. A neural network cannot multiply an integer "ID 5" by a weight matrix in any meaningful way; the ID is a name, not a quantity. The first thing every transformer does is replace each ID with a learned vector. That is the embedding table: a [vocab, d_model] matrix where row i is the vector for token i. "Embedding lookup" is literally X = [embed[tid] for tid in token_ids] — indexing rows.

Why it exists. We want tokens with related meaning to sit near each other in a continuous space so that the same downstream matmuls can act on "king" and "queen" coherently. The table is learned: gradient descent moves each row until the geometry is useful. d_model (the model width — 4096 for Llama-2-7B, 768 for GPT-2-small) is the dimensionality of that space and the single number that flows through the entire network unchanged.

Under the hood. The lab's gpt_forward starts exactly here:

token_ids = [1, 5, 9, 2]
X = [ embed[1],      # a d_model-length vector
      embed[5],
      embed[9],
      embed[2] ]     # X is now [T, d_model] = [4, d_model]

There is no matmul yet — just a gather. From here on, every position is a d_model vector and the sequence is a [T, d_model] matrix. (Many models tie the embedding table to the final LM head — the same [vocab, d_model] weights are reused, transposed, to turn the last hidden state back into logits, saving vocab·d_model parameters. The lab keeps them separate for clarity.)

Production significance. The embedding + LM-head rows are 2·vocab·d_model parameters — for a 128k-token vocab and d_model=4096 that is ~1B params before any transformer layer, which is why large vocabularies are expensive and why weight tying matters.

Common misconception. "The embedding encodes position." It does not — embed[5] is the same vector wherever token 5 appears. Position is injected separately (Chapter 7). The embedding encodes identity, not order.

Chapter 2: Self-Attention — Q, K, V and the Dot Product as Similarity

What it is. Self-attention lets every position look at every other position and pull in information from the relevant ones. It is the only operation in a transformer that mixes across positions — everything else (norm, MLP, residual) acts on each position independently. Without attention, the model would process each token in isolation.

Why Q, K, V. Borrow the analogy of a soft dictionary lookup. Each position emits three learned projections of its embedding:

• a query q — "what am I looking for?"
• a key k — "what do I offer to others searching?"
• a value v — "what do I actually contribute if attended to?"

They are just three linear maps of the same input X:

$$ Q = X W_Q, \qquad K = X W_K, \qquad V = X W_V. $$

In the lab these are three linear(X, Wq/Wk/Wv) calls. Splitting the roles is what lets a position ask for something different from what it offers.

The dot product is similarity. To decide how much query i should attend to key j, take their dot product q_i · k_j. A dot product is large and positive when two vectors point the same way — so it is a learned, content-based similarity score. Computing all of them at once is a matmul:

$$ S = Q K^\top \quad (\text{shape } [T, T]),\qquad S_{ij} = q_i \cdot k_j. $$

Row i of S is "how much does query i match every key." We then turn each row into weights and take a weighted average of the value rows:

$$ \text{Attention}(Q,K,V) = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}} + \text{mask}\right) V. $$

Under the hood — a worked picture.

                 keys (what each position offers)
                 k0    k1    k2
              ┌──────────────────┐
   query q1 ──┤ q1·k0 q1·k1  -∞  │  scores row for position 1 (k2 is the future → masked)
              └──────────────────┘
                     │ softmax over the row
                     ▼
              [ 0.6,  0.4,  0.0 ]   attention weights (sum to 1)
                     │
                     ▼  weighted sum of value rows
   output1 = 0.6·v0 + 0.4·v1 + 0.0·v2

That is exactly scaled_dot_product_attention in the lab: scores, scale, mask, softmax per row, then a weighted sum of V.

Production significance. This single op is where a model resolves coreference ("it" → which noun?), copies from context, does in-context learning, and attends to retrieved documents in RAG. When generation hallucinates or ignores the prompt, you are debugging these weights.

Common misconception. "Q, K, V are different inputs." In self-attention they are three projections of the same sequence X. (Cross-attention — encoder→decoder, or attending to image patches — is the variant where K/V come from a different source than Q. The lab is self-attention.)

Chapter 3: Why /√d_k and Why softmax

The softmax. We need to turn a row of raw scores into a set of non-negative weights that sum to 1 so the output is a convex combination of value rows (a weighted average, never an extrapolation). softmax does exactly that:

$$ \text{softmax}(s)_j = \frac{e^{s_j}}{\sum_k e^{s_k}}. $$

It is also differentiable and "soft" — small score changes nudge the weights smoothly, which is what gradient descent needs. The lab's softmax subtracts the max before exponentiating:

$$ \text{softmax}(s)_j = \frac{e^{s_j - \max(s)}}{\sum_k e^{s_k - \max(s)}}. $$

This changes nothing mathematically (softmax is shift-invariant) but everything numerically: e^{1002} overflows to inf and the result is NaN, while e^{1002-1002}=1 is fine. As a bonus, a masked -inf score maps to e^{-\infty}=0 — exactly zero weight, which is what the causal mask needs. That numerical-stability detail is part of the lesson, and test_softmax_numerically_stable_large_logits checks it.

Why divide by √d_k. The scores are dot products over d_k dimensions. If the query and key entries are roughly independent with unit variance, the dot product q·k = Σ_{i=1}^{d_k} q_i k_i has variance ≈ d_k and standard deviation ≈ √d_k. So as the head dimension grows, the raw scores grow too — and large scores push softmax toward a near one-hot distribution (it saturates). A saturated softmax has almost-zero gradient, so training stalls. Dividing by √d_k rescales the scores back to roughly unit variance regardless of d_k:

$$ \text{Var}!\left[\frac{q\cdot k}{\sqrt{d_k}}\right] \approx \frac{d_k}{(\sqrt{d_k})^2} = 1. $$

   raw scores (large d_k):  [ 8.0,  0.1, -7.9 ]  → softmax ≈ [1.00, 0.00, 0.00]  (saturated, dead gradient)
   /√d_k scaled:            [ 2.0,  0.0, -2.0 ]  → softmax ≈ [0.84, 0.11, 0.02]  (soft, trainable)

Production significance. This is the canonical "why √d_k" interview question. The honest answer is "to keep the softmax in its sensitive, high-gradient regime as the head dimension scales." test_sqrt_dk_scaling_is_present in the lab detects the factor by comparing to a hand-computed scaled softmax.

Common misconception. "√d_k is just an arbitrary normalization." It is a variance argument: without it, wider heads systematically saturate the softmax and the model trains worse. The constant comes from the standard deviation of a sum of d_k independent products.

Chapter 4: The Causal Mask & Autoregression

What it is. A decoder-only language model predicts the next token, so when computing the output at position i it must only see positions 0..i — never the future. The causal mask enforces this. It is a [T, T] matrix added to the scores before softmax: 0 on and below the diagonal, -∞ strictly above it.

$$ \text{mask}_{ij} = \begin{cases} 0 & j \le i \ -\infty & j > i \end{cases} $$

Why additive -∞. After adding the mask, every "future" score is -∞, and softmax(-\infty)=0. So position i puts exactly zero weight on every j > i — the future contributes nothing, and it contributes nothing to the gradient either. The lab builds this in causal_mask and softmax handles the -∞ cleanly because it subtracts the max.

Under the hood.

   causal_mask(4):                attention weights after softmax (row = a query):
   j:  0    1    2    3
   ┌─────────────────────┐        token 0 sees only {0}
 0 │  0   -∞   -∞   -∞   │        token 1 sees {0,1}
 1 │  0    0   -∞   -∞   │        token 2 sees {0,1,2}
 2 │  0    0    0   -∞   │        token 3 sees {0,1,2,3}
 3 │  0    0    0    0   │        the strict upper triangle is always 0 weight
   └─────────────────────┘

test_causal_attention_zero_in_upper_triangle proves the consequence: token 0 can only attend to V[0], so its output is exactly V[0].

Why it's called autoregression. "Autoregressive" means the model generates one token at a time, each conditioned only on what came before. The causal mask is what makes a parallel training pass (all positions at once) compute the same thing as the sequential generation it will later do — so you can train on a whole sequence in one shot and the model still learns next-token prediction honestly.

Production significance. Training on T positions in parallel — instead of T separate forward passes — is the entire reason transformers train efficiently. A bug here is catastrophic and silent: if the mask leaks even one future position, the model "cheats" during training (sees the answer), gets a deceptively low loss, and then generates garbage at inference because the future it relied on isn't there. This is a real, classic bug (Chapter: HITCHHIKERS war stories).

Common misconception. "The mask is set to a large negative number like -1e9." Many implementations do use a large finite negative (because -inf can produce NaN if a whole row is masked). The concept is -∞ → zero weight; the lab uses true -inf and guards the all-masked case by construction (the diagonal is always 0, so no row is ever fully masked).

Chapter 5: Multi-Head Attention — Many Subspaces at Once

What it is. Instead of one attention over the full d_model width, split it into h heads, each of width d_head = d_model / h, run attention independently in each, then concatenate the results and apply an output projection W_O. One attention can only express one weighted-average pattern per position; h heads express h of them at once.

Why multiple heads. Different relationships want different similarity functions. One head might track syntactic dependency (verb → its subject), another coreference (pronoun → antecedent), another positional/local patterns. With a single head, all of that has to be averaged into one set of weights and they interfere. Splitting into subspaces lets the model attend to several relationships in parallel and combine them.

Under the hood.

   X [T, d_model]
     │  project to Q, K, V  (each [T, d_model])
     ▼
   split each into h heads ──► Q_h, K_h, V_h   each [T, d_head]
     │
     ├─ head 0: SDPA(Q0, K0, V0, mask) ─► out0 [T, d_head]
     ├─ head 1: SDPA(Q1, K1, V1, mask) ─► out1 [T, d_head]
     └─ ...                                ...
     │  concat heads back to [T, d_model]
     ▼
   linear(concat, W_O)  ──►  [T, d_model]

That is multi_head_attention exactly: linear for the projections, _split_heads to slice each [T, d_model] into h matrices of [T, d_head], scaled_dot_product_attention per head, _concat_heads to glue them back, then the output projection W_O. Note d_model must be divisible by h — test_mha_rejects_indivisible_heads checks the guard.

Production significance. h is a real config knob (num_attention_heads). Mechanistic- interpretability work names specific heads ("induction heads" that do in-context copying). And the per-head split is precisely what GQA exploits next: share K/V across heads to shrink the cache.

Common misconception. "More heads = more parameters / more compute." Splitting d_model into heads is a reshape — the total Q/K/V projection size is unchanged. Heads change how the same budget is used (many narrow subspaces vs one wide one), not how big it is.

Chapter 6: MQA / GQA and the KV-Cache

The KV-cache (recap + mechanism). Phase 00 established the memory math; here is the mechanism. During generation ("decode"), the model emits one token at a time. To produce token t+1 it needs attention over all previous keys and values. Recomputing K/V for every past token at every step is O(T²) total work. Instead we cache the K and V vectors of every past token, so each new step computes K/V for just the one new token and reads the rest from the cache — O(T) total. The lab's attention_with_kv_cache is one such step: append the new K/V, attend the new query over the whole cache, return the context plus the grown cache.

The soul test. Why is the cache correct? Because under a causal mask, the last query attends to exactly tokens 0..t — which is exactly what the cache holds. So the incremental result is numerically identical (within float tolerance) to running full causal attention and taking the last row. test_kv_cache_equals_full_recompute proves this bit-for-(float)-bit, and it is the entire justification for the cache:

   full recompute:   SDPA(Q, K, V, causal_mask(T))  →  take row T-1
   kv-cache decode:  step t=0..T-1, attend new_q over cache  →  final context
   assertion:        context  ==  full[T-1]   (within approx)

MQA and GQA — shrinking the cache. The cache size is 2 · n_layers · T · d_model · bytes · batch (Phase 00). The lever is d_model in the cache — i.e. the K/V width. Two ideas reduce it:

• MQA (Multi-Query Attention): all query heads share a single K and V head. The cache stores one head's worth of K/V instead of h — an h× reduction — at a small quality cost.
• GQA (Grouped-Query Attention): the middle ground every modern model uses. Use g K/V heads (1 < g < h), each shared by a group of h/g query heads. With 8 K/V heads out of 32 query heads, the cache is 4× smaller with negligible quality loss.

$$ d_{kv} = d_{model}\cdot\frac{n_{kv}}{n_{head}} \quad\Rightarrow\quad \text{KV-cache} \propto n_{kv}. $$

   MHA   (h=4):  q0 q1 q2 q3        GQA (g=2):  q0 q1 | q2 q3      MQA (g=1):  q0 q1 q2 q3
                 k0 k1 k2 k3                     k0    |  k1                    k0
                 v0 v1 v2 v3                     v0    |  v1                    v0
                 cache = 4 K/V                   cache = 2 K/V                  cache = 1 K/V

Production significance. GQA is why models like Llama-2-70B and Mistral can serve long contexts at high concurrency without the KV-cache OOM from Phase 00. num_key_value_heads in a HF config is this exact knob. The lab implements plain multi-head; adding GQA is an extension (and ties straight back to Phase 00's cache math).

Common misconception. "GQA makes attention faster." Its main win is memory (a smaller KV-cache), which indirectly raises throughput because decode is memory-bandwidth-bound (Phase 00) — fewer bytes to read per token. The FLOPs barely change.

Chapter 7: Position — RoPE vs Learned Absolute

The problem. Self-attention is permutation-equivariant: shuffle the input positions and the outputs shuffle with them, unchanged in content. softmax(QKᵀ)V has no notion of "before" or "after." But language is ordered — "dog bites man" ≠ "man bites dog." So order must be injected.

Learned absolute positions (the old way). GPT-2 adds a learned vector pos_emb[t] to the embedding at position t. Simple, but: it can't represent positions past the trained maximum (no pos_emb[5000] if you trained to 1024), and it encodes absolute position when what attention really cares about is relative offset.

RoPE (Rotary Position Embedding, the modern default). Instead of adding a position vector, RoPE rotates the query and key vectors by an angle proportional to their position. Group the d dimensions into d/2 pairs (x_{2i}, x_{2i+1}); treat each pair as a 2-D vector and rotate it by angle pos · θ_i, where each pair has its own frequency:

$$ \theta_i = \text{base}^{-2i/d}, \qquad \begin{pmatrix} x'{2i} \ x'{2i+1} \end{pmatrix} = \begin{pmatrix} \cos(pos,\theta_i) & -\sin(pos,\theta_i) \ \sin(pos,\theta_i) & \cos(pos,\theta_i) \end{pmatrix} \begin{pmatrix} x_{2i} \ x_{2i+1} \end{pmatrix}. $$

That is _rope_vector in the lab. Low-i pairs rotate fast (high frequency, capture nearby offsets); high-i pairs rotate slowly (low frequency, capture long-range offsets) — a Fourier-feature ladder of position scales. base (default 10000; 1e6 for long-context models) sets how slowly the slowest pairs turn.

The key property — it encodes relative position. A rotation preserves length, so RoPE never changes a vector's norm (test_rope_preserves_norm), and position 0 is the identity (test_rope_position_dependent). The magic is in the dot product: rotating q by pos_q·θ and k by pos_k·θ, the attention score q'·k' depends only on the difference pos_q − pos_k, not on the absolute positions. Attention naturally cares about offset ("how far back is the thing I'm copying?"), and RoPE hands it relative offset for free.

   pair i of a vector, before and after rotation by angle pos·θ_i:

            ^                      the (even, odd) pair is a point in 2-D;
            │   • x'  (rotated)    rotation spins it by pos·θ_i around the origin.
            │  ╱                   |x'| = |x|  (norm preserved).
            │ ╱  ) angle = pos·θ_i  q rotated by pos_q·θ, k by pos_k·θ:
            │╱___________>          q'·k' depends only on (pos_q − pos_k).

Production significance. RoPE is in Llama, Mistral, Qwen, Gemma, and most current open models. Because it's relative and frequency-based, you can extend the context window after training by rescaling the frequencies — NTK-aware scaling and YaRN are exactly "tweak the RoPE base / frequencies so positions beyond the trained length still get sensible angles." That whole long-context literature is RoPE math.

Common misconception. "RoPE is added to the embedding like GPT-2's position embedding." It is a rotation applied to Q and K inside attention, not an additive vector on the residual stream. (The lab applies it once on the residual stream up front for simplicity — see the README's honest limits — but in real models the rotation lives inside each attention layer so it acts on the scores.)

Chapter 8: RMSNorm vs LayerNorm

Why normalize at all. Stacking dozens of layers, activations can drift to very large or very small magnitudes, which makes gradients explode or vanish and training unstable. A normalization layer rescales each position's vector to a controlled magnitude before each sublayer, keeping the numbers in a sane range so the network trains smoothly. It is a stability device, not a feature extractor.

LayerNorm (the original). For a vector x of width d, subtract the mean, divide by the standard deviation, then scale and shift with learned per-dimension weight and bias:

$$ y_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}},\gamma_i + \beta_i, \qquad \mu = \frac{1}{d}\sum_i x_i,\quad \sigma^2 = \frac{1}{d}\sum_i (x_i-\mu)^2. $$

It centers (zero mean) and scales (unit variance). The lab's layer_norm does exactly this; test_layer_norm_zero_mean_unit_var checks both.

RMSNorm (the modern default). Drop the mean-centering and the bias entirely; just divide by the root-mean-square and scale:

$$ y_i = \frac{x_i}{\sqrt{\frac{1}{d}\sum_j x_j^2 + \epsilon}},\gamma_i. $$

The lab's rms_norm does this; test_rms_norm_scales_to_unit_rms checks the output RMS is 1. It is cheaper (no mean, no variance subtraction, no bias term — fewer ops and fewer parameters) and, empirically, works just as well. That is why Llama, Mistral, and Gemma all use RMSNorm.

Pre-norm vs post-norm. The original transformer put the norm after the sublayer and residual add (post-norm); modern models put it before the sublayer (pre-norm), which is far more stable for deep stacks because the residual stream stays an unnormalized "clean" highway (Chapter 10). The lab is pre-norm.

   LayerNorm:  center (− mean) → scale (÷ std) → ×weight + bias      (4 reductions + bias)
   RMSNorm:                       scale (÷ rms) → ×weight            (1 reduction, no bias)
                ↑ skips mean-centering and the shift — cheaper, equally good

Production significance. It's a config choice with measurable wins in throughput and parameter count at scale, and "why did the field move LayerNorm → RMSNorm, pre-norm → post-norm" is a common architecture-depth interview question. The answer: cheaper + just as good (RMSNorm), and more stable for depth (pre-norm).

Common misconception. "RMSNorm is a lossy approximation of LayerNorm." It's a different normalization (no centering), not an approximation — and on transformer activations the centering turns out not to matter much, so you pay less for the same stability.

Chapter 9: The Feed-Forward / SwiGLU — Where the Parameters Live

What it is. After attention mixes information across positions, the feed-forward network (FFN / MLP) transforms each position independently — it's a per-token nonlinear function. The classic GPT FFN is two linear layers with a GELU in between: up to a wider hidden dimension d_ff (usually 4·d_model), a nonlinearity, then down back to d_model.

Why a nonlinearity (GELU). Without a nonlinearity between linear layers, the composition is just another linear layer — no extra expressive power. GELU is the classic smooth gate (the lab's gelu uses the tanh approximation): roughly the identity for large positive inputs, squashing negatives toward zero, test_gelu_shape_of_curve checks the shape.

SwiGLU (the modern FFN). Modern models replace the single up-projection + GELU with a gated unit. Two parallel projections of the input — a gate (passed through SiLU) and an up — are multiplied element-wise, then projected down:

$$ \text{SwiGLU}(x) = \big(\underbrace{\text{SiLU}(x W_{\text{gate}})}{\text{gate}} \odot \underbrace{(x W{\text{up}})}{\text{value}}\big), W{\text{down}}, \qquad \text{SiLU}(z) = z\cdot\sigma(z). $$

The gate lets the network modulate each hidden unit multiplicatively — a learned, data-dependent "how much of this feature passes" — which empirically beats a plain GELU MLP at equal parameters. That is swiglu_ffn (with _silu) in the lab. (Because SwiGLU has three matrices instead of two, real models shrink d_ff to ~8/3·d_model to keep the parameter count matched.)

Where the parameters live. Per layer, attention is four d_model × d_model matrices (4·d²), while the FFN is three matrices of roughly d_model × d_ff with d_ff ≈ 4·d_model (~8·d² for a plain MLP, ~3 · 8/3 · d² = 8d² for SwiGLU). So the FFN holds roughly 2/3 of every layer's parameters — and thus ~2/3 of the whole model's. This is why Phase 00's 2N/token rule treats the MLP as the dominant term, and why Mixture-of-Experts puts its experts in the FFN: that's where the parameters (and the capacity) are.

   per transformer layer (rough):
     attention:  Wq Wk Wv Wo            ≈ 4·d²          ~1/3 of the layer
     FFN:        W_gate W_up W_down      ≈ 8·d²          ~2/3 of the layer
                                        ───────
                              the MLP is where the model "stores what it knows"

Production significance. Quantization (P06), pruning, and MoE all target the FFN because that's the parameter mass. "Where are the parameters in a transformer?" with the 2/3 answer is a frequent senior interview question.

Common misconception. "Attention is the expensive, parameter-heavy part." Attention dominates the conversation and the long-context compute (O(T²)), but the FFN dominates the parameter count and most of the FLOPs at typical sequence lengths.

Chapter 10: The Residual Stream & the Pre-Norm Block

What it is. A transformer block does not replace its input; it adds a correction to it. The running sum that every block reads from and writes back to is the residual stream — a [T, d_model] highway running the full depth of the model.

The pre-norm block. Each block is two sublayers, each wrapped in norm → sublayer → add:

$$ X \leftarrow X + \text{MHA}(\text{RMSNorm}(X)), \qquad X \leftarrow X + \text{FFN}(\text{RMSNorm}(X)). $$

That is transformer_block in the lab: add_matrices(X, attn) and add_matrices(X, ffn) are the residual adds — the same add_matrices you wrote in Chapter 1's helpers.

Why the residual stream matters. Two reasons:

1. Gradients flow. The +X means the gradient has a direct, unobstructed path back through every layer (the derivative of X + f(X) w.r.t. X includes the identity). This is what lets you train 80-layer models without vanishing gradients — the original ResNet insight, inherited by transformers.
2. Layers compose by accumulation. Each sublayer reads the current stream, computes a small correction, and adds it back. The stream is a shared workspace; interpretability research views it as a "communication channel" where heads write features that later heads read.

   residual stream  ──┬─────────────────────────┬──────────────►  (continues to next block)
                      │                          │
                  RMSNorm                    RMSNorm
                      │                          │
                     MHA                        FFN
                      │                          │
                      └──►(+)                    └──►(+)
                      add back to the stream     add back to the stream

Production significance. Pre-norm + residual is the reason deep transformers are trainable; the choice of pre- vs post-norm visibly affects training stability and is a real architecture decision. Phase 03 (autograd) will differentiate exactly this structure, and the clean gradient path is why it works.

Common misconception. "Normalization rescales the residual stream." In pre-norm, the norm is applied to a copy fed into the sublayer; the residual stream itself stays un-normalized (that's the point — it's a clean highway). The lab does exactly this: n1 = rms_norm(X) is consumed by attention, but the add is X + attn, not n1 + attn.

Chapter 11: The Full Forward Pass — Embedding to Logits

Putting it together. Now every piece composes into gpt_forward:

$$ \text{ids} \xrightarrow{\text{embed}} X \xrightarrow{\text{+RoPE}} X \xrightarrow{\text{N blocks}} X \xrightarrow{\text{RMSNorm}} X \xrightarrow{\text{LM head}} \text{logits } [T, \text{vocab}]. $$

Step by step, exactly as the lab does it:

   1. embedding lookup:   X = [embed[tid] for tid in token_ids]      [T, d_model]
   2. position:           X = rope(X, [0,1,...,T-1])                 inject order
   3. mask:               mask = causal_mask(T)                      no peeking ahead
   4. blocks:             for block in blocks: X = transformer_block(X, block, mask)
   5. final norm:         X = [rms_norm(row, final_norm_w) for row in X]
   6. LM head:            logits = linear(X, lm_head)                [T, vocab]

What the output means. Row t of the [T, vocab] logits is the model's score for every possible next token, given tokens 0..t. To generate, you softmax-and-sample row t to pick token t+1 (sampling is the next phase's territory). During training, you compare all rows to the true next tokens at once — that's the parallelism the causal mask buys you (Chapter 4).

Determinism. build_params(seed=…) draws every weight from one seeded random.Random, so the same seed gives the same logits — test_gpt_forward_deterministic checks that same-seed runs match and different seeds differ. Determinism is a lab requirement (LAB-STANDARD) and a debugging superpower: a forward pass that's reproducible is one you can bisect.

Production significance. This is model(input_ids).logits in any framework. nanoGPT's GPT.forward is this exact sequence in ~30 lines; an HF LlamaForCausalLM is this plus GQA, real batching, and fused kernels. You now know what every one of those lines is doing.

Common misconception. "The model outputs a token." It outputs logits — a score per vocab entry, per position. Turning logits into a token (greedy / temperature / top-k / top-p sampling) is a separate, deliberate step, and a frequent source of "the model is repetitive / incoherent" bugs.

Chapter 12: FlashAttention — The IO-Aware Fused Version

The problem with naive attention. The attention you built materializes the full [T, T] score matrix in memory (scores = Q·Kᵀ), softmaxes it, then multiplies by V. For long sequences that T² matrix is huge — and worse, on a GPU it gets written to and read from slow HBM (high-bandwidth memory) multiple times. From Phase 00's roofline: attention is memory-bandwidth-bound, and the bottleneck is moving that T² matrix, not the FLOPs.

The FlashAttention idea. Never write the full score matrix to HBM at all. Tile Q, K, V into blocks that fit in fast on-chip SRAM, and compute attention block by block, maintaining a running ("online") softmax — keeping a running max and running normalizer so you can fold each new block into the output without ever having seen the whole row at once. The math is identical to the attention you wrote; the memory traffic is what changes.

   naive:    Q·Kᵀ → [T,T] in HBM → softmax in HBM → ·V        (reads/writes the T² matrix)
   flash:    for each block of K,V:                            (T² never touches HBM)
               load Q-block, K-block, V-block into SRAM
               update running max m, running sum l, running output o   (online softmax)
             → exact same result, a fraction of the HBM traffic

Why it's a big deal. Same answer, but it turns attention from HBM-bound to compute-bound by slashing memory traffic — multiple-× faster and O(T) memory instead of O(T²), which is what makes long context affordable. It is a pure systems optimization: no approximation, no change to the model, just IO-awareness. This is the bridge from "I understand the attention algorithm" (this lab) to "I understand why production attention is fast" (Phase 00's roofline made concrete).

Production significance. FlashAttention (and v2/v3) is the default attention kernel in PyTorch SDPA, vLLM, and every serious serving stack. When torch.nn.functional.scaled_dot_product_attention runs fast, this is usually why. The lab's "build a block/tiled softmax" extension is a from-scratch taste of the core trick.

Common misconception. "FlashAttention is a faster approximation of attention." It is exact — the online-softmax trick reorders the computation so the result is mathematically identical. It trades a bit of recomputation for far less memory traffic; it does not drop any terms.

Lab Walkthrough Guidance

The lab (lab-01-attention-forward-pass) turns these chapters into code. Suggested order (matches the file top-to-bottom; helpers precede their users):

1. Linear algebra (matmul, transpose, add_vectors, add_matrices) — Chapter 1's helpers. add_matrices is the residual op. Get the validation (ragged/shape) right; everything builds on it.
2. softmax — Chapter 3. The whole lesson is the max-subtraction; -inf must map to 0.
3. Norms (layer_norm, rms_norm) — Chapter 8. RMSNorm has no mean and no bias.
4. linear — the affine map used by every projection.
5. causal_mask + scaled_dot_product_attention — Chapters 2–4. The 1/√d_k scale, the mask add, softmax per row, weighted sum of V. This is the core.
6. Multi-head (_split_heads, _concat_heads, multi_head_attention) — Chapter 5. A reshape plus per-head SDPA plus W_O.
7. RoPE (_rope_vector, rope) — Chapter 7. Rotate (even, odd) pairs; preserve the norm.
8. FFN (gelu, _silu, swiglu_ffn) — Chapter 9. Gate ⊙ up, then down.
9. Block + forward (transformer_block, build_params, gpt_forward) — Chapters 10–11.
10. KV-cache (attention_with_kv_cache) — Chapter 6. The soul test — test_kv_cache_equals_full_recompute must pass to tolerance.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked numbers.

Success Criteria

• All tests pass against your lab.py and against LAB_MODULE=solution.
• You can write softmax(Q·Kᵀ/√d_k + mask)·V from memory and explain every term — including the variance argument for √d_k and why softmax subtracts the max.
• You can explain the causal mask as additive -∞ → zero softmax weight → autoregression, and why a mask leak is a silent, catastrophic training bug.
• You can explain multi-head as parallel subspaces (a reshape, not extra params) and how GQA/MQA shrink the KV-cache (tying back to Phase 00's cache math).
• You can explain RoPE's rotation, why it preserves the norm, and why the QK score is relative.
• You can state where the parameters live (~2/3 FFN) and contrast RMSNorm with LayerNorm and pre- vs post-norm.
• You can prove the KV-cache equals the full causal recompute's last row, and explain why (the last query under the mask sees exactly the cached tokens).
• You can sketch why FlashAttention is faster (IO-aware, never materializes T²) and exact.

Interview Q&A

• "Walk me through scaled dot-product attention." softmax(Q·Kᵀ/√d_k + mask)·V: Q/K/V are learned projections of the input; the dot product scores query-key similarity; √d_k keeps the softmax unsaturated; softmax makes a distribution; ·V is the weighted average of value rows.
• "Why divide by √d_k?" The dot product over d_k dims has variance ≈ d_k; without the scale, wider heads produce large scores that saturate softmax to near one-hot and kill the gradient. /√d_k restores ~unit variance so training stays in the soft, high-gradient regime.
• "What exactly does the causal mask do?" Adds -∞ to scores for future positions; softmax(-∞)=0, so position i puts zero weight on any j>i. It makes a parallel training pass compute honest next-token prediction; a leak lets the model cheat and then fail at inference.
• "Why multiple heads?" To attend to several relationships in parallel subspaces (syntax, coreference, position) instead of averaging them into one. It's a reshape of the same parameter budget, not more parameters.
• "RoPE vs learned absolute positions?" RoPE rotates Q/K by pos·θ_i so the attention score depends on relative offset, preserves the norm, and extrapolates to longer contexts (NTK/YaRN are RoPE-frequency tweaks). Learned absolute embeddings can't exceed their trained max and encode absolute, not relative, position.
• "RMSNorm vs LayerNorm? Pre- vs post-norm?" RMSNorm drops mean-centering and bias — cheaper, equally good, the modern default. Pre-norm (normalize the sublayer input, keep a clean residual highway) is more stable for deep stacks than the original post-norm.
• "Where do the parameters live?" ~2/3 in the feed-forward (~8d² vs attention's ~4d² per layer). That's why the 2N/token rule treats the MLP as dominant and why MoE puts experts there.
• "Why is the KV-cache correct, and what does it save?" Under a causal mask the current query attends to exactly the past tokens, which is the cache — so incremental decode equals the full recompute's last row. It turns decode from O(T²) to O(T) by never recomputing past K/V.
• "How does GQA help, given decode is memory-bound?" Fewer K/V heads → smaller KV-cache → fewer bytes read per token → higher throughput on a bandwidth-bound workload, at negligible quality cost.
• "What is FlashAttention and is it an approximation?" An IO-aware, tiled, online-softmax attention that never writes the T² score matrix to HBM — same exact result, far less memory traffic, O(T) memory. Not an approximation.
• "Why does the model output logits, not a token?" It scores every vocab entry per position; turning a logit row into a token (greedy/temperature/top-k/top-p) is a separate sampling step and a common source of generation-quality bugs.
• "Why apply normalization at all?" Deep stacks let activation magnitudes drift, exploding or vanishing gradients; the norm rescales each position to a controlled magnitude so training stays stable.

References

• Vaswani et al., Attention Is All You Need (2017) — the transformer, scaled dot-product and multi-head attention, the √d_k scale.
• Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021) — RoPE.
• Zhang & Sennrich, Root Mean Square Layer Normalization (2019) — RMSNorm.
• Shazeer, GLU Variants Improve Transformer (2020) — SwiGLU; and Fast Transformer Decoding: One Write-Head is All You Need (2019) — MQA.
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023).
• Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022), and FlashAttention-2 (2023).
• Karpathy, nanoGPT — the cleanest readable GPT forward pass; and Let's build GPT from scratch.
• Alammar, The Illustrated Transformer — the canonical visual walkthrough of Q/K/V and multi-head.
• PyTorch docs — torch.nn.functional.scaled_dot_product_attention, torch.nn.MultiheadAttention, torch.nn.RMSNorm.

Hitchhiker's Guide — The Transformer From Scratch

The compressed practitioner tour. If WARMUP.md is the professor, this is the senior who leans over and says "here's what you actually need to remember to build and debug a transformer."

The 30-second mental model

A transformer is an embedding lookup, then N identical blocks, then a norm + LM head. A block is two residual sublayers: X += Attn(norm(X)) and X += FFN(norm(X)). Attention is softmax(Q·Kᵀ/√d_k + mask)·V — content-based weighted retrieval; the causal mask (-∞ above the diagonal) makes it autoregressive. Multi-head runs that in h parallel subspaces. RoPE injects relative position by rotating Q/K. RMSNorm keeps magnitudes sane. The FFN holds ~2/3 of the parameters. At decode time the KV-cache stores past K/V so you never recompute — and it's provably identical to the full recompute because the last query, under the mask, sees exactly the cache. That's the whole machine.

The numbers to tattoo on your arm

Back-of-envelope one-liners

# "Head dimension for d_model=4096, 32 heads?"
d_k = 4096 / 32 = 128   → scale attention by 1/sqrt(128) ≈ 0.088

# "Where are the params in a layer (d=4096, SwiGLU d_ff≈11008)?"
attn  = 4 * 4096^2          ≈ 67M
ffn   = 3 * 4096 * 11008    ≈ 135M   → FFN ≈ 2/3.  Always.

# "Is my causal mask right? Token 0's attention output should equal V[0]."
SDPA(Q,K,V, causal_mask(T))[0] == V[0]   # row 0 sees only itself

# "RoPE sanity: position 0 is identity, and the norm never changes."
rope(v, 0) == v ;  ||rope(v, p)|| == ||v||  for all p

The framework one-liners (where these live in real tools)

import torch.nn.functional as F
# scaled dot-product attention, fused (FlashAttention under the hood):
out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)   # the mask, for free
torch.nn.MultiheadAttention(embed_dim, num_heads)               # the whole MHA module
torch.nn.RMSNorm(d_model)                                       # RMSNorm (HF: LlamaRMSNorm)
# RoPE: HF LlamaRotaryEmbedding + apply_rotary_pos_emb(q, k, cos, sin)
# KV-cache: model.generate(..., use_cache=True); past_key_values / DynamicCache
# config knobs that change the math: num_attention_heads, num_key_value_heads (GQA),
#   rope_theta, rms_norm_eps, intermediate_size (d_ff)
# nanoGPT model.py is this entire lab in ~150 readable lines.

War stories

• The mask that leaked the future. Off-by-one in the causal mask let each token see one future position. Training loss dropped beautifully (the model was peeking at the answer); inference was gibberish (the future it leaned on wasn't there). Symptom: great train loss, broken generation. Always test out[0] == V[0] — token 0 must see only itself.
• The RoPE base mismatch. Fine-tuned a rope_theta=1e6 long-context base model but loaded it with the default 10000. Short prompts were fine; anything past a few thousand tokens degraded into nonsense because the positional angles no longer matched what the weights expected. Position config must match the checkpoint exactly — RoPE base and max_position_embeddings.
• The forgotten √d_k. Someone "simplified" attention by dropping the scale. With d_k=128 the scores blew up, softmax went one-hot, gradients died, and the model wouldn't learn. The fix was one / math.sqrt(d_k).
• The KV-cache that didn't match. A custom decode loop diverged from the prefill path after a few tokens — they'd applied RoPE to the cached K again on each step. The cache stores already-rotated keys; rotate once. Golden test: cached decode must equal full-recompute's last row to tolerance.
• The pre-norm/post-norm swap. Ported a model assuming post-norm; the deep stack wouldn't train (activations blew up). Modern models are pre-norm — normalize the sublayer input, keep the residual stream clean.

Vocabulary (rapid-fire)

• Q / K / V — query/key/value, three learned linear projections of the input.
• SDPA — scaled dot-product attention, softmax(QKᵀ/√d_k + mask)·V.
• Causal mask — -∞ above the diagonal; makes attention autoregressive.
• Head / d_k — a subspace of width d_model/n_heads; the per-head key dimension.
• MQA / GQA — share K/V across query heads (1 / a few) to shrink the KV-cache.
• KV-cache — stored past K/V so decode is O(T); the inference memory wall (Phase 00).
• RoPE — rotary position embedding; rotates Q/K so the score is relative.
• RMSNorm — mean-free, bias-free normalization; the modern default.
• SwiGLU — gated FFN down(silu(gate) ⊙ up); where ~2/3 of params live.
• Residual stream — the running +X highway every block reads and writes.
• Pre-norm — normalize the sublayer input; stable for deep stacks.
• FlashAttention — IO-aware tiled attention; exact, never materializes T².
• Logits — the per-vocab scores; sampling turns a row into the next token.

Beginner mistakes

• Forgetting the 1/√d_k scale (softmax saturates, training dies).
• An off-by-one causal mask (future leak: low train loss, broken inference).
• Applying RoPE to the residual stream or re-rotating cached keys instead of rotating Q/K once.
• Thinking "the model outputs a token" — it outputs logits; sampling is a separate step.
• Quoting attention as the parameter-heavy part — it's the FFN (~2/3).
• Using == on floats in tests instead of a tolerance (softmax/RoPE are floating-point).
• Confusing GQA's win (memory / bandwidth) with a FLOP speedup.
• Treating FlashAttention as an approximation — it's exact, just IO-aware.

The one thing to take away

A transformer is embed → N×(residual attention + residual FFN) → norm → logits, and attention is softmax(Q·Kᵀ/√d_k + mask)·V. If you can write those two lines from memory, explain the causal mask and √d_k, and prove the KV-cache equals the full recompute, you can build, read, and debug any modern decoder model. Everything else in the curriculum is a change to this machine.

Brother Talk — The Transformer From Scratch

Off the record. The stuff I'd tell you over a beer, not in a design review.

This is the phase that separates "I use transformers" from "I understand them," and the gap is enormous. Almost everyone in this field can call model.generate(). A shocking number of people with "ML" or "AI" in their title cannot write attention on a whiteboard, cannot tell you what √d_k is for, and have never once thought about what the causal mask is actually doing. The day you can build a GPT forward pass from nothing — embeddings, attention, the mask, RoPE, the KV-cache — you stop being a user of the magic and become someone who can fix it. That's the whole game.

Write it by hand, even though it's "just calling PyTorch" in real life. I know it feels pointless to implement matmul with a triple loop when torch has a fused kernel. Do it anyway. The reason seniors can debug a broken model at 2 a.m. is that they once built the thing with their bare hands and know in their body what each line does. When the loss is suspiciously low, the senior thinks "mask leak" before the junior has finished re-reading the stack trace. You earn that reflex here, once, and keep it forever.

The causal mask will save your job someday. I'm serious. At some point you'll see a model with a gorgeous training curve that generates garbage in production, and a room full of people will be baffled because "the loss went down, it must be working." You'll ask "are we sure the mask isn't leaking the future?" — and you'll be right, because a model that can peek at the answer during training learns nothing useful for inference. Tattoo this: low train loss + broken generation = suspect the mask. That one diagnosis is a story you'll tell for years.

The KV-cache equality is the most beautiful thing in this lab — don't skip the soul test. It's easy to treat the cache as a performance hack and move on. But sit with why it's correct: the last query, under a causal mask, attends to exactly the tokens already in the cache. That's not an approximation or an engineering shortcut — it's an identity. Once you really feel that, the entire serving stack (PagedAttention, continuous batching, all of P09) makes sense, because it's all built on that one guaranteed-equal trick. Make test_kv_cache_equals_full_recompute pass and understand it cold.

Don't get hypnotized by attention. Attention gets all the press — the paper is literally called "Attention Is All You Need." But ~2/3 of the parameters live in the boring feed-forward, and most of the FLOPs at normal context lengths are there too. The model stores what it knows in the MLP and moves information around with attention. Knowing that the FFN is where the mass is will make you right in a dozen conversations about quantization, MoE, and pruning where everyone else is staring at the attention.

You will be tempted to memorize the equations. Don't — derive them. Anyone can recite softmax(QKᵀ/√d_k)·V. The senior can tell you why the √d_k is there (the dot-product variance grows with dimension and saturates softmax) and why RoPE encodes relative position (rotation, and the dot product of two rotated vectors depends only on the angle difference). Interviewers don't want the formula; they want to watch you reconstruct it. Practice explaining it out loud to a wall.

What's actually worth caring about: the attention equation and what each term means, the causal mask, why √d_k, where the parameters live (~2/3 FFN), and the KV-cache identity. What's not worth losing sleep over: memorizing the exact GELU constant, whether your toy uses pre- or post-norm, the precise SwiGLU d_ff ratio. Get the load-bearing ideas into your bones; look the constants up.

The career framing. This phase is where the curriculum stops being about numbers around the model (Phase 00) and becomes about the model itself. Everything ahead — autograd differentiates this forward pass, fine-tuning patches these linear layers, quantization shrinks these weights, serving caches these keys and values — is a modification of the machine you build right here. Own this one cold and every later phase is "a change to a thing I already understand," not new magic. That's the difference between learning ten disconnected tools and understanding one system deeply. Now go make pytest green — and don't skip the soul test.

Lab 01 — Attention, RoPE & a Tiny GPT Forward Pass

Phase: 02 — The Transformer From Scratch Difficulty: ⭐⭐⭐☆☆ (the algebra is small; the mechanism is the whole point) Time: 4–6 hours

Every engineer says "I use transformers." Far fewer can build one with no framework — and that gap is exactly what a senior interview probes. This lab makes you write the mechanism by hand: scaled dot-product attention, the causal mask that makes it autoregressive, multi-head attention, RoPE rotary positions, RMSNorm/LayerNorm, a SwiGLU feed-forward, a pre-norm transformer block, the full GPT forward pass (embedding → blocks → norm → LM head → logits), and — the soul of the lab — the KV-cache, proven numerically identical to a full causal recompute for the last position. Pure stdlib lists and math; no numpy, no torch; deterministic from a seed. When generation is wrong, slow, or leaks the future, this is the muscle memory that lets you fix it.

What you build

• matmul / transpose / add_vectors / add_matrices — the linear algebra a transformer is built from. add_matrices is the residual-stream op X + sublayer(X); we write the matmul triple loop so the m·n·k MAC cost is visible.
• softmax — the numerically stable, max-subtracting softmax that turns scores into a distribution that sums to 1 and maps a masked -inf to exactly zero weight.
• layer_norm / rms_norm — the two normalizations: LayerNorm (mean + variance + bias) and the cheaper, mean-free RMSNorm that Llama/Mistral/Gemma use.
• linear — the affine projection y = x·W (+ b) used everywhere (Q/K/V/O, the MLP, the LM head).
• causal_mask — the [T, T] additive mask (0 on/below the diagonal, -inf above) that is autoregression: token i may attend to j only if j <= i.
• scaled_dot_product_attention — softmax(Q·Kᵀ / sqrt(d_k) + mask) · V: the core mechanism, including the 1/sqrt(d_k) scale that keeps the softmax from saturating.
• multi_head_attention (with _split_heads / _concat_heads) — project, split into heads, attend per head in parallel subspaces, concat, output-project with Wo.
• rope (with _rope_vector) — rotary position embedding: rotate each (even, odd) dim pair by pos·θ_i; preserves the norm and makes the attention dot product depend on relative offset.
• gelu / swiglu_ffn (with _silu) — the activations and the gated down(silu(x·W_gate) ⊙ (x·W_up)) feed-forward where ~2/3 of the parameters live.
• transformer_block — one pre-norm block: X += MHA(norm1(X)); X += FFN(norm2(X)).
• build_params / gpt_forward — deterministically build a tiny GPT from a seed, then run the full forward pass: token ids → [T, vocab] logits.
• attention_with_kv_cache — one incremental decode step against a key/value cache, proven equal to the full causal recompute's last row — the entire justification for the cache.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why softmax subtracts the max and why test_softmax_numerically_stable_large_logits would NaN without it.
• You can explain why a masked (-inf) score contributes exactly zero probability, and why that makes test_causal_attention_zero_in_upper_triangle give out[0] == V[0].
• You can explain what the 1/sqrt(d_k) factor does and why test_sqrt_dk_scaling_is_present detects it (the unscaled distribution is sharper).
• You can explain why RoPE leaves the vector norm unchanged (test_rope_preserves_norm) and why position 0 is the identity (test_rope_position_dependent).
• You can explain why test_kv_cache_equals_full_recompute is the soul of the lab: the last query under a causal mask attends to exactly tokens 0..t, which is exactly the cache, so the incremental result must equal the full recompute's last row.
• You can state where the parameters live (~2/3 in the FFN) and why GQA/MQA shrink the KV-cache.

How this maps to the real stack

Limits of the miniature (be honest in the interview): we apply RoPE once on the residual stream for simplicity, whereas real models rotate Q and K inside each attention layer (the rotation must touch the scores, not the residual); we use plain multi-head, not GQA/MQA (the lever that actually shrinks the KV-cache in production); attention here is the naive O(T²) materialize-the- scores form, not the IO-aware fused FlashAttention that never writes the full score matrix to HBM; there is no training, no dropout, no batching, and tiny dims; and our matmul is a Python triple loop, not a tensor-core GEMM. The algorithm is faithful; the systems engineering is the rest of the curriculum (P06 quantization, P09 serving/PagedAttention, P17 kernels).

Extensions (build these on real hardware)

• Move RoPE inside the attention layer: rotate Q and K per head before the dot product, and confirm the QK score depends only on the relative offset (pos_q − pos_k).
• Add GQA: give K/V fewer heads than Q (share K/V across query-head groups) and measure the KV-cache shrink (ties back to Phase 00's KV-cache math).
• Implement a block (tiled) softmax that computes attention without ever materializing the full [T, T] score matrix — the core idea of FlashAttention's online softmax — and check it matches scaled_dot_product_attention to tolerance.
• Add a greedy / temperature / top-k sampler on top of gpt_forward and generate a sequence token-by-token using attention_with_kv_cache; verify it equals the prefill-then-decode path.
• Port one block to PyTorch and assert your gpt_forward logits match torch.nn.functional.scaled_dot_product_attention + nn.RMSNorm to a tight tolerance.

Interview / resume

• Talking points: "Why divide attention scores by √d_k?" "How does the causal mask make a transformer autoregressive — what exactly is -inf doing?" "Why does RoPE encode relative position, and why is that better than learned absolute embeddings?" "Where do most of the parameters live, and why is the FFN ~2/3?" "Prove the KV-cache is correct — why is it identical to recomputing?"
• Resume bullet: Implemented a transformer from scratch in pure Python — scaled dot-product attention with the causal mask, multi-head attention, RoPE, RMSNorm, a SwiGLU feed-forward, the full GPT forward pass, and an incremental KV-cache proven numerically identical to the full causal recompute — mapping each piece to its PyTorch / nanoGPT / FlashAttention equivalent.

Phase 03 — Autograd & Training Dynamics

The phase where the model stops being a fixed function and starts to learn. Phase 02 built the forward pass — attention, RoPE, a GPT that turns tokens into logits. This phase builds the machinery that makes those weights correct: reverse-mode automatic differentiation (the thing behind loss.backward()), the optimizers that turn gradients into updates, and the training-loop dynamics — learning-rate schedules, gradient clipping, numerical precision — that decide whether a run converges, plateaus, or blows up at 2 a.m. on step 4000.

Why this phase exists

Ask a candidate "what does loss.backward() do?" and you separate the people who use PyTorch from the people who understand it. The honest answer is small and beautiful: the framework records every operation as a node in a computation graph, each node knows its own local derivative, and one reverse pass over the graph applies the chain rule to produce every gradient at once. That is the entire idea of backpropagation, and it is the engine under every model you will ever train. Everything else in training — Adam, weight decay, warmup, clipping, mixed precision — is a refinement on top of "compute the gradient, take a step."

You cannot reason about a training run you do not understand. When the loss spikes, when it won't go down, when fp16 produces NaN, when Adam diverges but SGD doesn't — the engineer who built the autograd tape and the optimizer by hand knows where to look. The five things to own here:

1. Reverse-mode autodiff — why one backward pass yields all gradients, and how the topological order + grad accumulation make it work (and why forward-mode is the opposite tradeoff).
2. The softmax + cross-entropy gradient — the loss that trains every classifier and LM, and why it collapses to the clean p − y.
3. Optimizers: SGD → momentum → Adam → AdamW — the moment estimates, bias correction, and why decoupled weight decay (the "W") was a real fix, not a rebrand.
4. Training dynamics — warmup + cosine schedules, gradient clipping, loss spikes: the knobs that decide convergence.
5. Numerical precision — fp32 / fp16 / bf16, mixed precision, loss scaling, gradient accumulation: the difference between a run that fits and one that OOMs or NaNs.

Concept map

                         ┌──────────────────────────────────────────┐
                         │  Forward pass (Phase 02): tokens → logits │
                         └──────────────────────────────────────────┘
                                            │  loss = cross_entropy(logits, target)
                                            ▼
                    ┌───────────────────────────────────────────────┐
                    │     COMPUTATION GRAPH  (every op = a node)      │
                    │   each node knows its LOCAL derivative          │
                    └───────────────────────────────────────────────┘
                                            │  .backward()
                          reverse topological order + chain rule
                                            ▼
                    ┌───────────────────────────────────────────────┐
                    │   GRADIENTS for every parameter (one pass)      │
                    └───────────────────────────────────────────────┘
                          │              │                │
                  clip_grad_norm   lr_schedule       OPTIMIZER
                  (tame spikes)    (warmup+cosine)   SGD→Adam→AdamW
                          └──────────────┴────────────────┘
                                            ▼
                              θ ← θ − lr · update(grad)   ← one step
                                            │
                            repeat → loss strictly decreases
                                            │
                 (precision: fp32 / fp16+loss-scaling / bf16 ; grad accumulation)

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py.

Anatomy of one training step

Every training loop in existence — from this lab to a 70B run on a thousand GPUs — is the same six moves. Burn this into muscle memory; the rest of the phase is just how each move works:

1. zero_grad        clear last step's gradients (.grad accumulates by design)
2. forward          inputs → logits → loss        (builds the computation graph)
3. backward         loss.backward()               (reverse-mode: one pass, all grads)
4. clip             clip_grad_norm(params, τ)      (cap step length, keep direction)
5. step             optimizer.step()               (SGD / AdamW turns grad → update)
6. schedule         lr_scheduler.step()            (warmup + cosine: how big the step)

Miss step 1 and gradients pile up (the classic silent bug). Get step 3's order wrong and you read half-finished gradients. Skip step 4 on a bad batch and the loss spikes. The lab implements all six.

Integrated scenario ideas

• Explain loss.backward() on a whiteboard: draw the graph for loss = CE(W·x, y), label each node's local derivative, and walk the reverse pass that fills W.grad. This is the autograd interview question.
• Debug a loss spike: a run is fine until step 4000, then loss explodes. Enumerate causes (LR too high after warmup, a bad batch, fp16 overflow, no grad clipping) and the fix order.
• Defend AdamW over Adam: show on a weight-decayed objective why coupling L2 into the gradient (Adam) interacts badly with the adaptive denominator, and why decoupling (AdamW) fixes it.
• Size a precision decision: when do you reach for bf16 vs fp16+loss-scaling, and what does gradient accumulation buy you when the batch won't fit?

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive the gradient of z = x*y + x by hand and trace it through .backward().
• You can explain why reverse-mode gives all gradients in one pass and when forward-mode wins.
• You can derive the softmax + cross-entropy gradient and state why it equals p − y.
• You can explain SGD vs Adam vs AdamW, the moment estimates, bias correction, and decoupled weight decay — without notes.
• You can explain warmup, cosine decay, gradient clipping, and the precision/loss-scaling story.
• You can tie every piece of the Value engine to its torch equivalent.

Key takeaways

• One backward pass, all gradients. Reverse-mode autodiff is the whole trick: record the graph, give each op a local derivative, walk it backward applying the chain rule. Gradients accumulate, which is exactly why you zero_grad() every step.
• softmax + cross-entropy is p − y. The most common loss in the field has the cleanest gradient in the field — probabilities minus the one-hot target. Know the derivation cold.
• AdamW ≠ Adam + L2. Decoupled weight decay is a genuine correction; it's the default optimizer for transformers for a reason. Momentum, adaptive per-parameter steps, bias correction — own each.
• Training dynamics are engineering, not luck. Warmup avoids early divergence, cosine anneals to a minimum, clipping survives the giant gradient, and precision/loss-scaling/accumulation decide whether the run fits the hardware at all.
• This is PyTorch. torch.Tensor.requires_grad + loss.backward() is the Value engine you built, vectorized onto tensors and a GPU. The algorithm is identical; only the data type changed.

Next: Phase 04 — Vision-Language Models & Multimodal Fusion.

Warmup Guide — Autograd & Training Dynamics

Zero-to-senior primer for Phase 03. We start from "what is a gradient and why does a model need one" and end at the full training loop: the computation graph, reverse-mode automatic differentiation (the engine behind loss.backward()), the softmax+cross-entropy (p − y) gradient, the optimizer family SGD → momentum → Adam → AdamW, learning-rate schedules, gradient clipping and loss spikes, numerical precision and mixed precision, and gradient accumulation. Then we build all of it as a tiny scalar autograd engine and watch a real net learn XOR. When you finish, loss.backward() is no longer a black box — it is the thing you wrote.

Table of Contents

• Chapter 1: Why Gradients — Learning as Descent
• Chapter 2: The Computation Graph
• Chapter 3: The Chain Rule and Reverse-Mode Autodiff
• Chapter 4: Backprop on the Graph — One Pass, All Gradients
• Chapter 5: Softmax + Cross-Entropy and the Clean (p − y) Gradient
• Chapter 6: Optimizers — SGD → Momentum → Adam → AdamW
• Chapter 7: Learning-Rate Schedules — Warmup + Cosine
• Chapter 8: Gradient Clipping and Loss Spikes
• Chapter 9: Numerical Precision, Mixed Precision, Gradient Accumulation
• Chapter 10: This Is PyTorch Autograd
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Why Gradients — Learning as Descent

From zero. A neural network is a function f(x; θ) with a pile of tunable numbers θ (the parameters). "Training" means: pick θ so the model's outputs are good. We make "good" precise with a loss L(θ) — a single number that is small when the model is right and large when it's wrong. Training is then a minimization problem: find θ that makes L small.

The one tool we have. θ has billions of dimensions, so we can't search it. But we can ask a local question: if I nudge each parameter a tiny bit, does the loss go up or down, and how fast? That is exactly the gradient ∇L, the vector of partial derivatives ∂L/∂θ_i. The gradient points in the direction of steepest increase, so to decrease the loss we step the other way:

$$ \theta \leftarrow \theta - \eta \cdot \nabla L(\theta). $$

This is gradient descent, and η (the learning rate) is how big a step we take. The entire field of training is (a) computing ∇L efficiently and (b) choosing a good step from it. Chapters 2–4 are (a); chapters 6–8 are (b).

Why this framing matters. Once you see training as "compute the gradient, take a step, repeat," everything else slots in: backprop is how we compute the gradient cheaply; Adam is a smarter step; the LR schedule is how the step size changes over time; clipping is what to do when the gradient is huge. Hold this picture and nothing in the phase is mysterious.

Common misconception. "The model figures out the weights itself." No — gradient descent moves the weights, and backprop supplies the gradients it descends on. There is no magic; there is calculus and a loop. The senior move is to know the loop well enough to debug it.

Chapter 2: The Computation Graph

What it is. Any expression you compute is a directed graph: leaves are inputs/parameters, internal nodes are operations, and edges point from an operation to the values it produced. Take

$$ L = (w \cdot x + b - y)^2. $$

As a graph: w and x feed a *; that and b feed a +; subtract y; square the result. Each node holds a value (the forward result) and remembers which nodes it came from.

   w ──┐
       ├─(*)── u ──┐
   x ──┘           ├─(+)── p ──┐
   b ──────────────┘           ├─(−)── e ──(square)── L
   y ──────────────────────────┘

Why build the graph at all? Because the graph is the program, and the derivative of a program is computed by walking that same graph backward (Chapter 3). In our lab, every arithmetic operation on a Value creates a new node that stores its parents (_prev) and a closure (_backward) describing its local derivative. The graph is built as a side effect of the forward pass — you just compute L normally, and the tape records itself.

Two flavors. Frameworks build this graph either statically (define the whole graph once, then run it many times — TensorFlow 1.x, JAX's jit) or dynamically (rebuild it every forward pass, so control flow is plain Python — PyTorch eager, and our Value engine). Dynamic graphs are what make PyTorch feel like normal code; the price is rebuilding the tape each step.

Common misconception. "The graph stores the math symbolically and differentiates it like a CAS (computer algebra system)." It does not. Autodiff is not symbolic differentiation and not finite differences. It stores numeric local derivatives at each node and composes them — exact (to floating point) and cheap. That distinction is a favorite interview probe.

Chapter 3: The Chain Rule and Reverse-Mode Autodiff

The chain rule. If L depends on u and u depends on w, then

$$ \frac{\partial L}{\partial w} = \frac{\partial L}{\partial u} \cdot \frac{\partial u}{\partial w}. $$

In words: the sensitivity of the loss to w is the sensitivity to u times the local sensitivity of u to w. Every node only needs to know its local derivative (∂u/∂w for the op it performed); the chain rule stitches the locals into the global gradient.

Two ways to apply it. Suppose the computation goes inputs → ... → loss.

• Forward-mode computes, for one chosen input, how every downstream value changes — it pushes derivatives forward alongside the values. To get gradients w.r.t. all n parameters you'd run it n times. Cheap when you have few inputs, many outputs.
• Reverse-mode computes, for one chosen output (the loss, a scalar!), how it changes w.r.t. every input — it pulls derivatives backward from the output. One pass gives all n gradients. Cheap when you have many inputs, one output — which is exactly a neural network: billions of parameters, one scalar loss.

This is the whole reason backprop exists. We have N parameters and 1 loss, so reverse-mode is the right algorithm: a single backward pass yields ∂L/∂θ_i for every parameter, at roughly the cost of one extra forward pass. (Recall the 6ND from Phase 00 — the "× 3" is forward + the ~2× of this backward pass.) Forward-mode would cost N passes. That asymmetry is why every framework does reverse-mode and why this is the single most important algorithm in the curriculum.

The local derivatives you need. For each op, the local rule the node stores:

Each row is one _backward closure in the lab. That's the entire engine.

Common misconception. "Reverse-mode is always better." It's better for our shape (many in, one out). If you had one input and many outputs (rare in ML, common in sensitivity analysis), forward-mode wins. Knowing why you pick reverse-mode — the input/output count — is the senior answer, not "because PyTorch does."

Chapter 4: Backprop on the Graph — One Pass, All Gradients

The algorithm. Backpropagation is reverse-mode autodiff run on the computation graph:

1. Forward pass — compute every node's value, building the graph (each node remembers its parents and its _backward closure).
2. Seed the output node's gradient to 1.0 — because ∂L/∂L = 1.
3. Reverse topological order — process nodes so that a node is handled only after every node that uses it. Call each node's _backward, which uses the chain rule to add this node's contribution into its parents' .grad.

Why topological order, exactly? A node's gradient is the sum over all paths from it to the loss. If node x feeds two later nodes, x.grad isn't complete until both of those have pushed their contributions. Processing in reverse-topo order guarantees that by the time we read a node's grad to push it upstream, every downstream consumer has already added to it. Get the order wrong and you read a half-finished gradient.

forward:   build graph, fill .data           ──►
                                              loss
backward:  seed loss.grad = 1.0              ◄──
           reverse-topo: each node does
           parent.grad += local_deriv * node.grad

Gradients accumulate — this matters twice. Inside one backward pass, a shared node (the diamond: x used in two places) must sum its incoming gradients, so every _backward uses +=, never =. And across training steps, grads keep accumulating in .grad until you reset them — which is precisely why PyTorch makes you call optimizer.zero_grad() (or set grads to 0.0 at the top of the loop, as our train does) every step. Forget it and you sum this step's gradient onto last step's: the classic silent training bug.

A concrete hand-check (the lab's first test). For z = x*y + x at x=2, y=3: dz/dx = y + 1 = 4 (one path through the *, one through the bare + x), dz/dy = x = 2. Run .backward() and you should get exactly x.grad == 4, y.grad == 2. If you can derive that by hand and the engine agrees, your chain rule is right.

The soul test: finite differences. The independent ground truth is the definition of a derivative:

$$ \frac{\partial L}{\partial x} \approx \frac{L(x+h) - L(x-h)}{2h} \quad (\text{central difference}). $$

It's slow (O(params) re-evaluations) and a little inexact, but it needs no calculus — so it's the perfect referee. If your analytic backprop matches central finite differences on a gnarly composite like tanh(x)² · e^x + log(x+3), your engine is correct. This is the test that proves an autograd implementation; every framework's test suite has it (torch.autograd.gradcheck).

Common misconception. "Backprop is its own algorithm separate from autodiff." Backprop is reverse-mode autodiff specialized to neural-net graphs. The 1986 "backprop" paper popularized it for nets; the underlying reverse-mode idea is older and fully general.

Chapter 5: Softmax + Cross-Entropy and the Clean (p − y) Gradient

Turning logits into probabilities. A classifier/LM outputs raw scores (logits) z. Softmax maps them to a probability distribution:

$$ p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}. $$

The stability trick (non-negotiable). e^{1000} overflows to inf. Subtract the max logit m = max_j z_j first:

$$ p_i = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}. $$

This is mathematically identical (the e^{-m} cancels top and bottom) but now every exponent is ≤ 0, so nothing overflows. Our softmax does this, and a test feeds it [1000, 1001, 1002] to prove it. Never write a softmax without the max-subtraction — it's a guaranteed interview ding.

The loss. Cross-entropy measures how far the predicted distribution p is from the true one-hot target y (a 1 at the correct class, 0 elsewhere):

$$ \text{CE} = -\sum_i y_i \log p_i = -\log p_{\text{target}}. $$

For a one-hot target it collapses to "negative log-probability of the right answer." Build it on the stable softmax so log never receives 0.

The beautiful gradient. Differentiate CE through softmax w.r.t. the logits and almost everything cancels:

$$ \boxed{;\frac{\partial \text{CE}}{\partial z_i} = p_i - y_i;} $$

— predicted probability minus the one-hot target. If the model is confident and right, p ≈ y and the gradient is near zero (nothing to fix); if it's confidently wrong, the gradient is large. This (p − y) form is why the softmax+CE pair is the default classification loss everywhere: the backward is a single subtraction. Our lab verifies it numerically — the analytic (p − y) must match finite differences, and the gradient on the target logit is p_target − 1.

Why frameworks fuse them. F.cross_entropy in PyTorch takes raw logits, not probabilities, and computes softmax + log + the loss together using log-sum-exp — both for stability and to use the clean fused gradient directly. If you ever softmax-then-CE manually, you've reimplemented it less safely.

Common misconception. "Cross-entropy needs probabilities as input." The math does, but the production op takes logits and does the softmax internally — feeding it already-softmaxed values double-softmaxes and silently wrecks training. State this; interviewers love it.

Chapter 6: Optimizers — SGD → Momentum → Adam → AdamW

We have the gradient. How do we turn it into a step? This is a short evolutionary ladder; know each rung and what the next one fixed.

1. SGD. The literal descent rule:

$$ \theta \leftarrow \theta - \eta , g, \qquad g = \nabla L. $$

Simple, memory-free, and with a good schedule it generalizes beautifully (it's still used to train many vision models). Weakness: one global η for all parameters, and it zig-zags in ravines.

2. Momentum. Accumulate an exponentially-weighted average of past gradients (a "velocity") and step along it:

$$ v \leftarrow \mu v + g, \qquad \theta \leftarrow \theta - \eta v. $$

This damps oscillation across a ravine and accelerates along its floor — like a heavy ball rolling downhill. μ ≈ 0.9 is typical.

3. Adam (Adaptive Moment estimation). Keep two running averages: the 1st moment m (mean of gradients, like momentum) and the 2nd moment v (mean of squared gradients, a per-parameter scale):

$$ m \leftarrow \beta_1 m + (1-\beta_1) g, \qquad v \leftarrow \beta_2 v + (1-\beta_2) g^2. $$

Because m, v start at 0 they're biased toward zero early, so bias-correct:

$$ \hat m = \frac{m}{1-\beta_1^{,t}}, \qquad \hat v = \frac{v}{1-\beta_2^{,t}}, $$

then take a per-parameter step scaled by the inverse root of the 2nd moment:

$$ \theta \leftarrow \theta - \eta , \frac{\hat m}{\sqrt{\hat v} + \epsilon}. $$

The √v̂ denominator gives each parameter its own effective learning rate — big-gradient parameters get smaller steps, small-gradient ones get larger — which is why Adam "just works" on transformers with almost no tuning. t is the 1-based step count (for bias correction); defaults β = (0.9, 0.999), ε = 1e-8.

4. AdamW — decoupled weight decay. Weight decay shrinks weights toward zero to regularize. The "obvious" way (plain Adam with L2) adds λθ to the gradient — but then Adam's adaptive √v̂ denominator rescales that decay differently per parameter, so the regularization you asked for isn't the regularization you get. AdamW decouples it: do the Adam step and, separately, shrink the weight directly:

$$ \theta \leftarrow \theta - \eta \lambda \theta \quad (\text{decoupled decay}), \qquad \theta \leftarrow \theta - \eta \frac{\hat m}{\sqrt{\hat v}+\epsilon}. $$

This small change measurably improves generalization and is the default optimizer for training transformers. Our adamw_step implements exactly this — and a test proves that with zero gradient the weight still shrinks (decay acting alone).

Common misconception. "AdamW is just Adam with L2 weight decay." No — the whole point is that L2-in-the-gradient and decoupled decay are not equivalent under an adaptive optimizer. The "W" is a genuine algorithmic fix, which is exactly why interviewers ask the difference.

Chapter 7: Learning-Rate Schedules — Warmup + Cosine

Why the LR shouldn't be constant. Early in training the weights are random and gradients are wild; a big step can send the loss to NaN. Late in training you want small steps to settle into a minimum rather than bounce around it. So the canonical schedule has two phases:

• Linear warmup — ramp η from ~0 up to base_lr over the first few hundred/thousand steps. This is especially important for Adam, whose 2nd-moment estimate is noisy and unreliable at the very start (t small) — a warmup hides that fragile early period.
• Cosine decay — after warmup, anneal η from base_lr down toward 0 following a half-cosine:

$$ \eta(t) = \eta_{\max} \cdot \tfrac12\Big(1 + \cos\big(\pi \cdot \text{progress}\big)\Big), \quad \text{progress} \in [0,1]. $$

 lr
  base ┤        ____
       │     __/    \__
       │    /          \__
       │   / warmup        \___ cosine → ~0
       └──┴──────────────────────► step

Cosine is smooth (no abrupt drops to destabilize Adam's moments) and spends a lot of time at low LR near the end, which empirically finds better minima. Our lr_schedule implements both branches; the tests check the warmup ramp, the peak at base_lr, and that it decays to ~0.

Common misconception. "Warmup is a minor trick." Skip warmup on a large transformer with Adam and you can diverge in the first few hundred steps. It's load-bearing, not decorative.

Chapter 8: Gradient Clipping and Loss Spikes

The problem. Occasionally a batch produces a huge gradient — a hard example, a numerical edge, a rare token. One oversized step can launch the weights into a bad region and spike the loss (you watch it shoot from 2.0 to 9.0 on a dashboard). On recurrent/transformer training this is common enough to need a standing defense.

Global-norm clipping. Treat all parameter gradients as one long vector, compute its L2 norm, and if it exceeds a threshold, scale the whole vector down so its norm equals the threshold:

$$ \text{norm} = \sqrt{\textstyle\sum_i g_i^2}, \qquad \text{if norm} > \tau:; g_i \leftarrow g_i \cdot \frac{\tau}{\text{norm}}. $$

Scaling the whole vector preserves its direction (so you still descend correctly) but caps its length (so no single step is catastrophic). Our clip_grad_norm does exactly this and returns the pre-clip norm — which is itself a useful signal: log it, and a sudden spike in the gradient norm is your early warning that a loss spike is coming.

The debugging ladder for a loss spike. When loss explodes mid-run, check, in order: (1) is the LR too high after warmup? (2) is gradient clipping on, and is the pre-clip norm spiking? (3) is this fp16 overflow producing inf/NaN (Chapter 9)? (4) is it a single pathological batch? The engineer who logs gradient norm and LR finds the cause in minutes.

Common misconception. "Clipping changes the optimum / hurts the model." Global-norm clipping preserves gradient direction; it only bounds step length on the rare giant-gradient step. It's a stability tool, not a regularizer that distorts where you converge.

Chapter 9: Numerical Precision, Mixed Precision, Gradient Accumulation

This chapter is conceptual (the lab is fp32 stdlib floats), but the ideas are senior table-stakes.

The precision menu.

Why range matters more than precision for training. fp16's 5 exponent bits mean small gradients underflow to 0 (they vanish) and large activations overflow to inf. bf16 keeps fp32's exponent, so it almost never overflows/underflows — you trade mantissa precision (usually fine, since gradients are noisy anyway) for range. This is why modern accelerators train in bf16 and why it "just works" where fp16 needs babysitting.

Mixed precision. Do the heavy matmuls in fp16/bf16 (fast, half the bandwidth — recall Phase 00's roofline: decode is bandwidth-bound) but keep a master copy of the weights and the optimizer state in fp32, and accumulate sensitive reductions in fp32. You get most of the speed/memory win without the accuracy loss.

Loss scaling (the fp16 fix). Because fp16 gradients underflow to zero, multiply the loss by a big constant S before backward (which scales every gradient by S, lifting them out of the underflow zone), then divide the gradients by S before the optimizer step. Dynamic loss scaling auto-tunes S: raise it when things are fine, halve it when an inf/NaN appears. bf16's range usually makes loss scaling unnecessary — another reason it's preferred.

Gradient accumulation. Want an effective batch of 256 but only 32 fits in memory? Run 8 micro-batches, summing (accumulating) their gradients without stepping, then take one optimizer step and zero the grads. Mathematically it's a 256-batch update; in memory it's a 32-batch footprint. This is the same grad-accumulation behavior (+=) you implemented in backward — now exploited deliberately. The cost is time (more forward/backward passes per step), not memory.

Common misconception. "fp16 and bf16 are interchangeable 16-bit formats." They are not — bf16 trades precision for range and is far more forgiving for training; fp16 needs loss scaling to avoid gradient underflow. Naming which one and why is the senior answer.

Chapter 10: This Is PyTorch Autograd

Everything you built maps one-to-one onto the real framework — same algorithm, just vectorized.

The one thing to internalize: a torch.Tensor with requires_grad=True is a Value whose .data is a multi-dimensional array instead of one float, and whose _backward operates on whole tensors instead of scalars. The reason GPUs help is not a different algorithm — it's that one node is now a matmul over millions of numbers that the hardware does in parallel. You built the real thing; production just made the nodes bigger.

Lab Walkthrough Guidance

The lab (lab-01-autograd-engine) turns this guide into code. Suggested order (matches the file):

1. Value core ops — __add__, __mul__, __pow__, exp/log/tanh/relu (Ch 2–3). For each, set out.data and an out._backward closure that does parent.grad += local · out.grad. Use += everywhere.
2. Derived ops — __neg__/__sub__/__truediv__ and the r-variants, all in terms of the core ops (no new _backward needed).
3. backward() (Ch 4) — build the reverse topological order with a visited set, seed self.grad = 1.0, walk reversed. Hand-check z = x*y + x first, then make the finite-difference soul test pass — that single green test means your chain rule is correct.
4. softmax / cross_entropy (Ch 5) — subtract the max; build CE on the stable softmax; verify the gradient is p − y.
5. Neuron/Layer/MLP — deterministic init from random.Random(seed); the final layer is linear (logits). Check the parameter count.
6. clip_grad_norm, sgd_step, adamw_step, lr_schedule (Ch 6–8) — clip returns the pre-clip norm; AdamW does decoupled decay and the bias-corrected Adam step.
7. train (the soul training test) — zero grads, sum CE over the 4 XOR examples, one .backward(), optional clip, optimizer step, record loss. It must strictly decrease and be deterministic for a fixed seed.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked numbers.

Success Criteria

• All tests pass against your lab.py and against LAB_MODULE=solution (31 tests).
• You can derive dz/dx, dz/dy for z = x*y + x and match .backward().
• The finite-difference test is green — your analytic gradients match central differences on a gnarly composite (the proof your engine is correct).
• You can explain why a reused node accumulates its gradient (and why that forces zero_grad).
• You can write a stable softmax (max-subtraction) and derive the (p − y) CE gradient from memory.
• You can state, without notes, the SGD → momentum → Adam → AdamW ladder and what decoupled weight decay fixes.
• Your XOR training loss is strictly decreasing and deterministic for a fixed seed.

Interview Q&A

• "What does loss.backward() actually do?" — Walks the computation graph in reverse topological order, applying the chain rule: each node adds its local-derivative × upstream-grad into its parents' .grad. One pass fills every parameter's gradient because the loss is a single scalar (reverse-mode).
• "Why does reverse-mode give all gradients in one pass, and when is forward-mode better?" — Reverse-mode is O(1) passes for many-inputs-one-output (nets: N params, 1 loss). Forward-mode is O(inputs) but O(1) in outputs, so it wins when inputs are few and outputs many (the opposite shape).
• "Derive the softmax + cross-entropy gradient." — p_i = softmax(z)_i, CE = −log p_target; ∂CE/∂z_i = p_i − y_i (probabilities minus one-hot). Near-zero when confident-and-right, large when confident-and-wrong.
• "Why subtract the max in softmax?" — e^{large} overflows; subtracting max makes every exponent ≤ 0 with identical math (the factor cancels). Without it, large logits give inf/NaN.
• "SGD vs Adam vs AdamW?" — SGD: one global LR. Adam: per-parameter LR via the 2nd-moment estimate, plus momentum and bias correction. AdamW: Adam with decoupled weight decay (shrink the weight directly rather than adding L2 to the gradient, which the adaptive denominator would distort). AdamW is the transformer default.
• "Why bias-correct in Adam?" — m, v start at 0, so early estimates are biased toward zero; dividing by 1 − β^t corrects it, which matters most in the first steps (and is why warmup helps).
• "Why warmup + cosine?" — Warmup avoids early divergence while Adam's moment estimates are noisy; cosine anneals the step smoothly so late training settles into a minimum.
• "Your loss spiked at step 4000 — what do you check?" — LR too high post-warmup; gradient norm spiking (clip it, log the pre-clip norm); fp16 overflow → NaN (loss scaling / bf16); a pathological batch. Global-norm clipping caps the giant step without changing direction.
• "fp16 vs bf16?" — Same size; bf16 keeps fp32's exponent range (rarely overflows/underflows, modern training default), fp16 has more mantissa but a tiny range so it needs loss scaling to stop gradient underflow.
• "What is gradient accumulation and why use it?" — Sum gradients over several micro-batches before one optimizer step to simulate a larger batch within a small memory budget; it's the same += accumulation, used deliberately. Costs time, not memory.
• "Why must you zero_grad() every step?" — .grad accumulates by design (so shared nodes and micro-batches sum correctly); without zeroing, this step's gradient piles onto the last step's.
• "How do you know your autograd is correct?" — gradcheck: compare analytic gradients against central finite differences. If they match on a nontrivial composite, the engine is right.

References

• Rumelhart, Hinton & Williams, Learning representations by back-propagating errors (1986) — the backprop paper.
• Baydin et al., Automatic Differentiation in Machine Learning: a Survey (2018) — forward vs reverse mode, the definitive overview.
• Andrej Karpathy, micrograd (github.com/karpathy/micrograd) and Neural Networks: Zero to Hero — the spiritual source of this lab; watch "The spelled-out intro to neural networks and backpropagation."
• Kingma & Ba, Adam: A Method for Stochastic Optimization (2015).
• Loshchilov & Hutter, Decoupled Weight Decay Regularization ("AdamW", 2019).
• Micikevicius et al., Mixed Precision Training (2018) — fp16, master weights, loss scaling.
• PyTorch docs — Autograd mechanics, torch.optim.AdamW, torch.nn.utils.clip_grad_norm_, torch.cuda.amp (automatic mixed precision), and torch.autograd.gradcheck.
• Goodfellow, Bengio & Courville, Deep Learning (2016), ch. 6.5 (backprop) and ch. 8 (optimization).

Hitchhiker's Guide — Autograd & Training Dynamics

The compressed practitioner tour. If WARMUP.md is the professor, this is the senior who leans over and says "here's what you actually need to remember about loss.backward() and the loop."

The 30-second mental model

A model's program is a computation graph; every op is a node that knows its local derivative. loss.backward() walks that graph in reverse topological order applying the chain rule, so one pass fills every gradient (reverse-mode wins because nets are many-params-in, one-loss-out). Gradients accumulate (+=) — that's why you zero_grad() every step. Then an optimizer turns grads into a step: SGD → momentum → Adam → AdamW (decoupled weight decay). Wrap it with **warmup

• cosine** LR, clip the global grad norm to survive spikes, and pick a precision (bf16 ≫ fp16 for training because of range). softmax+CE has the cleanest gradient in ML: p − y.

The numbers to tattoo on your arm

Back-of-envelope one-liners

# "Does my autograd work?" → gradcheck against central finite differences
analytic == (L(x+h) - L(x-h)) / (2h)   within ~1e-4   → correct

# "Why is one backward pass enough?" → reverse-mode, scalar loss
N params, 1 loss  → reverse-mode: O(1) passes for all N grads (forward-mode: O(N))

# "Loss spiked at step 4000" → the ladder
LR too high post-warmup? grad-norm spiking (clip!)? fp16 overflow→NaN (bf16/loss-scale)? bad batch?

# "Batch won't fit" → gradient accumulation
want batch 256, fit 32 → 8 micro-batches, accumulate grads, ONE step  (time cost, not memory)

# softmax+CE gradient (no derivation needed at the whiteboard)
d(CE)/d(logit_i) = softmax(logit)_i - onehot_i

The framework one-liners (where these live in real tools)

import torch, torch.nn.functional as F
x = torch.tensor(2.0, requires_grad=True)
y = (x*x + x); y.backward(); x.grad          # = 5.0  (your Value engine, vectorized)

loss = F.cross_entropy(logits, target)        # logits in! softmax+log-sum-exp fused → (p-y) grad
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)  # decoupled WD
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total)    # or HF warmup+cosine

opt.zero_grad(set_to_none=True)               # grads accumulate — clear them every step
loss.backward()                               # reverse-mode autodiff
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)        # returns pre-clip norm
opt.step(); sched.step()

with torch.autocast("cuda", dtype=torch.bfloat16):  # mixed precision; bf16 ⇒ no loss scaling
    loss = model(batch)
torch.autograd.gradcheck(fn, inputs)          # the finite-difference referee

War stories

• The forgotten zero_grad. Loss looked like it was learning then chaotically diverged. Cause: no optimizer.zero_grad() — gradients accumulated across steps, so each update was the sum of all prior gradients. One line fixed a week of confusion. Grads accumulate by design.
• The 2 a.m. NaN. fp16 run, loss → NaN around step 1500. Small gradients had underflowed to zero and a big activation overflowed to inf. Fix: enable dynamic loss scaling (or just switch to bf16, whose fp32-range exponent dodged the whole problem).
• The AdamW-vs-Adam regression. Team "added weight decay" by setting Adam's weight_decay and saw worse generalization than expected. The L2-in-the-gradient was being rescaled by Adam's adaptive denominator. Switching to AdamW (decoupled) recovered the expected regularization.
• The loss spike that clipping ate. Long training, periodic loss spikes from rare hard batches. Logging the pre-clip gradient norm showed the spikes coming; turning on global-norm clipping at 1.0 flattened them without changing the final loss — direction preserved, length capped.
• The double-softmax. Someone fed softmax(logits) into F.cross_entropy (which softmaxes again). Training "worked" but underperformed for weeks. CE takes logits, not probabilities.

Vocabulary (rapid-fire)

• Computation graph — nodes = ops, edges = data dependencies; differentiated by walking backward.
• Reverse-mode autodiff / backprop — one backward pass → all gradients (the right mode for nets).
• Local derivative — an op's own ∂out/∂in; the chain rule composes them.
• Topological order — process a node only after every consumer; lets grads finish before use.
• Grad accumulation — += into .grad: required within a pass (shared nodes), exploited across micro-batches, and the reason zero_grad exists.
• (p − y) — the softmax+cross-entropy gradient w.r.t. logits.
• Moments (Adam) — 1st = mean of grads (momentum), 2nd = mean of grad² (per-param scale).
• Decoupled weight decay (AdamW) — shrink the weight directly, not via the gradient.
• Bias correction — 1/(1−β^t) fixup for zero-initialized moments.
• Warmup / cosine — ramp the LR up, then anneal it down smoothly.
• Global-norm clipping — cap ‖grad‖₂ to a threshold; keeps direction, bounds step length.
• fp16 / bf16 / loss scaling / AMP — 16-bit training; bf16 has the range, fp16 needs loss scaling.
• gradcheck — verify autograd against central finite differences.

Beginner mistakes

• Forgetting optimizer.zero_grad() (grads accumulate across steps).
• Writing softmax without max-subtraction (overflow on large logits).
• Feeding probabilities to cross_entropy instead of logits (double softmax).
• Overwriting (=) instead of accumulating (+=) gradients in a custom backward (shared nodes break).
• Treating AdamW as "Adam + L2" — the decoupling is the point.
• Skipping warmup on a transformer with Adam (early divergence).
• Using fp16 without loss scaling and chasing the resulting NaNs (reach for bf16).
• Not logging the gradient norm — then having no idea why the loss spiked.

The one thing to take away

Before you trust any training run, be able to (1) explain loss.backward() as "reverse-topo chain rule, one pass, all grads," (2) verify a gradient against finite differences, and (3) name your optimizer/schedule/clip/precision and why. If you can, you can debug any loss curve. If you can't, you're praying to a black box.

Brother Talk — Autograd & Training Dynamics

Off the record. The stuff I'd tell you over a beer, not in a design review.

Build the autograd engine by hand once and loss.backward() stops being a god you pray to. I'm serious — this is the single most clarifying weekend in the whole track. Right now, somewhere in your head, there's a little altar where you put "and then the framework computes the gradients ✨." Burn it down. It's a graph, a reverse walk, and the chain rule — maybe 150 lines. The day you write it, every training bug you'll ever hit becomes debuggable instead of mystical. That shift is the difference between "I run training scripts" and "I own training."

The (p − y) gradient is a party trick that signals seniority. When someone asks you to derive the softmax+cross-entropy gradient and you just say "it's `p minus the one-hot, the derivation cancels almost everything" — calmly, like you've known it forever — the room recalibrates how senior you are. It's a two-minute thing to memorize and it pays off in every ML interview you'll ever take. Learn it cold.

Forgetting zero_grad() will get you, and then it'll get you again. Everyone does it. The loss does something deranged, you stare at the model architecture for an hour, and the bug is one missing line because gradients accumulate by design. Now that you've built the engine and felt the += yourself, you'll recognize the symptom instantly: "this looks like grads piling up." That recognition is worth a hundred Stack Overflow tabs.

AdamW vs Adam is a free interview point that most people fumble. They'll ask "what's the difference?" and the weak answer is "AdamW has weight decay." The right answer — "L2-in-the-gradient gets rescaled by Adam's adaptive denominator, so it's not the regularization you asked for; AdamW decouples the decay and applies it to the weight directly" — takes ten more seconds and marks you as someone who reads papers, not just docs. Decoupled. Say "decoupled."

bf16 is the answer to half the "my training NaN'd" questions, and almost nobody leads with range. People obsess over precision (mantissa bits) when the thing that actually kills fp16 training is range (exponent bits) — gradients underflow to zero, activations overflow to inf. bf16 keeps fp32's exponent, so it just... doesn't do that. Knowing it's about range, not precision is the kind of detail that makes a staff engineer nod.

Loss spikes are a vibe, and you can learn to read them. Once you've trained a few things you stop panicking at a spike and start running the ladder: LR too hot after warmup? grad-norm blowing up (clip it)? fp16 overflow? one cursed batch? Log the gradient norm from day one — it's your seismograph; you'll see the earthquake coming before the loss does. The people who look calm during a bad run aren't calm because they're geniuses; they're calm because they have a checklist.

The finite-difference check is the most honest friend you'll ever have. When you write any custom backward — a new layer, a fused op, a weird loss — don't trust your calculus. Nudge the input by a hair, see how the output moves, compare to your analytic gradient. If they match, you're right; if they don't, you have a bug, full stop, no debate. torch.autograd.gradcheck exists for exactly this and senior engineers reach for it reflexively. Cheap insurance against the worst kind of bug — the one where training "works" but slowly, quietly, wrongly.

Don't let the optimizer zoo intimidate you — it's one idea with three patches. SGD is "step downhill." Momentum is "keep some speed so you don't zig-zag." Adam is "give each parameter its own step size based on how noisy its gradient has been." AdamW is "and decay the weights separately so the adaptive part doesn't mangle your regularization." That's the whole story. Everyone treats this as arcane; it's four sentences. Once you can say those four sentences and mean them, you've got the part that actually matters in interviews and in practice.

What's actually worth caring about: the four reflexes — explain backward(), derive (p−y), gradcheck against finite differences, and name your optimizer/schedule/clip/precision with a reason. What's not worth losing sleep over: memorizing Adam's exact bias-correction algebra to five decimals, or hand-deriving tanh's derivative under pressure (you'll look it up; everyone does). Get the shape of the ideas; the constants live in the docs.

The career framing. This is the phase that makes you trustworthy with a training run — which is where the real money and the real 2 a.m. pages live. Anyone can fine-tune with a recipe; the person who can look at a diverging loss curve and say "that's an LR-after-warmup problem, here's the fix" is the person the team can't lose. The forward pass (Phase 02) made you understand the model. This phase makes you understand the making of the model. That's the senior jump. Now go make pytest green and watch your little XOR net actually learn — it's a genuinely good feeling.

Lab 01 — Reverse-Mode Autograd Engine that Trains a Tiny LM

Phase: 03 — Autograd & Training Dynamics Difficulty: ⭐⭐⭐☆☆ (the engine is ~150 lines; the insight is ⭐⭐⭐⭐⭐) Time: 4–5 hours

Every deep-learning framework is, at its core, the thing you build here: a graph of scalar operations, each of which knows its local derivative, plus one reverse pass that applies the chain rule to fill in every gradient at once. This lab is micrograd with teeth — a Value autograd engine, a numerically-stable softmax + cross-entropy with the clean (p − y) gradient, a tiny MLP, SGD and AdamW, global-norm gradient clipping, a warmup+cosine LR schedule, and a deterministic training loop that actually learns XOR with a loss that strictly decreases. When you finish, torch.Tensor.backward() is no longer magic — it's this, vectorized onto a GPU.

What you build

• class Value — a scalar node that records data, grad, its parents _prev, and a local _backward closure. Operators + * ** - / neg (and the r-variants) and methods exp / log / tanh / relu, each defining its own local derivative. .backward() builds the reverse topological order, seeds the output grad to 1.0, and runs every closure once.
• softmax(values) — max-subtracted, stable for huge logits, inside the graph so gradients flow.
• cross_entropy(logits, target_index) — −log(softmax(logits)[target]), built on the stable softmax so log never sees 0; its logit gradient is the famous p − onehot.
• Neuron / Layer / MLP(nin, nouts, seed) — a deterministic tiny net with hidden tanh layers and a linear output layer (so the outputs are usable logits), plus .parameters().
• sgd_step / adamw_step — the two optimizers, operating directly on Value.grad. AdamW carries 1st/2nd-moment state and applies decoupled weight decay.
• clip_grad_norm(params, max_norm) — global-norm clipping; returns the pre-clip norm.
• lr_schedule / train(...) — warmup+cosine LR, and a full-batch loop that learns XOR deterministically with a strictly-decreasing loss.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # against the reference (must be green)
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can derive dz/dx and dz/dy for z = x*y + x by hand and match .backward().
• test_grad_matches_finite_difference_* passes — the soul test: your analytic gradients match a central finite-difference estimate. If this is green, your chain rule is correct.
• You can explain why test_reused_node_accumulates_grad needs += (not =) in every _backward.
• test_softmax_stable_for_large_logits passes — you subtracted the max.
• test_cross_entropy_softmax_gradient_is_p_minus_onehot passes — you can state why the gradient is p − y without deriving it on the spot.
• test_training_loss_strictly_decreases passes — the soul training test: a real net, a real loss, learning a real (non-linearly-separable) task, monotonically.
• test_training_is_deterministic passes — same seed → byte-identical loss history (this is why _prev is an ordered tuple, not a set).

How this maps to the real stack

Limits of the miniature (say these in the interview): it's scalar autograd — real engines operate on tensors so one node is a whole matmul, not one number, which is the entire reason GPUs help; we keep the whole graph alive every step (no detach/checkpointing, no graph freeing), so memory grows with graph size; there is no mixed precision / loss scaling (fp32 only here); and the train loop is full-batch with no data loader, no minibatch noise, and no validation split.

Extensions (build these on real hardware)

• Add Tensor (a 2-D Value analog backed by array) so a "node" is a matmul and you can see why vectorization—not a different algorithm—is what GPUs buy you.
• Add momentum SGD and plain Adam (coupled L2) next to AdamW and show the difference decoupling makes on a weight-decayed quadratic.
• Add gradient accumulation: backward over k micro-batches before one optimizer step; show the loss curve matches a k×-larger batch.
• Add activation checkpointing: free intermediate Values and recompute them in backward; show the time/memory tradeoff.
• Port the exact same net to PyTorch and assert your gradients match loss.backward() to 1e-6.

Interview / resume

• Talking points: "Walk me through loss.backward() — what is it actually doing?" "Why does one backward pass give every gradient (reverse-mode), and when would forward-mode win?" "Derive the softmax+cross-entropy gradient." "SGD vs Adam vs AdamW — what does decoupled weight decay fix?" "Your loss spiked at step 4000 — what knobs?" (clip, warmup, lower LR, check fp16 overflow).
• Resume bullet: Implemented a reverse-mode automatic-differentiation engine from scratch (computation graph, topological backprop, the softmax+cross-entropy (p−y) gradient), with SGD and AdamW optimizers, global-norm gradient clipping, and a warmup+cosine schedule, and used it to train a tiny MLP to a strictly-decreasing, deterministic loss — i.e. rebuilt the core of PyTorch autograd.

Phase 04 — Vision-Language Models

The phase where the model stops being a language model and starts being a perception model. Up to now everything entered as token IDs. Now the input is a grid of pixels — and the senior question is mechanical, not magical: how do you turn an image into tokens an LLM can attend to, how do you teach vision and text to share one embedding space, and how do you splice the picture into the prompt without blowing the context budget? Answer that on a whiteboard and you can reason about GPT-4V, Claude's vision, LLaVA, and Flamingo as systems, not products.

Why this phase exists

Every multimodal model you'll be asked about in a senior interview is the same three pieces in a trench coat: a vision encoder (almost always a Vision Transformer), a contrastive alignment step (CLIP/SigLIP) that puts images and text in one space, and a fusion mechanism that hands the image to an LLM. The reason this deserves its own phase is that the interesting failures and costs all live in the seams:

1. An image is not free tokens. A 336×336 image at 14×14 patches is 576 tokens — before a single word of prompt. Multiply by the KV-cache math from Phase 00 and you see why high-resolution and multi-image inputs are the real cost driver of VLMs.
2. CLIP's whole trick is the loss. The encoders are ordinary transformers; what makes a cat-photo embedding land near the word "cat" is the symmetric InfoNCE objective over in-batch negatives and a temperature. Understand the loss and you understand zero-shot.
3. Fusion is an architecture decision. Projection (LLaVA: vision → MLP → LLM token space) is cheap and dominant; cross-attention (Flamingo) is more flexible but heavier. Knowing which one a model uses tells you its token-budget and training profile.
4. You train the connector, not the world. Production VLMs freeze the vision encoder, train the projector, and LoRA-tune the LLM. That "what's frozen, what's trained" answer is a senior tell — and it forward-references Phase 05.

Get fluent here and "the model can see now" stops being a slogan and becomes a token count, a loss, and a splice you can draw.

Concept map

                         ┌───────────────────────────────────┐
                         │   Pixels → tokens an LLM can read   │
                         └───────────────────────────────────┘
                            │              │             │
                ┌───────────┘       ┌──────┘      ┌──────┘
                ▼                    ▼             ▼
          VISION ENCODER      ALIGNMENT (CLIP)   FUSION
          ViT: image→patches  shared embedding   project (LLaVA)
          →patch embed→tokens  InfoNCE + τ        vs cross-attn (Flamingo)
          + [CLS] + pos        in-batch negatives  splice at <image>
                │                    │             │
                └────────┬───────────┴──────┬──────┘
                         ▼                   ▼
                  zero-shot / retrieval   image tokens inflate
                  (argmax cosine to text)  the KV-cache  ── ties to Phase 00
                         │                   │
                         └────────┬──────────┘
                                  ▼
                    VQA · OCR · grounding · captioning
                  (encoder frozen, projector + LoRA trained → Phase 05)

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py.

Integrated scenario ideas

• Size a VLM request: a prompt with two 768×768 images at 14×14 patches — count the image tokens, add them to the text, and use the Phase 00 KV-cache formula to show why the images, not the words, set the memory.
• Defend "freeze the encoder": explain why you train only the projector (+ LoRA on the LLM) on a small instruction set, and what would break if you fine-tuned the whole vision tower (catastrophic forgetting, cost, data hunger).
• Pick a fusion family: a model that must handle interleaved image/text documents vs one that answers a single-image VQA — argue projection vs cross-attention for each.
• Explain zero-shot to a new hire: why "no training" classification works because the classifier weights are text embeddings, and how prompt wording moves accuracy.
• OCR/grounding reality check: why a 224×224 patchified image cannot read small text, and how higher resolution (more patches) trades cost for legibility.
• Debug a contrastive run: the loss won't drop — walk through the suspects (desynced image/caption pairs so the diagonal is no longer the positive, a temperature that's collapsed, a too-small batch starving the negatives) using the similarity matrix as your microscope.
• Build a zero-shot classifier with no training data: turn a list of class names into "a photo of a {class}" prompts, embed them as the classifier weights, and argmax cosine — then show how prompt wording and template ensembling move accuracy.

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive the token count of an image from (H, W, patch_size) and explain the resolution↔token-count↔cost tradeoff.
• You can write the symmetric InfoNCE loss from the similarity matrix and explain why the diagonal is the positive and off-diagonals are in-batch negatives.
• You can explain temperature, why softmax/InfoNCE use max-subtraction, and why CLIP learns its temperature.
• You can contrast projection (LLaVA) vs cross-attention (Flamingo) fusion and say what's frozen vs trained in a real VLM.
• You can take any "the model can see" prompt and turn it into ViT → CLIP → projector → LLM + a token-budget number.

Key takeaways

• A Vision Transformer is a tokenizer for pixels. Patchify → linear embed → add position and [CLS] → feed an ordinary transformer. Once an image is tokens, the LLM machinery you already built (attention, KV-cache) applies unchanged.
• CLIP's power is the loss, not the encoders. The symmetric InfoNCE objective over in-batch negatives, scaled by a learned temperature, is what forges a shared image/text space — and that shared space is what makes zero-shot classification and retrieval free.
• Fusion is a budget decision. Projection (LLaVA) is the cheap, dominant default; cross-attention (Flamingo) buys flexibility at cost. Either way, images become tokens and tokens inflate the KV-cache — the Phase 00 wall returns, now driven by image resolution.
• Zero-shot is the shared space cashing out. Because images and text live in one space, you build a classifier from text prompts and pick the argmax cosine — no task-specific training, and you add a class by writing a prompt. Retrieval is the same operation transposed.
• You train the connector, freeze the rest. The senior mental model of "how is a VLM trained" is: vision encoder frozen, projector trained, LLM LoRA-tuned — which is exactly the PEFT story of the next phase.

Next: Phase 05 — PEFT: LoRA, Adapters & Distillation.

Warmup Guide — Vision-Language Models

Zero-to-senior primer for Phase 04. We start from "what does it even mean for a model to see an image" and end with the full pipeline that lets an LLM look at a photo and answer a question about it: the Vision Transformer that turns pixels into tokens, the CLIP contrastive objective that puts images and text in one embedding space, zero-shot classification and retrieval that fall out of that space for free, the two fusion families that hand the image to an LLM (projection à la LLaVA vs cross-attention à la Flamingo), and the systems reality — images inflate the token budget and the KV-cache, so the Phase 00 cost math comes straight back. Every term is built from first principles, and the lab makes each one runnable on nested lists.

Table of Contents

• Chapter 1: What "Seeing" Means for a Transformer
• Chapter 2: The Vision Transformer — Image → Patches → Tokens
• Chapter 3: CLIP — Forging a Shared Image/Text Space
• Chapter 4: The InfoNCE Loss, the Temperature, and In-Batch Negatives
• Chapter 5: Zero-Shot Classification & Retrieval
• Chapter 6: Fusion — Projection (LLaVA) vs Cross-Attention (Flamingo)
• Chapter 7: The Systems Reality — Tokens, KV-Cache, Resolution & Training
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: What "Seeing" Means for a Transformer

From zero. A transformer does not know what a pixel is. It only knows how to process a sequence of vectors — tokens — by mixing them with attention and transforming them with MLPs. Everything you learned in Phases 01–03 (tokenize, embed, attend, predict) assumes the input already is a sequence of vectors. So the entire problem of "making a model see" reduces to one question: how do you turn a 2D grid of pixels into a sequence of vectors that a transformer can chew on the same way it chews on words?

Why this framing matters. It collapses a scary-sounding topic ("multimodal AI") into a familiar one. There is no new kind of network. There is a front-end that converts pixels to tokens, and then it's transformers all the way down. Once you internalize that, a vision-language model stops being a black box and becomes three boxes you can name: encode the image to tokens, align those tokens with text, feed them to the LLM.

The three jobs. Pull any VLM apart and you find:

The senior's habit. When someone says "we'll add vision to the assistant," the senior does not picture a new model. They picture: which vision encoder (a ViT, probably CLIP-pretrained), how many tokens per image at our resolution, which fusion (projector, almost certainly), and what that does to our context length and KV-cache. Then they have an opinion with numbers in it.

Common misconception. "Multimodal models have some special architecture that understands images." No — the understanding is learned (Chapters 3–5), and the architecture is a vision encoder plus the same transformer you already know. The novelty is the front-end and the training objective, not a new flavor of neural network.

Chapter 2: The Vision Transformer — Image → Patches → Tokens

The problem it solves. An image of size H×W (say 224×224) has 50,176 pixels. You cannot make each pixel a token — attention is O(T²), so 50k tokens is 2.5 billion attention pairs per layer. You need far fewer, coarser tokens. The Vision Transformer's answer (Dosovitskiy et al., 2020) is gloriously blunt: cut the image into a grid of square patches and treat each patch as a token. The paper's title says it all — "An Image is Worth 16×16 Words."

Step 1 — patchify. Slice the image into non-overlapping P×P patches. An H×W image yields

$$ n_{\text{patches}} = \frac{H}{P} \cdot \frac{W}{P} $$

patches, each flattened to a vector of length P·P (for grayscale; P·P·3 for RGB). A 224×224 image with P = 16 gives 14·14 = 196 patches. This is the only place pixels enter; from here on it's vectors.

  H×W image, P×P patches            flatten each patch (row-major)
  ┌──┬──┬──┬──┐                      patch (0,0) = [p00 p01 p10 p11 ...]
  │P0│P1│P2│P3│   raster order →     patch (0,1) = [...]
  ├──┼──┼──┼──┤   left→right,        ...
  │P4│P5│P6│P7│   top→bottom         → n_patches vectors of length P·P
  └──┴──┴──┴──┘

Step 2 — patch embedding. Each flat patch vector is run through a single linear layer W of shape (P·P, d_model):

$$ \text{token}_i = \text{patch}_i \cdot W. $$

In real ViTs this is implemented as a Conv2d with kernel_size = stride = P — which is exactly a per-patch linear projection, just phrased as a strided convolution. After this step each patch is a d_model-dim token, indistinguishable in shape from a word embedding.

Step 3 — the [CLS] token and position embeddings. Two finishing touches make it a real ViT:

• [CLS] token. A single learned vector is prepended to the sequence (borrowed from BERT). After the transformer runs, the output at the [CLS] position is used as the whole-image embedding — the vector CLIP compares against text. (Modern variants sometimes mean-pool the patch tokens instead.)
• Position embeddings. Attention is permutation-invariant — it has no idea patch 5 is below patch 1. So a learned position vector is added to each token, encoding where the patch was in the grid. (Wrap the token in backticks in prose: the [CLS] token is special.)

Step 4 — it's just a transformer now. Feed the n_patches + 1 tokens into a standard transformer encoder (the attention + MLP blocks from Phase 02). Patches attend to each other; the [CLS] token aggregates the image into one vector.

Why this is the universal vision front-end. Almost every modern VLM uses a ViT (often a CLIP-pretrained one) as its eyes. The patch grid is the knob that controls everything downstream: more patches (smaller P or higher resolution) → more detail, more tokens, more cost.

Common misconception. "The ViT does some clever feature extraction like a CNN's edge detectors." Out of the box, the patch embed is a single linear layer — there's no hand-crafted vision prior. The ViT learns everything from data (which is why ViTs are data-hungry compared to CNNs, and why CLIP-pretraining on hundreds of millions of pairs is so valuable). The lab stops at the patch embed precisely to make that minimalism visible.

Chapter 3: CLIP — Forging a Shared Image/Text Space

The problem it solves. A ViT gives you an image embedding. A text encoder gives you a text embedding. But they live in different, unrelated coordinate systems — there's no reason the image-vector for a cat photo is anywhere near the text-vector for "cat." CLIP (Radford et al., 2021 — Contrastive Language-Image Pretraining) trains both encoders jointly so that a matching image/text pair lands close together and a non-matching pair lands far apart. The result is one shared embedding space where similarity means "these mean the same thing."

The setup. Take a giant dataset of (image, caption) pairs — CLIP used ~400M scraped from the web. For a batch of n pairs:

1. Encode all n images with the image encoder → n image embeddings.
2. Encode all n captions with the text encoder → n text embeddings.
3. L2-normalize every embedding to the unit sphere, so a dot product is the cosine similarity.
4. Compute the n×n similarity matrix S where S[i][j] = cos(image_i, text_j), scaled by a temperature.

Cosine similarity, precisely. For two vectors,

$$ \cos(a, b) = \frac{a \cdot b}{\lVert a\rVert , \lVert b\rVert} \in [-1, 1]. $$

It is 1 when they point the same way, -1 opposite, 0 orthogonal. The senior detail: guard the zero-norm case — if either vector is all zeros the angle is undefined, so return 0 rather than dividing by zero. (The lab tests this explicitly.)

            texts →
          t0   t1   t2   t3        the diagonal S[i][i] is the
        ┌────┬────┬────┬────┐      MATCHED pair (image i ↔ its caption).
   i0   │ ★  │    │    │    │      everything off-diagonal is a
   i1   │    │ ★  │    │    │      NON-match — a free "negative".
images  ├────┼────┼────┼────┤
   i2   │    │    │ ★  │    │      training pushes ★ up and the
   i3   │    │    │    │ ★  │      rest of each row/column down.
        └────┴────┴────┴────┘

Why "contrastive." You never tell the model the right answer in words. You only tell it which pairs go together (the diagonal) and which don't (everything else). Pulling positives together while pushing negatives apart is contrastive learning, and the magic is that the "everything else" — the in-batch negatives — comes for free: in a batch of 256, each image has 1 positive caption and 255 negatives, no extra labeling.

Common misconception. "CLIP is a classifier." It's not — it has no fixed label set. It's a similarity model. The classifier is something you build at inference time out of text (Chapter 5). That indirection is the whole reason it generalizes zero-shot.

Chapter 4: The InfoNCE Loss, the Temperature, and In-Batch Negatives

What we're optimizing. Given the n×n similarity matrix S, we want the diagonal to dominate every row and every column. CLIP frames this as two classification problems glued together:

• Image → text: treat row i as a classifier over the n texts; the correct class is i (the matched caption). Apply softmax + cross-entropy.
• Text → image: treat column j as a classifier over the n images; the correct class is j. Same loss, other direction.

The total is the symmetric InfoNCE (a.k.a. the CLIP loss), the average of the two:

$$ \mathcal{L} = \frac12\Big(\underbrace{\frac1n\sum_i \text{CE}(S_{i,:},, i)}{\text{image→text}} + \underbrace{\frac1n\sum_j \text{CE}(S{:,j},, j)}_{\text{text→image}}\Big). $$

For a single row with logits z and correct index t, cross-entropy is

$$ \text{CE}(z, t) = -\log \frac{e^{z_t}}{\sum_k e^{z_k}} = \log!\Big(\textstyle\sum_k e^{z_k}\Big) - z_t. $$

The numerical-stability rule (this is part of the lesson). e^{1000} overflows to infinity in floating point. So never exponentiate raw logits — subtract the row max first (log-sum-exp / max-subtraction):

$$ \log!\sum_k e^{z_k} = m + \log!\sum_k e^{,z_k - m}, \qquad m = \max_k z_k. $$

This is mathematically identical but every exponent is ≤ 0, so it never overflows. The lab's test_info_nce_numerically_stable_for_large_logits throws 1000.0 logits at you precisely to check you did this. (Same trick you used for softmax in Phase 08's sampler.)

The temperature τ. Before the softmax, CLIP divides the cosine similarities by a small temperature (or equivalently multiplies by a learned logit_scale):

$$ S[i][j] = \frac{\cos(\text{img}_i, \text{txt}_j)}{\tau}. $$

Cosine lives in [-1, 1] — far too flat for a sharp softmax. Dividing by τ = 0.07 rescales the gap to [-14, 14], so the softmax can actually separate the positive from the negatives. CLIP learns τ rather than fixing it (clamped to avoid collapse). Small τ → sharp, confident distribution; large τ → soft, forgiving one.

Why in-batch negatives + big batches matter. The loss's "push apart" signal comes entirely from the off-diagonal negatives in the batch. More negatives = a harder, more informative contrastive task, which is why CLIP trained with batch size 32,768. (SigLIP later swapped the softmax InfoNCE for a pairwise sigmoid loss specifically so it scales to huge batches without the softmax's global normalization — same idea, friendlier arithmetic.)

The soul of the lab. info_nce_loss is low when the diagonal dominates (matched pairs are the most similar) and high when you shuffle the text side so the positives fall off the diagonal. Those two tests — test_info_nce_low_when_pairs_matched and ...high_when_pairs_shuffled — are the lesson: a loss is just a number that's small when the model is right and big when it's wrong, and you can feel the contrastive objective by watching it move.

Common misconception. "Higher temperature is always smoother and safer." Temperature is a lever, and CLIP learns it because the right sharpness depends on the data; too small collapses training, too large gives a flat, uninformative gradient. It's tuned, not guessed.

Chapter 5: Zero-Shot Classification & Retrieval

The payoff. Once images and text share one space, you get zero-shot capabilities — solving tasks the model was never explicitly trained on — basically for free. This is the result that made CLIP famous: competitive ImageNet accuracy with zero ImageNet training images.

Zero-shot classification, mechanically. To classify an image into one of K classes:

1. Turn each class name into a text prompt — e.g. "a photo of a {class}" for cat, dog, car.
2. Encode each prompt with the text encoder → K text embeddings. These are your classifier weights, constructed on the fly from words.
3. Encode the image → one image embedding.
4. Predict argmax_k cos(image, text_k).

$$ \hat{y} = \arg\max_{k} ; \cos(\text{img}, \text{txt}_k). $$

That's clip_zero_shot_classify in the lab. No training, no fixed label set, no fine-tuning — change the prompts and you've changed the classifier. Want to add "airplane" as a class? Add a prompt. That flexibility is impossible with a softmax head trained on fixed labels.

Prompt engineering is real here. "a photo of a {class}" beats the bare class name, and prompt ensembling (averaging the embeddings of 80 templates per class) gave CLIP several extra points on ImageNet. The wording moves accuracy because it nudges where the text lands in the shared space.

Retrieval is the same operation, transposed. Instead of "which of these texts matches this image," ask "which of these images matches this text" (or vice versa) — encode everything, rank by cosine similarity, return the top-k. Image search by text query, and finding the caption for an image, are both just nearest-neighbor in the shared space (which connects forward to the vector search of Phase 11).

Common misconception. "Zero-shot means the model is guessing." No — it's doing genuine nearest-neighbor in a space that was trained to put meaning-equivalent things together. The "zero" refers to zero task-specific training examples, not zero knowledge. The knowledge is in the pretrained embedding space.

Chapter 6: Fusion — Projection (LLaVA) vs Cross-Attention (Flamingo)

The problem it solves. A CLIP-style encoder gives you image tokens (or one image embedding). An LLM expects word-embedding-shaped tokens in its embedding space. The gap between "vision encoder output space" and "LLM token space" is bridged by a connector, and there are two main families.

Family 1 — Projection (LLaVA). The simplest and now-dominant approach (Liu et al., 2023). Take the vision encoder's patch features, run each through a small MLP projector that maps vision_dim → llm_token_dim, and insert the resulting vectors directly into the LLM's input sequence as if they were ordinary tokens.

  image ──▶ ViT (frozen) ──▶ patch features (vision_dim)
                                   │
                                   ▼
                            PROJECTOR (2-layer MLP, GELU)   ← the part you train
                                   │
                                   ▼  image tokens (llm_token_dim)
  text:  "What is in <image> ?"  tokenize ──▶ [w0 w1 <IMG> w3 w4]
                                                       │
                            splice at the <image> slot ▼
                        [w0  w1  img0 img1 ... imgN  w3  w4]  ──▶  LLM

The LLM then attends over text and image tokens uniformly — it doesn't know which tokens came from pixels. That's Projector.project + assemble_multimodal_sequence in the lab: project the features, then splice them in at the <image> placeholder, giving a fused sequence of length len(text) - 1 + len(image). Cheap, simple, and shockingly effective.

Family 2 — Cross-attention (Flamingo). The earlier, heavier approach (Alayrac et al., 2022). Instead of stuffing image tokens into the text sequence, you interleave new cross-attention layers into the frozen LLM. The text tokens stay as the query stream; the image tokens become the keys and values that the text attends to. The image never enters the main sequence — the LLM reaches out to it through gated cross-attention blocks.

The senior takeaway. Projection won the popularity contest because it's simple, trains fast, and reuses the LLM unchanged — but it spends your context budget on image tokens. Cross- attention keeps the image out of the sequence (cheaper context) at the cost of more parameters and architectural surgery. Knowing which one a model uses tells you its cost profile and its multi-image story.

Common misconception. "The LLM has special image-handling logic." In the LLaVA family it absolutely does not — the projected image vectors are just tokens, and the (unmodified) LLM attends to them with the exact same mechanism it uses for words. That's why projection is so appealing: zero changes to the language model itself.

Chapter 7: The Systems Reality — Tokens, KV-Cache, Resolution & Training

Images are expensive tokens. This is where Phase 00 comes roaring back. Each image becomes (H/P)·(W/P) tokens before a word of prompt:

A single high-res image can be thousands of tokens. Recall the Phase 00 KV-cache formula — 2 · n_layers · T · d_model · bytes · batch — and note that T now includes every image token. The image, not the prompt, often dominates the context length and the KV-cache. This is why multi-image and high-resolution inputs are the real cost driver of VLMs, and why techniques like token pooling, perceiver resamplers, and tiling (split a big image, encode tiles, downsample) exist — they're all token-budget tricks.

The resolution ↔ token-count ↔ task tradeoff. Resolution is the dial that decides which tasks a VLM can even attempt:

• OCR / document reading needs high resolution — small text vanishes at 224×224. More patches = more legible glyphs = more tokens = more cost. This is the core tension in document VLMs.
• Grounding (pointing at where an object is — bounding boxes, "click the blue button") needs enough spatial tokens to localize. Too few patches and the model can't tell left from right.
• VQA / captioning (answering questions, describing) tolerate lower resolution because they're about what, not where or what-it-says.

So "what resolution?" is not a quality knob — it's a capability and cost decision. Doubling resolution roughly quadruples image tokens (it's H/P · W/P), which quadruples that part of the KV-cache and the attention cost.

How a VLM is actually trained (the senior tell). You do not train the whole thing end to end from scratch. The standard recipe (LLaVA-style):

1. Freeze the vision encoder — it's already a good CLIP-pretrained ViT; retraining it is expensive and risks wrecking the features.
2. Train the projector — the connector is the cheap, small piece that learns to translate vision features into LLM token space. Stage 1 trains only this on image-caption pairs.
3. Instruction-tune with the LLM via LoRA — Stage 2 unfreezes the LLM with LoRA (low-rank adapters) on visual-instruction data, while the encoder stays frozen and the projector keeps training.

That "encoder frozen, projector trained, LLM LoRA-tuned" answer is a strong senior signal — and the LoRA half is exactly Phase 05. The whole VLM is a small amount of new training on top of two big frozen models, which is why a lab can build a tiny working version on nested lists.

Common misconception. "Adding vision means retraining a giant model." For most teams it means training a projector (a few-million-parameter MLP) and LoRA adapters on top of frozen encoders and a frozen-ish LLM — orders of magnitude cheaper than pretraining, and entirely the point of doing it this way.

Lab Walkthrough Guidance

The lab (lab-01-clip-vit-projector) turns these chapters into code. Suggested order (matches the file top-to-bottom):

1. ViT front-end (patchify, patch_embed) — Chapter 2. patchify is careful indexing plus the divisibility check (ValueError); patch_embed is a matmul (patch · W) with an optional prepended [CLS]. Watch the raster order — left→right, top→bottom.
2. Similarity primitives (cosine_similarity, l2_normalize) — Chapter 3. The only traps are the zero-norm guard (return 0 / return zeros, never divide by zero) and clamping cosine to [-1, 1] to absorb float drift.
3. CLIP matrix + loss (clip_similarity_matrix, info_nce_loss) — Chapter 4. The matrix is scaled cosine (cos / temperature). The loss is symmetric (rows + columns) and must use max-subtraction — the large-logit test depends on it. The diagonal is the positive.
4. Zero-shot (clip_zero_shot_classify) — Chapter 5. Argmax cosine to the class text embeddings. One line of logic, but it's the whole "build the classifier from text" idea.
5. Connector (Projector, assemble_multimodal_sequence) — Chapter 6. The projector is a seeded 2-layer GELU MLP (vision_dim → hidden → llm_dim); determinism comes from random.Random(seed). The assembler splices image tokens at the placeholder index — mind the length: len(text) - 1 + len(image), and reject out-of-range indices.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked numbers. The two InfoNCE tests are the soul — make those green and you've felt contrastive learning.

Success Criteria

• All tests pass against your lab.py and against LAB_MODULE=solution.
• You can compute the token count of an image from (H, W, P) and explain why doubling resolution roughly quadruples image tokens — and what that does to the KV-cache.
• You can write the symmetric InfoNCE loss from a similarity matrix from memory, say why the diagonal is the positive, and explain the max-subtraction trick that keeps it stable.
• You can explain the temperature: why cosine is too flat for softmax, why dividing by τ fixes it, and why CLIP learns τ.
• You can describe zero-shot classification as "build the classifier from text prompts, then argmax cosine," and why that generalizes where a fixed softmax head can't.
• You can contrast projection (LLaVA) vs cross-attention (Flamingo) fusion on token budget, LLM changes, and multi-image flexibility.
• You can state the training recipe — encoder frozen, projector trained, LLM LoRA-tuned — and why it's cheap.

Interview Q&A

• "How does a transformer process an image at all?" — Patchify it into a grid of P×P patches, linearly embed each patch into a d_model token, add position embeddings and a [CLS] token, then run a standard transformer. An image becomes (H/P)·(W/P) tokens; after that it's ordinary attention.
• "What is CLIP and what does it actually learn?" — Two encoders (image + text) trained jointly with a contrastive loss so matching pairs have high cosine similarity in a shared space. It learns alignment, not classification — the classifier is built later from text.
• "Write the CLIP loss." — Symmetric InfoNCE: softmax cross-entropy over the rows (image→text) plus over the columns (text→image), diagonal = positive, averaged; logits are cos / temperature; use log-sum-exp for stability.
• "Why divide by a temperature?" — Cosine sits in [-1, 1], too flat for the softmax to separate positive from negatives. Dividing by τ≈0.07 sharpens it; CLIP learns τ (clamped) so the sharpness fits the data.
• "What are in-batch negatives and why use a huge batch?" — Every non-matching pair in the batch is a free negative; bigger batches give more negatives → a harder, more informative contrastive task. CLIP used batch 32k; SigLIP's sigmoid loss scales this even further.
• "Explain zero-shot classification." — Encode "a photo of a {class}" for each class to get classifier weights from text, encode the image, and pick the argmax cosine. No task-specific training; change the prompts to change the classes.
• "How does an LLM 'see' the image — what's the fusion?" — Two families: projection (LLaVA — project image features with an MLP and splice them into the token sequence; the LLM attends to them like words) and cross-attention (Flamingo — image tokens are K/V the LLM attends to via inserted cross-attn layers). Projection is simpler and spends context; cross-attn keeps the image out of the sequence.
• "What does adding an image do to your serving cost?" — Each image is hundreds–thousands of tokens, so it inflates context length, the KV-cache (2·L·T·d·bytes·batch), and the attention cost. The image often dominates the prompt — size your KV-cache for it.
• "What resolution should a VLM use?" — A capability decision, not a quality knob. OCR and grounding need high resolution (small text/precise location need many patches); VQA/captioning tolerate less. Doubling resolution ~4×'s the image tokens and cost.
• "How is a VLM trained — what's frozen?" — Freeze the CLIP-pretrained vision encoder, train the projector on image-caption pairs, then instruction-tune with LoRA on the LLM. You train the connector, not the world (forward-ref Phase 05).
• "Projection vs cross-attention — when would you pick each?" — Projection for single/few-image, simplest training, reuse the LLM unchanged. Cross-attention for many interleaved images or when you can't afford image tokens eating the context budget.
• "What's the difference between CLIP and SigLIP?" — Same idea (contrastive image/text alignment), different loss: CLIP uses softmax InfoNCE (needs the full batch for normalization); SigLIP uses a pairwise sigmoid loss that's cheaper to scale to enormous batches.

References

• Dosovitskiy et al., An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale (ViT, 2020) — patches as tokens, the [CLS] token, position embeddings.
• Radford et al., Learning Transferable Visual Models From Natural Language Supervision (CLIP, 2021) — contrastive image/text pretraining, the symmetric InfoNCE loss, zero-shot.
• Liu et al., Visual Instruction Tuning (LLaVA, 2023) and Improved Baselines with Visual Instruction Tuning (LLaVA-1.5) — the MLP projector connector and the two-stage training recipe.
• Alayrac et al., Flamingo: a Visual Language Model for Few-Shot Learning (2022) — gated cross-attention fusion and interleaved image/text.
• Zhai et al., Sigmoid Loss for Language Image Pre-Training (SigLIP, 2023) — the sigmoid alternative to softmax InfoNCE that scales to larger batches.
• van den Oord et al., Representation Learning with Contrastive Predictive Coding (2018) — the original InfoNCE objective CLIP's loss descends from.
• HF transformers docs — ViTModel, CLIPModel, LlavaForConditionalGeneration — to verify the miniature against the real classes; OpenCLIP for the contrastive-loss and zero-shot code.

Hitchhiker's Guide — Vision-Language Models

The compressed practitioner tour. If WARMUP.md is the professor, this is the senior who leans over and says "here's what you actually need to remember about how a model sees."

The 30-second mental model

A VLM is three machines: a ViT that cuts an image into patches and embeds each patch as a token; a CLIP-style contrastive objective that trains image and text encoders to share one embedding space (matching pairs → high cosine); and a connector that hands image tokens to an LLM — either projection (LLaVA: MLP the features into LLM token space and splice them into the sequence) or cross-attention (Flamingo: image is K/V the LLM attends to). Zero-shot classification is "build the classifier from text prompts, argmax cosine." And the catch: images are hundreds-to-thousands of tokens, so they inflate the KV-cache — the Phase 00 wall, back again.

The numbers to tattoo on your arm

Back-of-envelope one-liners

# "How many tokens is a 336×336 image at patch 14?"
(336/14) * (336/14) = 24 * 24 = 576 tokens   (before any text)

# "I send 4 images at 768×768 — what's my context cost?"
4 * (768/14)^2 ≈ 4 * 3009 ≈ 12,000 image tokens → that's your KV-cache driver

# "KV-cache of those image tokens?" (Phase 00 formula)
2 * n_layers * T_image * d_model * bytes * batch   ← T now dominated by the image

# "Why is the temperature 0.07?"
cos ∈ [-1,1] is too flat for softmax; /0.07 stretches the gap to [-14,14] → sharp

The framework one-liners (where these live in real tools)

# Encode an image / text with CLIP (embeddings already L2-normalized for cosine)
from transformers import CLIPModel, CLIPProcessor
img_emb = model.get_image_features(**proc(images=img, return_tensors="pt"))
txt_emb = model.get_text_features(**proc(text=["a photo of a cat"], return_tensors="pt"))
# Zero-shot = argmax cosine (== dot product on normalized embeddings)
logits = (img_emb @ txt_emb.T) * model.logit_scale.exp()   # the learned temperature

# The ViT patch embed IS a strided conv
nn.Conv2d(in_ch, d_model, kernel_size=patch, stride=patch)   # one token per patch

# LLaVA fusion: project vision features, splice at the <image> token
image_features = mm_projector(vision_tower(pixel_values))     # vision_dim → llm_dim
# inputs_embeds = [text_embeds[:idx], image_features, text_embeds[idx+1:]]

War stories

• The context blew up on the first screenshot. Demo worked on tiny test images; a user uploaded a 4K screenshot and the request 10×'d the token count and OOM'd. Nobody had multiplied resolution into the token budget. Image tokens = (H/P)²; size context and KV-cache for the image, not the prompt.
• The "it can't read the receipt" ticket. VQA worked great; OCR was garbage. The encoder ran at 224×224 — small text literally didn't survive patchification. Fix was higher resolution / tiling, which cost 4× the tokens. Resolution is a capability decision, not a quality slider.
• The frozen-encoder save. A team tried to fine-tune the whole vision tower on 5k in-house images and torched the CLIP features (catastrophic forgetting), tanking everything. Reverting to freeze encoder, train projector, LoRA the LLM fixed it in a day. Train the connector, not the world.
• The temperature collapse. Someone hard-coded τ too small "to make it confident"; training diverged. CLIP learns and clamps τ for a reason — it's a tuned lever, not a confidence dial.
• The shuffled-batch bug. A data loader desynced images and captions; the contrastive loss stayed high and nobody knew why. The diagonal was the positive — except the diagonal pairs no longer matched. (This is literally the lab's "shuffled ⇒ high loss" soul test, in production.)

Vocabulary (rapid-fire)

• ViT — Vision Transformer; image → P×P patches → linear embed → transformer.
• Patch embedding — the Conv2d(kernel=stride=P) that turns each patch into a token.
• [CLS] token — learned vector prepended; its output is the whole-image embedding.
• CLIP — Contrastive Language-Image Pretraining; aligns image & text in one space.
• Cosine similarity — a·b/(|a||b|) ∈ [-1,1]; the distance metric of the shared space.
• InfoNCE — the symmetric contrastive loss; diagonal = positive, off-diagonal = negatives.
• Temperature / logit_scale — sharpens the softmax over similarities; learned in CLIP.
• In-batch negatives — every non-matching pair in the batch; free training signal.
• Zero-shot — build the classifier from text prompts; argmax cosine; no task training.
• Projector / connector — the MLP that maps vision features into LLM token space (LLaVA).
• Cross-attention fusion — Flamingo-style; image is K/V the LLM attends to.
• SigLIP — CLIP with a sigmoid loss instead of softmax InfoNCE; scales to bigger batches.
• Grounding — localizing where (boxes, points); needs spatial token resolution.

Beginner mistakes

• Forgetting an image is hundreds-to-thousands of tokens, so it dominates the KV-cache.
• Treating resolution as a quality knob instead of a capability + cost decision (OCR needs it).
• Skipping the zero-norm guard in cosine similarity (a degenerate embedding → divide by zero).
• Not subtracting the max in softmax/InfoNCE → overflow on large logits.
• Computing InfoNCE one-directionally — it's symmetric (rows and columns).
• Fine-tuning the whole vision encoder instead of freezing it and training the projector.
• Hard-coding the temperature instead of letting it be a learned/tuned lever.
• Thinking the LLM has special image logic — in LLaVA the image tokens are just tokens.

The one thing to take away

A VLM is ViT (image→tokens) + CLIP (shared space) + connector (image→LLM) — and the moment you add an image you've added a pile of tokens to your context and KV-cache. If you can name those three boxes, write the symmetric InfoNCE loss, and count an image's tokens from (H, W, P), you can reason about any multimodal model as a system. If you can't, "the model can see now" is just a demo.

Brother Talk — Vision-Language Models

Off the record. The stuff I'd tell you over a beer, not in a design review.

Vision-language sounds like a different universe, and it absolutely is not. When I first heard "multimodal model" I pictured some exotic architecture I'd never understand. Then I took one apart and laughed: it's a ViT that chops the image into squares, a contrastive loss that's just cross-entropy on a similarity matrix, and an MLP that glues the image tokens into the prompt. That's it. Three pieces you already half-know. The mystique is marketing; the mechanism is a weekend. Once you see "an image is just more tokens," the whole field opens up and stops being intimidating.

The thing nobody tells you: images are token bombs. Everyone obsesses over "can the model see" and nobody runs the arithmetic on what seeing costs. One high-res image can be three thousand tokens — bigger than most prompts. I've watched a team's beautiful vision demo fall over the instant a real user uploaded a 4K screenshot, because they sized the context for the words and forgot the picture is the expensive part. Be the person in the room who says "what resolution, and how many image tokens, and what does that do to the KV-cache?" That single question is the whole Phase 00 cost math coming back to save you, and almost nobody asks it.

Resolution is where the real fights happen. People treat it like a quality slider — "bump it up, looks sharper." No. Resolution is a capability decision. Want OCR? You need the resolution or the small text literally doesn't survive being cut into patches. Want to point at a button (grounding)? Same. But every doubling roughly quadruples the tokens and the cost. The senior move is to match resolution to the task, not to vibes — and to know that "it can't read the receipt" is usually a resolution problem, not an intelligence problem.

"What's frozen?" is the question that makes you sound senior in five words. When someone asks how a VLM is trained, the junior answer is hand-wavy "you train it on image-text data." The senior answer is precise: freeze the vision encoder, train the projector, LoRA the LLM. You're training a tiny connector on top of two big frozen models — which is why a startup with one GPU can build a VLM and why your nested-list lab can build a toy one. Internalize "train the connector, not the world." It's the truth about how this entire field actually ships, and it walks straight into the LoRA phase next.

The contrastive loss is the most satisfying thing in this whole track to actually feel. You build the similarity matrix, you run the loss, it's a sane number. You shuffle the captions so the matched pairs fall off the diagonal, you run it again, and the number jumps. That's it — that's contrastive learning, live, on your screen. No GPU, no 400 million images, no mystery. The lab's "matched ⇒ low, shuffled ⇒ high" test is the soul of CLIP in ten lines, and the day it clicks is the day "embedding space" stops being a buzzword.

Don't get hypnotized by the frontier demos. GPT-4V and Claude reading a whiteboard photo is genuinely magic to watch, and it's easy to think the secret sauce is unknowable. It isn't. It's these mechanisms at enormous scale with great data. Your job isn't to reproduce the scale — it's to understand the mechanism so well that when a VLM is wrong, slow, or expensive, you know exactly which box (encoder, alignment, fusion, token budget) to poke. That understanding is what they're actually paying a senior for.

Zero-shot is the trick that makes you look like a wizard to product people. "We can add a new category without retraining — just write a prompt for it" sounds like science fiction in a planning meeting, and it's true, because the classifier weights are text embeddings. But here's the part nobody mentions: the wording matters a lot. "a photo of a {x}" beats the bare word, and ensembling templates buys real accuracy. So when zero-shot underperforms, don't assume the model is dumb — half the time it's the prompt. Treat the prompt as a tunable, because it is.

The failure modes are weirdly mundane. You'd expect multimodal bugs to be exotic. They're not. The big ones are: images silently desynced from captions in the data loader (loss won't drop and you'll burn a week), resolution too low for the actual task (OCR garbage), and the context quietly exploding because nobody counted image tokens. None of those are deep ML problems — they're plumbing and arithmetic. Which is great news: it means careful engineering, not genius, is what ships a working VLM.

What's worth caring about: the three boxes (ViT, CLIP, connector), the token-count arithmetic, the symmetric InfoNCE loss, and "what's frozen." What's not worth losing sleep over: memorizing the exact CLIP architecture hyperparameters, or whether your toy projector matches LLaVA's weight init to the decimal. Estimates and mechanism win interviews; trivia doesn't.

The career framing. Multimodal is where the field is going — agents that see screens, robots that see rooms, assistants that read documents. The engineers who can reason about how a model sees, as a system with a token budget and a training recipe, are going to be the ones designing those systems instead of just calling the API. Get this phase cold and you're not "the person who uses the vision model" — you're "the person who knows what it costs and why it works." Now go make the soul test green.

Lab 01 — ViT Patch Embedder + CLIP Contrastive Encoder + LLaVA Projector

Phase: 04 — Vision-Language Models Difficulty: ⭐⭐⭐☆☆ (the math is small; the fusion plumbing is what trips people) Time: 3–4 hours

A vision-language model is three machines bolted together: a Vision Transformer that turns pixels into tokens, a CLIP objective that teaches an image encoder and a text encoder to live in the same embedding space, and a connector that splices image tokens into an LLM's token stream so the language model can "read" the picture. This lab builds all three from scratch on nested lists — patchify + patch_embed (ViT), cosine_similarity + clip_similarity_matrix + info_nce_loss + clip_zero_shot_classify (CLIP), and Projector + assemble_multimodal_sequence (LLaVA) — so the mechanism that lets GPT-4V, Claude, and LLaVA see is muscle memory, not magic.

What you build

• patchify(image, patch_size) — ViT's tokenizer: chop an H×W grayscale image into non-overlapping P×P patches in raster order, each flattened to a vector; reject non-divisible sizes. An image becomes (H/P)·(W/P) tokens — "an image is worth 16×16 words."
• patch_embed(patches, W) — the linear projection that lifts each flat patch into a d-dim token, with an optional prepended learned [CLS] vector.
• cosine_similarity / l2_normalize — the similarity metric CLIP lives on, with the zero-norm guard interviewers love to probe.
• clip_similarity_matrix(image_embs, text_embs, temperature) — the temperature-scaled cosine logits: entry [i][j] is how much image i matches text j, divided by temperature.
• info_nce_loss(sim_matrix) — the symmetric InfoNCE loss (image→text + text→image cross-entropy, diagonal = positive), with the log-sum-exp trick so large logits never overflow. The soul of the lab: low when matched pairs sit on the diagonal, high when you shuffle.
• clip_zero_shot_classify(image_emb, class_text_embs) — classify by argmax cosine similarity to text prompts: no training, no fixed label set.
• Projector(in_dim, out_dim, hidden, seed) — LLaVA's connector: a seeded 2-layer GELU MLP that maps vision features into the LLM's token-embedding width.
• assemble_multimodal_sequence(text_token_embs, image_token_embs, idx) — splice the projected image tokens into the text stream at the <image> placeholder, yielding one fused sequence of length len(text) - 1 + len(image).

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why patchify of a 4×4 image with patch_size=2 gives 4 patches of length 4, and why a non-divisible size must raise.
• You can state why cosine_similarity is 1 for identical, -1 for opposite, 0 for orthogonal vectors, and why a zero-norm input returns 0 instead of dividing by zero.
• You can explain why test_info_nce_low_when_pairs_matched and test_info_nce_high_when_pairs_shuffled are the whole point — the loss is low only when the diagonal (the matched pair) dominates its row and column.
• You can explain why info_nce_loss([[1000,0],[0,1000]]) is ~0 and does not overflow (max-subtraction / log-sum-exp).
• You can explain why assemble_multimodal_sequence returns length len(text) - 1 + len(image) and where in the sequence the image tokens land.
• Given any "how does an LLM see an image?" prompt you can sketch ViT → CLIP-style encoder → projector → LLM, and name the token-budget cost.

How this maps to the real stack

Limits of the miniature (be honest in the interview): real ViTs add learned position embeddings and run a full transformer over the patch tokens — we stop at the linear patch embed; CLIP trains the encoders for hours on hundreds of millions of pairs — we just score fixed vectors; the projector here is randomly initialized, not trained (training it, with the encoder frozen and LoRA on the LLM, is Phase 05); and we model grayscale images as nested lists, so there is no 3-channel Conv2d, no interpolation of position embeddings for new resolutions, and no O(T²) attention cost — the parts that make high-resolution VLMs expensive.

Extensions (build these on real hardware)

• Add learned position embeddings and a [CLS] token to patch_embed, then run a one-layer self-attention block over the patch tokens (reuse your Phase 02 attention) to make it a real (tiny) ViT.
• Add a real temperature parameter that is learned (CLIP's logit_scale) and show the gradient of info_nce_loss w.r.t. it (reuse your Phase 03 autograd).
• Swap InfoNCE for SigLIP's pairwise sigmoid loss and compare batch-size sensitivity.
• Implement Flamingo-style cross-attention fusion (image tokens as keys/values the LLM attends to) and contrast its token-budget cost with the LLaVA projection approach.
• Compute the KV-cache inflation: an image at 336×336 with 14×14 patches is 576 tokens — plug that into your Phase 00 kv_cache_bytes and show how many images blow the context budget.
• Load a real openai/clip-vit-base-patch32 and check your zero-shot classification picks the same class on a handful of images.

Interview / resume

• Talking points: "How does an LLM actually see an image?" "What is contrastive pretraining and why does the temperature matter?" "Projection vs cross-attention fusion — when would you pick each?" "How does adding an image change the KV-cache and the cost?"
• Resume bullet: Built a from-scratch vision-language pipeline — ViT patch embedding, CLIP-style contrastive encoder with a numerically stable symmetric InfoNCE loss and zero-shot text-prompt classification, and a LLaVA-style MLP connector that fuses projected image tokens into an LLM's token stream — clarifying the exact mechanism behind multimodal models.

Phase 05 — PEFT: LoRA, QLoRA & Knowledge Distillation

The phase where adaptation stops being a luxury. Phase 00 taught you that a full fine-tune holds the weights + gradients + optimizer state + activations in memory — ~16 bytes/param, so a 7B model needs ~112 GB before activations, and a 65B model is hopeless on commodity hardware. This phase is how the industry got around that: don't move the giant weights — freeze them and learn a tiny correction. LoRA, QLoRA, and distillation are the three techniques that turn "fine-tuning needs a cluster" into "fine-tuning needs one GPU and an afternoon."

Why this phase exists

A Senior AI Engineer is asked, constantly: "We need this model to do our task — how?" The naive answer (full fine-tune) is usually wrong on cost, and the lazy answer ("just prompt it harder") is usually wrong on quality. The senior answer is a menu:

• LoRA — freeze W, learn a rank-r update ΔW=(α/r)·A·B. You train r·(d+k) params instead of d·k — often <1% of the model — and at inference you merge() the delta back so there is zero added latency. One base, many cheap swappable adapters.
• QLoRA — store the frozen base in 4-bit NF4, keep the small LoRA adapters in 16-bit, and you can fine-tune a 65B model on a single 48 GB GPU. The base never trains, so its quantization error never compounds.
• Distillation — when you want a permanently smaller model, train a small student on a big teacher's softened logits. The student inherits the teacher's "dark knowledge" (the relative probabilities of the wrong answers), often matching it at a fraction of the inference cost — which, recall Phase 00, is the bill you pay forever.

These are not interchangeable. LoRA/QLoRA adapt a frozen model cheaply; distillation shrinks one permanently. Knowing which lever a problem wants — and being able to defend it with the param-count and memory math — is the skill this phase installs.

What to recall from earlier phases

• Phase 00 (cost math). Training memory ≈ 16 bytes/param (weights + gradients + Adam moments + fp32 master), and inference cost is 2N/token paid forever. Both facts are the reason PEFT and distillation exist; this phase is their direct payoff.
• Phase 00 (quantization). Weights at 0.5 bytes/param (int4) vs 2 (fp16) is the difference between two GPUs and one — QLoRA puts that lever under fine-tuning.
• Phase 02/03 (the transformer & autograd). You know a block is attention projections (q,k,v,o) plus an MLP, and what a gradient/optimizer step costs — PEFT decides which of those weights get gradients at all.

Concept map

                  ┌──────────────────────────────────────────────┐
                  │  Full fine-tune is too expensive (P00 math):  │
                  │  weights + grads + Adam state ≈ 16 B/param     │
                  └──────────────────────────────────────────────┘
                          │                          │
            ADAPT a frozen model              SHRINK to a new model
                          │                          │
            ┌─────────────┴─────────────┐            ▼
            ▼                            ▼      DISTILLATION
          LoRA                        QLoRA      teacher → student
   ΔW = (α/r)·A·B               4-bit NF4 base    softmax(z/T)
   A:(d×r) gauss                + 16-bit LoRA     L = T²·KL(t‖s)
   B:(r×k) ZERO  ─── start = pretrained          + α·CE(s, y)
   train r·(d+k) ≪ d·k          double-quant      "dark knowledge"
            │                   paged optimizer          │
            ▼                          │                 ▼
        merge() ⇒ 0 latency            ▼          smaller N ⇒ less
   swappable adapters         65B on one GPU       inference cost forever
            └──────────────┬───────────┴─────────────────┘
                           ▼
                THE PEFT DECISION (recall: inference cost is forever)
        which layers? what rank? adapt-and-serve, or distill-and-replace?

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py.

Integrated scenario ideas

• Size a LoRA fine-tune: for a 7B model, count the trainable params if you adapt q,v at r=8 across 32 layers; show it is well under 1% and that the optimizer state shrinks by the same factor — the reason it fits one GPU.
• Defend QLoRA on a 65B: compute the 4-bit base footprint (~35 GB) vs the fp16 base (130 GB), add the tiny adapter, and show why one 48 GB GPU now works.
• Adapter zoo: one frozen base + 20 task adapters at a few MB each vs 20 full fine-tuned copies; compute the storage and the swap cost.
• Distill, don't scale (Phase 00 callback): a 70B teacher vs a distilled 7B student at your QPS — compute the perpetual $/1M gap and the breakeven volume.
• Pick the lever: given "we need 5 niche tasks on shared infra" vs "we need one task served at 50k QPS forever," argue LoRA-zoo vs distillation with numbers.

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive ΔW=(α/r)·A·B and explain the orientation (A:d×r, B:r×k).
• You can explain why B is zero-initialized in one sentence and what breaks if it is not.
• You can compute the LoRA trainable fraction for a given (d,k,r) in your head.
• You can explain QLoRA's three tricks (NF4, double quantization, paged optimizer) and why the base never training is what makes 4-bit safe.
• You can write the distillation loss from memory and explain T and the T² factor.

Key takeaways

• Adapt the update, not the weights. Weight updates during fine-tuning are empirically low-rank, so a rank-r A·B captures most of the gain at r·(d+k) params. This is the whole hypothesis behind LoRA — and why it works at all.
• B zero-init is the design, not a detail. It makes the adapted model identical to the pretrained one at step 0, so training only ever adds skill from a known-good starting point. Lose this and you start from a randomly corrupted model.
• QLoRA decouples storage precision from training precision. The base is frozen, so its 4-bit error is fixed and never compounds; only the high-precision adapters learn. That decoupling is what fits a 65B fine-tune on one GPU.
• Distillation moves capability into a smaller N. Softened logits carry far more signal than hard labels, so the student learns the teacher's function, not just its answers — and a smaller N is cheaper to serve on every request, forever (P00).
• The lever depends on lifetime, not vibes. Many cheap swappable tasks → LoRA zoo. One task served at huge scale → distill and replace. Memory-constrained training → QLoRA. Defend it with the math, not the hype.

Next: Phase 06 — Quantization & Model Compression.

Warmup Guide — PEFT: LoRA, QLoRA & Knowledge Distillation

Zero-to-senior primer for Phase 05. We start from the brutal arithmetic of a full fine-tune (recall Phase 00: weights + gradients + optimizer state ≈ 16 bytes/param) and build, from first principles, the three techniques that made adaptation cheap: LoRA (the low-rank update ΔW=(α/r)·A·B, the zero-initialized B, and the mergeable adapter), QLoRA (a 4-bit NF4 frozen base, double quantization, paged optimizers — a 65B fine-tune on one GPU), the other PEFT methods (adapters, prefix/prompt tuning, IA³) and when each, and knowledge distillation (training a small student on a big teacher's softened logits). Every section gives you the mechanism under the hood, not the framework call.

Table of Contents

• Chapter 1: Why PEFT — The Cost of a Full Fine-Tune
• Chapter 2: The LoRA Hypothesis — Updates Are Low-Rank
• Chapter 3: The LoRA Mechanism — A, B, r, α, and Merging
• Chapter 4: Why B Is Zero-Initialized (The Soul of LoRA)
• Chapter 5: Which Layers to Adapt, and Choosing r and α
• Chapter 6: QLoRA — 4-Bit NF4, Double Quantization, Paged Optimizers
• Chapter 7: The Rest of the PEFT Zoo — Adapters, Prefix/Prompt, IA³
• Chapter 8: Knowledge Distillation — Softened Logits & Dark Knowledge
• Chapter 9: The Practical Recipe — SFT → Data → Eval
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Why PEFT — The Cost of a Full Fine-Tune

From zero. "Fine-tuning" means continuing to train a pretrained model on your data so its weights shift toward your task. The obvious way is to update every weight — full fine-tuning. The problem is not the idea; it is the memory bill, and that bill is arithmetic you already did in Phase 00.

The memory of training (recall P00, Chapter 3). To train a parameter you hold, per parameter, much more than the parameter itself:

…plus activations (which scale with batch × sequence length). So a 7B full fine-tune is 7e9 × 16 = 112 GB before activations — already over one 80 GB GPU — and a 65B model is ~1 TB, hopelessly out of reach without a cluster.

The insight that unlocks PEFT. Almost all of those 16 bytes are spent on the optimizer state and gradients of weights you barely change. What if you froze the pretrained weights entirely (no gradient, no optimizer state, no master copy — 0 extra bytes for them) and only trained a small set of new parameters? Then the expensive ~14 bytes/param of training overhead applies only to that small set. That is the entire premise of Parameter-Efficient Fine-Tuning (PEFT): keep the giant frozen, learn something tiny.

Why this matters in production. PEFT is not a research curiosity — it is how teams ship. One frozen base + a few-MB adapter per task means you can serve dozens of specialized behaviors from one set of weights, version adapters like code, and fine-tune on a single GPU. And recall the asymmetry from P00: training cost is paid once, but inference cost is paid on every request forever — so a technique like LoRA that adds zero inference cost (after merging) is doubly attractive.

Common misconception. "PEFT trades quality for cheapness." On most adaptation tasks, a well-tuned LoRA matches full fine-tuning to within noise — because (next chapter) the update you needed was low-rank all along. You are not approximating a full fine-tune; you are exploiting its structure.

Chapter 2: The LoRA Hypothesis — Updates Are Low-Rank

The claim. When you fine-tune a pretrained model, the change to each weight matrix — ΔW = W_finetuned − W_pretrained — has a low intrinsic rank. Empirically, ΔW can be approximated well by a matrix of rank r far smaller than the matrix's dimensions. Hu et al. (2021) measured this directly and found r as small as 1–8 often suffices.

What "rank" means here, from zero. The rank of a matrix is the number of independent directions it actually uses. A d×k matrix can have rank up to min(d,k), but a low-rank matrix — say rank r — can be written as the product of a d×r and an r×k matrix:

$$ \underbrace{\Delta W}{d\times k} = \underbrace{A}{d\times r};\underbrace{B}_{r\times k}, \qquad r \ll \min(d,k). $$

The full d×k matrix has d·k numbers; the factored form has only r·(d+k). For a 4096×4096 matrix at r=8: 16.8M vs 65.5K — a 256× reduction.

Why fine-tuning updates are low-rank (the intuition). Pretraining already taught the model general language. Your fine-tune is a comparatively small, directional nudge — "prefer this style," "answer in this format," "know this domain." A nudge does not need to rewrite all min(d,k) directions of a weight matrix; it needs to push a handful. The data confirms the intuition: the useful signal in ΔW concentrates in a few singular directions, so a rank-r factorization captures it.

   full fine-tune learns ΔW directly (d·k params, all of it trainable)
        ┌──────────────┐
   W +  │     ΔW        │          rank up to min(d,k)
        └──────────────┘

   LoRA learns ΔW = A·B  (only r·(d+k) params)
        ┌──┐
   W +  │A │ · [  B  ]              rank exactly r ≪ min(d,k)
        │  │   (r×k)
        └──┘
        (d×r)

Common misconception. "Low-rank means low-quality." Rank limits the number of directions the update can move, not the magnitude of the move (that is what α controls — Chapter 3). For an adaptation that genuinely needs few directions, a rank-r update is not lossy; it is exactly the right capacity, and it also regularizes.

Chapter 3: The LoRA Mechanism — A, B, r, α, and Merging

The setup. Take any frozen linear layer with weight W of shape (d×k) that maps an input x (a row vector of width d) to x·W (width k). LoRA leaves W frozen and adds a parallel low-rank path:

$$ h = x W + \frac{\alpha}{r}, x A B, \qquad A\in\mathbb{R}^{d\times r},; B\in\mathbb{R}^{r\times k}. $$

(Orientation note — this lab fixes W as (d×k), inputs as (n×d). The LoRA paper writes W as (k×d) acting on column vectors; the algebra is identical, just transposed. Pick one and be consistent — the lab is consistent.)

The roles of the two new matrices.

• A (d×r) projects the input down into an r-dimensional bottleneck.
• B (r×k) projects that bottleneck up to the output width.
• Their product A·B is (d×k) — the same shape as W — but constrained to rank ≤ r.

The role of r (rank = capacity). r is the width of the bottleneck — how many independent directions the update may use. Bigger r = more capacity (and more params, r·(d+k)). Typical values: 4, 8, 16, 64. Past a point, more r stops helping because the true update's rank is the ceiling.

The role of α (scaling = magnitude). The update is scaled by α/r before it is added. Why divide by r? So the effective magnitude of the update is decoupled from the rank: if you raise r to get more capacity, the 1/r keeps the output scale roughly constant, so you do not have to re-tune the learning rate. In practice people fix α (often α = 2r or α = r) and tune r. Doubling α doubles the delta linearly — the lab tests this exactly.

$$ \Delta W = \frac{\alpha}{r} A B \quad\Rightarrow\quad \Delta W \big|{2\alpha} = 2,\Delta W \big|{\alpha}. $$

The cheap forward path. Notice the lab computes (x·A)·B, not x·(A·B). Both are equal, but x·A is (n×r) — a skinny intermediate — so the adapter's extra compute is O(n·d·r + n·r·k) instead of materializing a (d×k) delta. During training this skinny path is what makes LoRA cheap to run.

Merging — the zero-latency trick. Because the adapter is linear and parallel, you can fold it into the base once training is done:

$$ W_{\text{merged}} = W + \frac{\alpha}{r} A B. $$

Now x · W_merged is a single matmul — identical to the LoRA forward, with zero extra inference cost. This is a defining advantage over adapters (Chapter 7), which insert new layers you cannot fold away. (You can also keep adapters unmerged to hot-swap many tasks against one base — you trade a little latency for flexibility.)

Common misconception. "LoRA is slower at inference." Only if you leave it unmerged. Merged, it is exactly as fast as the original model — that is the whole point.

Chapter 4: Why B Is Zero-Initialized (The Soul of LoRA)

This is the single detail interviewers love and beginners miss, so it gets its own chapter.

The mechanism. At the start of training, LoRA initializes:

• A with small random Gaussian values, and
• B with exactly zeros.

Therefore the initial update is

$$ \Delta W \big|_{t=0} = \frac{\alpha}{r}, A \cdot \mathbf{0} = \mathbf{0}. $$

The product of anything with the zero matrix is the zero matrix. So at step 0:

$$ h = x W + \frac{\alpha}{r} x A B = x W + 0 = x W. $$

The adapted model is bit-for-bit the pretrained model at initialization. The lab's soul test asserts exactly this: delta() is all zeros and forward(x) == x · W_base with no floating-point drift.

Why this is the right design, not a convenience. Fine-tuning is continued training from a known-good model. If your adapter started by injecting random noise into W (which a random B would do), step 0 would corrupt a model that took millions of dollars to pretrain — you would spend the first part of training just climbing back to where you started, and you might never recover. Zero-init means training begins at the pretrained optimum and can only ever add skill, gradient step by gradient step.

Why not zero-init both A and B? If both were zero, the gradient of the loss with respect to both would also be zero (the product is zero and so are its partial derivatives in this symmetric configuration), and the adapter would be stuck — it could never start learning. You need exactly one side non-zero to break the symmetry and let gradients flow: random A, zero B. (PEFT does exactly this.)

  step 0:   A = random,  B = 0   →   ΔW = (α/r)·A·0 = 0   →   model == pretrained
  step 1+:  gradients flow into B (A non-zero breaks symmetry) → ΔW grows from 0

Common misconception. "It doesn't matter which one is zero." It must be exactly one. Zero A + random B also gives ΔW=0 at start, but PEFT conventionally zeros B; the non-negotiable part is that one is zero (so you start at the pretrained model) and the other is non-zero (so learning can begin).