Generative models · the from-scratch course

CS336 — Language Modeling from Scratch

A linearized rebuild of Stanford's CS336: how a modern large language model is built end to end — tokenizer, architecture, kernels, parallelism, running the run, scaling laws, data, alignment, long context, inference — derived as one chain of decisions under a fixed compute budget.

This is the rigorous, full-stack sibling of the six-lesson Mini GPT track. Mini GPT is the compact tour of post-training; this course builds the whole thing from scratch, the way Stanford's CS336 does — and reorganizes the five course units (basics → systems → scaling → data → alignment & inference) into a single derivation where every lesson is forced by the wall the previous one hit.

The one idea the whole track turns on

You are handed a fixed compute budget — C ≈ 6 · N · D FLOPs (N parameters, D training tokens), paid in GPU-hours. Every decision in the stack — how you tokenize, what architecture, which kernel, how you split across GPUs, how you survive a weeks-long run, how big to train, what data, which alignment method, how far you stretch the context, how you serve — is the answer to one question: what is the most capable model this budget can buy? Efficiency is not a chapter; it is the spine.

Who this is for

You can read Python and PyTorch, you know what a neural network and gradient descent are, and you want to understand the full LLM pipeline well enough to reason about — or reproduce — any single decision in it. No prior systems background assumed; the GPU, the cluster, and the math are all built up from first principles. By the end you should be able to look at any frontier-model choice and say which constraint forced it.

The spine — one forced chain

Read the diagram left to right, top row then bottom. Each box is a wall; the arrow into the next box is the idea it is forced to introduce. Nothing is here "because LLMs have it" — each link exists because, without it, the chain breaks.

The lessons

The efficiency game

The premise: a fixed compute budget, C ≈ 6ND, paid in GPU-hours. The five frontiers and the wall→fix→new-wall chain of the whole course. Interactive: split a fixed C between N and D and watch a toy loss bottom out near ~20 tokens/param.

Part I · Basics

Tokenization — BPE from raw bytes

A model trains on integers, not text. Byte-level BPE: start from 256 UTF-8 bytes, greedily merge the most frequent pair, repeat to vocab V. The vocab-size trade-off, compression ratio, and the pathologies (digits, glitch tokens, "count the r's"). Interactive: step BPE merges and watch the sequence shrink.

Part I · Basics

The Transformer, from scratch

A parallel, differentiable next-token predictor: the residual stream, scaled dot-product attention (and why the √d), the causal mask derived from the autoregressive factorization, multi-head split/merge, the MLP, pre-norm residuals, weight tying. Parameter count ≈ 12·L·d². Interactive: sliders → every intermediate shape and the param/FLOP breakdown.

Part I · Basics

The modern recipe — what changed since GPT-2

Five swaps that each buy efficiency or stability at ≈zero cost: pre-norm + RMSNorm, RoPE (with the relative-distance proof), SwiGLU, GQA, and bias-free / QK-norm. Interactive: toggle each and watch param count, KV-cache size, stability, and usable context move.

Part I · Basics

Training — AdamW, schedule, stability

Move θ down the loss without diverging: AdamW with decoupled decay, warmup → cosine schedule (why warmup), init, gradient clipping, bf16 mixed precision (why bf16 over fp16), μP hyperparameter transfer, and an fp8-training outlook. Interactive: drag the schedule and clip, watch a toy run spike or converge.

Part I · Basics

Resource accounting — FLOPs and the memory wall

The two back-of-envelope models: compute (derive 6ND) and memory (the 16-bytes/param AdamW rule + activations + checkpointing). The hinge: a 7B's training state overflows one 80 GB GPU. Interactive: a memory calculator that turns red on overflow.

Part I · Basics

GPUs — the hardware you're renting

The memory hierarchy (HBM ↔ SRAM ↔ registers) and tensor cores; the master concept of arithmetic intensity and the roofline: memory-bound vs compute-bound. Why most LM ops waste the chip. Interactive: a roofline explorer.

Part II · Systems

Kernels — fusion and FlashAttention

One kernel per op means one HBM round-trip per op; fusion fixes it. FlashAttention: the online-softmax recurrence that tiles attention in SRAM and never materializes the T×T scores — O(T) memory, long context. Interactive: naive O(T²) vs Flash O(T) memory.

Part II · Systems

Data parallelism — DDP → ZeRO → FSDP

DDP replicates and all-reduces gradients (scales throughput, not memory). ZeRO/FSDP shard the optimizer state, gradients, and parameters across data-parallel ranks, cutting per-GPU memory ~G×. Interactive: pick the ZeRO stage, watch per-GPU memory fit or overflow.

Part II · Systems

Model parallelism — tensor, pipeline, the 3D mesh

When the model itself won't fit: tensor parallelism (split the matmuls, one all-reduce/layer, intra-node), pipeline parallelism (stages + micro-batches, the bubble), sequence parallelism, composed into a DP×TP×PP mesh. Interactive: a 3D-mesh planner.

Part II · Systems

Mixture of Experts — parameters without the FLOPs

Capability tracks parameters, but you pay FLOPs only for active ones. Route each token to its top-k experts: params grow E×, FLOPs stay flat. Load balancing, capacity, expert parallelism / all-to-all. Interactive: MoE economics — the total-vs-active gap.

Part II · Systems

Running the job — surviving a multi-week run

You can build the distributed machine; now run it for weeks on thousands of GPUs where hardware fails and the loss spikes. Checkpoint/resume (params + optimizer + RNG + data cursor), deterministic sharded loading, fault tolerance, and the spike runbook. Interactive: a run-survival simulator — find the checkpoint-interval sweet spot.

Part II · Systems

Scaling laws — Chinchilla-optimal

One shot at the big run, so predict: loss is a power law in N, D, C (empirical form, honestly flagged). The Chinchilla IsoFLOP result — ~20 tokens/param — and why you deviate from it for inference cost. Interactive: an IsoFLOP explorer that pays back lesson 00's widget.

Part III · Scaling

Data — the pipeline that decides quality

Compute-optimal D is worthless on garbage tokens. The pipeline: source → extract → filter → deduplicate → decontaminate → mix; plus synthetic data and the data wall. Why dedup and mixing dominate quality. Interactive: a data funnel from raw crawl to usable tokens vs your D target.

Part IV · Data & evaluation

Evaluation — measuring what you can't see

Perplexity vs benchmarks (MMLU/HellaSwag/GSM8K/HumanEval) and how each is scored; the eval crisis — contamination, prompt sensitivity, gameability, Goodhart. Interactive: a contamination inflator showing how a small leak buys big fake skill.

Part IV · Data & evaluation

Long context — extending the window

You trained at a short window to respect the activation budget; production needs 10–100× more. Extend it after the fact: RoPE position interpolation / NTK / YaRN + a short continued-training phase, served with FlashAttention and sequence parallelism. Interactive: push the window and watch naive RoPE collapse while YaRN holds.

Part V · Alignment & inference

SFT — from completer to assistant

The base model maximizes p(text); it has the knowledge but not the behavior. SFT keeps the exact loss but masks it to response tokens, wrapped in a chat template; plus LoRA / QLoRA parameter-efficient tuning. Interactive: a loss-mask visualizer — unmask the prompt and watch it degrade.

Part V · Alignment & inference

Preferences & rewards — RLHF, DPO, GRPO

SFT imitates; it can't say one answer is better. Optimize a preference/reward signal under a KL anchor: RLHF (reward model + PPO), DPO (the closed-form shortcut), GRPO/RLVR (group baseline + verifier); plus a candid look at safety & red-teaming. Interactive: the KL–reward trade-off.

Part V · Alignment & inference

Inference — serving the model you paid for

Prefill is compute-bound; decode is memory-bound — and decode is the cost. The KV cache as the bottleneck (GQA, then MLA compression), continuous batching, speculative decoding, quantization. Interactive: decode arithmetic — the bandwidth wall and the KV cap.

Part V · Alignment & inference

The whole budget — bytes → served assistant

One concrete budget walked end to end: every earlier decision shown as a forced consequence of the last. The three motifs restated; what's not covered and where it lives in the library. Interactive: an end-to-end budget flow from compute to served tokens/s.

Part VI · Synthesis

How to use this

Read it linearly. Each lesson assumes the last. Lesson 08 only makes sense as the answer to the wall lesson 05 hits; lesson 17 (preferences) only as the thing lesson 16 (SFT) cannot do.
Touch every knob. Each widget has a configuration that breaks — a run that diverges, a model that overflows the GPU, a leak that fakes a benchmark, a context that collapses. Find it. The failure is the lesson.
Follow the cross-links. This track is a hub. Where another series goes deeper — kernels, RL, serving — it links out instead of repeating.

Companion tracks

This course is the from-scratch sibling of Mini GPT. It links into the tracks it overlaps rather than duplicating them: kernels & parallelism → GPU Kernels and Triton; alignment → Reinforcement Learning; serving → vLLM and SGLang; compression → Knowledge Distillation; multimodal generation → Generative Continuous.