The premise of CS336 is contrarian: in a world of open-weight downloads and API calls, why would anyone build a language model from the ground up? Because there are three kinds of knowledge, and only one of them survives copy-paste.
The course's thesis: efficiency is the master skill. Every decision — architecture, precision, parallelism strategy, data pipeline — is ultimately about getting the best model for a given compute budget. The bitter lesson (Sutton, 2019) says that scale wins over cleverness. But scale without efficiency is just waste.
Five assignments map the full pipeline: A1 Tokenizer + Transformer + Training loop · A2 Systems — kernels and parallelism · A3 Scaling laws — predict loss before training · A4 Data — build a pipeline from Common Crawl · A5 Alignment — post-training with RLHF/RLVR. Each assignment builds on the last. By the end, you've built every piece.
A language model doesn't see text — it sees a sequence of integer IDs. Tokenization is the bridge: convert a raw string into a sequence of tokens from a fixed vocabulary that the model can process.
The problem sounds simple but hides real trade-offs. The vocabulary size V directly affects model size (the embedding matrix is V × d), sequence length (fewer tokens = less compute per document), and coverage (can you represent any input?).
The path from raw text to subword tokens goes through three layers:
BPE is the algorithm behind GPT, LLaMA, and most modern tokenizers. The idea is beautifully simple: start with bytes, then greedily merge the most frequent pair, over and over, until your vocabulary is big enough.
Each merge reduces sequence length by replacing two tokens with one. Early merges capture universal patterns ("t"+"h" → "th", "e"+" " → "e "). Later merges capture common words ("the", "ing", "tion"). The merge list is the tokenizer — applying the same merges in the same order to new text produces the same tokenization.
A good BPE tokenizer achieves ~4:1 compression — 4 bytes of text per token on average for English. This means a 100K-byte document becomes ~25K tokens. The compression ratio is the key efficiency metric: higher compression = shorter sequences = less compute.
Watch BPE build up tokens from individual characters. Each step merges the most frequent adjacent pair.
Two problems. First, the vocabulary is effectively unbounded — new words, misspellings, technical jargon, and multilingual text all create unknown tokens. Second, even with a large fixed vocabulary, the embedding matrix (V × d) becomes enormous — a 500K-word vocabulary with d=4096 is 2 billion parameters just for embeddings. Subword tokenization (BPE) solves both: any byte sequence can be represented, and V is tunable to 32K–128K.
Sequence length decreases — more merges means more common substrings become single tokens, so fewer tokens are needed per document. But the returns diminish: going from 32K to 64K vocab gives a meaningful compression gain; going from 128K to 256K gives much less. And the embedding matrix grows linearly with V, so there's a trade-off: shorter sequences (less compute in attention) vs. larger embeddings (more parameters).
Mechanics (how things work), mindset (how to think about problems), and intuitions (feel for scale). Only mechanics transfers from reading — you can learn what a transformer is from a paper. Mindset (e.g., always thinking about resource costs) and intuitions (e.g., knowing that 1e22 FLOPs is "small") only come from building things yourself. That's the argument for building from scratch.
Everything in training is a tensor — parameters, gradients, optimizer states, data, activations. A tensor's memory footprint is simply the number of elements times the bytes per element. The bytes per element depends on the precision you choose, and that choice has cascading effects on memory, speed, and numerical stability.
| Type | Bits | Bytes | Layout | Use case |
|---|---|---|---|---|
| float32 | 32 | 4 | Optimizer states, small-model training. Safe but slow and memory-heavy. | |
| float16 | 16 | 2 | Risky — only 5 exponent bits. 1e-8 rounds to zero. Underflow and NaN problems. | |
| bfloat16 | 16 | 2 | The sweet spot. Same dynamic range as float32, less resolution. Standard for training. | |
| fp8 | 8 | 1 | Emerging. Two variants (E4M3 / E5M2). Hardware support via NVIDIA Transformer Engine. | |
| fp4 | 4 | 0.5 | Block-scaled — groups of values share a scaling factor. Only 16 possible values per block. |
The key insight behind bfloat16: it trades resolution (fewer mantissa bits) for dynamic range (same exponent bits as float32). In deep learning, you care more about not overflowing or underflowing than about the difference between 3.14159 and 3.14160. Stochastic gradient descent is inherently noisy — the precision of individual values matters less than the range they can represent.
Standard practice: use bf16 for parameters, activations, and gradients. Use fp32 for optimizer states (Adam's first and second moments need the stability). PyTorch's AMP (Automatic Mixed Precision) handles the casting — it uses bf16 for safe ops like MatMuls and keeps fp32 for numerically sensitive ops like exponentiation and layer norm.
Reading x @ y.transpose(-2, -1) and figuring out what -2 means is a recipe for bugs. Einops (Einstein operations) replaces index arithmetic with named dimensions — you say what you mean, and the library handles the mechanics.
Three primitives cover nearly everything:
The payoff: you never write a transpose. The dimension names make the operation self-documenting. And the shapes are checked at runtime — if your dimensions don't match, you get a clear error instead of a silent wrong result.
A FLOP is one floating-point operation — an addition or a multiplication. The number of FLOPs tells you how much work a computation requires, independent of hardware.
FLOPs (lowercase s) = floating-point operations, a count of work done. FLOP/s = floating-point operations per second, a measure of hardware speed. When NVIDIA says an H100 does "989 teraflops," that's FLOP/s — and read the fine print: it's with sparsity, so divide by 2 for dense workloads.
For a matrix multiply of shapes (B, D) × (D, K), the FLOPs are:
This formula scales with the product of all three dimensions. And it dominates: elementwise operations (ReLU, GELU, addition) cost O(n) FLOPs — negligible compared to the O(n³) of matrix multiplies for large enough matrices.
Another way to read the formula: for a linear layer with D×K parameters processing B data points, the FLOPs are 2 × tokens × parameters. This shape generalizes to transformers.
Counting FLOPs tells you how much work to do. But whether that work is fast or slow depends on something subtler: the ratio of compute to data movement.
Here's the hardware picture: tensors live in HBM (high-bandwidth memory). To compute on them, you ship them to the accelerator cores, do the math, and ship results back. Two speeds govern this:
H100 bf16: ~989 TFLOP/s (dense). How fast the cores can multiply and add.
H100 HBM3: 3.35 TB/s. How fast data moves between memory and compute.
The arithmetic intensity of an operation is: FLOPs performed / bytes moved. The accelerator intensity is: peak FLOP/s / memory bandwidth. For the H100: ~989e12 / 3.35e12 ≈ 295. This is the breakeven point.
The practical implication: transformers are designed to live in the compute-bound regime. The core operation is large matrix multiplies (attention, feedforward layers), which have O(n³) compute but only O(n²) data movement. Everything between the MatMuls (LayerNorm, ReLU, softmax) is memory bound but fast in absolute terms.
At inference, you generate one token at a time — matrix-vector products, not matrix-matrix. Intensity drops from ~n/3 to ~0.5, and you become memory bound. This is why inference is so much slower per token than training, and why batching inference requests matters enormously.
Enter matrix dimensions to see whether the operation is memory-bound or compute-bound on an H100.
How many FLOPs does one training step cost? For a model with N parameters processing a batch of D tokens:
Where does the 6 come from? Consider a single linear layer W of shape (D, K). The forward pass is a MatMul: 2·B·D·K FLOPs. The backward pass computes two gradients — one with respect to the input (for backpropagation) and one with respect to the parameters (for the update). Each is also a MatMul with the same three dimensions, just contracted differently. So backward = 2 × forward.
Sum across all layers: forward = 2ND, backward = 4ND, total = 6ND.
How long to train a 70B model on 15T tokens on 1024 H100s? FLOPs = 6 × 70e9 × 15e12 = 6.3e24. At 989 TFLOP/s per GPU, MFU 0.5: effective FLOP/s = 1024 × 989e12 × 0.5 ≈ 5.1e17. Time = 6.3e24 / 5.1e17 ≈ 12.4M seconds ≈ 143 days. This is the kind of calculation the course wants you to do reflexively.
Training memory breaks into four buckets, each scaling differently:
For Adam in bf16 mixed precision: 2 + 2 + 4 + 4 = 12 bytes per parameter just for model state (not counting activations). An H100 has 80 GB of HBM. Ignoring activations: 80e9 / 12 ≈ 6.7B parameters on a single GPU. With activations, substantially less.
Two standard techniques to reduce memory pressure:
You want a large effective batch size for training stability, but large batches eat activation memory. Gradient accumulation splits the batch into micro-batches: compute gradients on each micro-batch, accumulate them (don't zero between micro-batches), and update parameters only after processing all micro-batches.
Mathematically equivalent to a full batch. The memory savings come from only holding one micro-batch of activations at a time.
In standard training, you store activations for every layer during the forward pass (needed for the backward pass). Activation checkpointing only stores activations at a subset of layers and recomputes the missing ones during backward. Classic compute-for-memory trade-off.
In PyTorch: wrap a layer with torch.utils.checkpoint and it handles the save/recompute logic. For a deep network with linear + ReLU blocks, checkpointing each block saves roughly half the activation memory.
MFU measures how much of the hardware's theoretical compute you're actually using:
Why only 50%? Memory-bound operations (LayerNorm, activations, softmax) between the MatMuls. Communication overhead in distributed training. Kernel launch latency. Memory bandwidth ceilings on non-MatMul operations. The gap between 0.5 and 1.0 is the systems engineering challenge of the course.
Float16 has only 5 exponent bits, giving it poor dynamic range — values like 1e-8 round to zero, causing underflow and NaN instabilities during training. Bfloat16 trades mantissa bits for exponent bits: 8 exponent bits (same as float32) with only 7 mantissa bits. You lose precision but keep the full dynamic range. Since gradient descent is inherently noisy, the precision loss doesn't matter much, but overflow/underflow kills training. Google developed bf16 specifically for deep learning in 2018.
They take the same time. Both are deeply memory-bound (intensity ≈ 0.25 for ReLU, ≈ 5 for GELU — both far below the H100's threshold of 295). The bottleneck isn't the compute; it's shipping the tensor to the accelerator and back. Since both move the same number of bytes (read input + write output), the wall-clock time is identical. The 20× more FLOPs in GELU is invisible because the cores are idle waiting for data anyway.
For each linear layer, the forward pass does one MatMul: h = x @ W (2BDK FLOPs). The backward pass computes two MatMuls: (1) the gradient w.r.t. the input x_grad = h_grad @ W.T (for backpropagation to earlier layers) and (2) the gradient w.r.t. the parameters W_grad = x.T @ h_grad (for the optimizer update). Each backward MatMul has the same three dimensions as the forward, just contracted over a different pair. So backward = 2 × forward.
With Adam in mixed precision: 2 (params, bf16) + 2 (gradients, bf16) + 4 (first moment, fp32) + 4 (second moment, fp32) = 12 bytes per parameter. So 80 GB / 12 ≈ 6.7B parameters — but this ignores activations, which scale with batch size and sequence length. With a typical batch, you're limited to roughly 3–5B parameters. The bottleneck is the optimizer states at 8 bytes per parameter (fp32) — they consume more memory than the model itself. This is why techniques like ZeRO (sharding optimizer state across GPUs) exist.
With resource accounting under your belt, Module 3 opens the transformer itself. Self-attention, multi-head attention, feedforward layers, layer normalization, positional encodings — and the FLOPs/memory cost of each component. You'll count every parameter and every multiply, and see exactly where compute goes in a real model. Lecture by Tatsunori Hashimoto.