Reading Stanford CS336 Language Modeling from Scratch
A lecture-by-lecture reading of Stanford CS336: tokenizers, data, scaling, training, parallelism, evaluation, and alignment across the full language-model pipeline.
CS336 is Stanford's course on building a language model from scratch: five assignments cover the tokenizer and Transformer, training-systems optimization, scaling laws, data processing, and alignment and reasoning RL, with no code skeleton, only tests. The lectures, recordings, all five assignment repos and the leaderboards are public; from Assignment 2 on you supply your own compute, and the course-sponsored compute and grading are not available.
CS336's first lecture does not treat building a language model from scratch as reenacting every old technique. It separates mechanics, mindset, and intuitions, then uses BPE to show how raw bytes become trainable tokens.
Lecture 2 reduces model training to tensors, FLOPs, bytes, and time: use einops to track dimensions, arithmetic intensity and roofline analysis to identify bottlenecks, then trade compute for memory with gradient accumulation and activation checkpointing.
Lecture 3 does not turn its survey of modern LLMs into a single best recipe. It finds a conservative consensus—pre-norm, RMSNorm, no biases, SwiGLU, and RoPE—plus a small set of deviations justified by inference cost or stability.
Lecture 4 studies two kinds of sparsity: linear/recurrent attention reduces sequence-length cost, while MoE activates only part of a model for each token. Both turn saved FLOPs into routing, balancing, communication, and kernel problems.
Lecture 5 explains GPUs through SMs, warps, and the memory hierarchy, then unifies common optimization under low precision, fusion, recomputation, coalescing, and tiling. FlashAttention combines those principles for attention.
Lecture 6 turns GPU principles into kernels: benchmark scaling across shapes, profile actual calls and time, then implement GeLU, softmax, reductions, and tiled matrix multiplication in Triton. Speed begins with measuring correctly.
Lecture 7 starts below FSDP APIs, building a communication language from broadcast, all-reduce, all-gather, reduce-scatter, and all-to-all before assembling data, tensor, and pipeline parallelism.
Lecture 8 moves from parallel primitives to system design: ZeRO progressively shards optimizer state, gradients, and parameters; TP, PP, SP, and EP split width, depth, sequence, and experts. Their composition must follow topology and dynamic activation memory.
Lecture 9 begins with log-log linear relationships between data and error, then uses scaling laws to compare architectures, optimizers, batches, and model-data allocations. The Chinchilla dispute shows how fitting methods, observed ranges, and deployment objectives change the answer.
Lecture 10 separates prefill from decode: prefill parallelizes and is often compute-bound, while decode is sequential and commonly bandwidth-bound. GQA/MLA, quantization, speculative decoding, continuous batching, and PagedAttention reshape that cost.
Lecture 11 reads public recipes from MiniCPM, DeepSeek, Qwen, and Llama 3: hold most architectural ratios fixed, sweep learning rate and batch at small scale, then choose model/data allocation with IsoFLOPs. μP helps, but normalization, optimizers, and weight decay can break transfer.
Lecture 12 moves from perplexity to exams, chat preferences, agents, reasoning, and safety. Every benchmark changes the capability definition, scaffold, judge, and contamination risk, so evaluation must first say whether it compares a method, model, or complete system.
Lecture 13 traces training sources through Common Crawl, Wikipedia, GitHub, arXiv, books, and open datasets. Technically accessible is not the same as licensed, and raw data is not training data; provenance must precede cleaning and mixing.
Lecture 14 moves raw documents through language, quality, and safety filtering; exact and near deduplication; and source mixing. Each stage reshapes model behavior, while synthetic instruction and agent trajectories extend the pipeline into executable environments.
Lecture 15 divides post-training into imitation and optimization. SFT extracts pretrained capabilities from instruction-response data; RLHF uses pairwise feedback to bridge demonstrations and preferences. PPO and DPO both inherit data bias, reward overoptimization, and mode collapse.
Lecture 16 moves from PPO to GRPO and RLVR. Math, code, and environment outcomes provide scalable rewards and avoid some preference-model overoptimization, but group-normalized advantages introduce difficulty and length bias while rollout infrastructure becomes the dominant cost.
Lecture 17 organizes CLIP/SigLIP, LLaVA, Qwen-VL, and Chameleon into three paths: contrastive encoders learn semantics, vision-encoder/projector/LM stacks provide understanding, and discrete image tokens enable generation. Resolution, token budgets, and modality balance constrain them all.