Guus Bouwens
← Writing & notes

FROM THE NOTEBOOK

Language Modeling from Scratch - Stanford CS336

Stanford CS336 Language Modeling from Scratch I 2025 - YouTube

This course provides a rigorous, end-to-end masterclass in building modern large language models from first principles. Eschewing high-level APIs, the course philosophy is that "to understand it, you have to build it," forcing a deep, mechanical understanding of the entire LLM pipeline. The curriculum begins with the foundational components, including tokenization (BPE) and the modern Transformer architecture (Pre-Norm, RMS Norm, SwiGLU). It then dives deep into the systems-level challenges of achieving computational efficiency, covering single-GPU optimization through custom CUDA/Triton kernels (culminating in FlashAttention) and multi-GPU parallelism (Data, Tensor, and Pipeline parallelism). With the "how" established, the course explores the "what," covering the science of scaling laws to predictably trade-off compute, model size, and data, guided by the principles of the Chinchilla paper. The focus then shifts to the crucial, often-overlooked role of data, detailing the massive curation effort of filtering and deduplicating web-scale corpora like Common Crawl. Finally, the course covers the post-training process of making models useful and safe through alignment, starting with Supervised Fine-Tuning (SFT) and Reinforcement Learning from Human Feedback (RLHF) using algorithms like DPO, and progressing to the frontier of training advanced reasoning capabilities through RL from verifiable rewards using algorithms like GRPO. The course concludes by critically examining the complex and evolving field of evaluation, equipping students not just to use LLMs, but to innovate and push the fundamental boundaries of AI research.

L1: Overview and Tokenization

The course is introduced by the instructors, Percy Liang and Tatsu Hashimoto, and the course staff. They frame the class with a core philosophy: "to understand it, you have to build it." The motivation for the course stems from a perceived "crisis" in AI research, where researchers are becoming increasingly disconnected from the underlying technology by simply prompting proprietary, closed-off models. This class is designed to provide a deep, foundational understanding of the entire language model pipeline—from data processing and systems to modeling—to enable fundamental research.
The Challenge: Industrialization of Language Models
The instructors acknowledge the immense challenge posed by the scale of modern AI. Frontier models like GPT-4 are rumored to have trillions of parameters, cost hundreds of millions of dollars to train, and are developed in secret. Since students cannot replicate this, the course focuses on building smaller models. However, this raises the issue that small models may not be representative due to phenomena like emergent behaviors, where capabilities like in-context learning only appear at massive scale.
What Can Be Taught: Mechanics, Mindset, and Intuitions
Despite the scale gap, the course aims to impart three crucial types of knowledge:

  1. Mechanics: The raw ingredients of how things work, such as the Transformer architecture and model parallelism.
  2. Mindset: A critical focus on scaling and squeezing the maximum possible performance out of hardware. This "scaling mindset" is identified as the key driver behind the recent AI breakthroughs.
  3. Intuitions: (Partially) An understanding of which data and modeling decisions lead to good models, while acknowledging that these intuitions may not always transfer from small to large scales.

The "Bitter Lesson" Re-Framed: Algorithms at Scale
The lecture reinterprets Richard Sutton's "bitter lesson," arguing it doesn't mean "scale is all that matters." Instead, the key is "algorithms at scale." Efficiency becomes even more critical at large scales where waste is unaffordable. The instructors highlight that algorithmic improvements have outpaced Moore's Law, demonstrating their importance. The central question of the course is therefore: "What is the best model one can build given a certain compute and data budget?"
Course Structure and Logistics
The course is a demanding five-unit class with a heavy workload. Assignments require students to build systems from scratch, starting with blank files but guided by unit tests. The five core units of the course are:

  1. Basics: Implementing a full, basic training pipeline, including a tokenizer, a Transformer model, and a training loop.
  2. Systems: Optimizing the pipeline for hardware. This involves writing low-level GPU kernels using Triton, implementing parallelism (data and model parallelism), and speeding up inference using techniques like speculative decoding.
  3. Scaling Laws: Learning to use small-scale experiments to predict performance at a large scale, allowing for the optimal allocation of a fixed compute budget (e.g., "Chinchilla optimal" rules for balancing model size and data amount).
  4. Data: A deep dive into data acquisition and curation. This covers processing raw internet data (like Common Crawl, which is mostly "trash"), filtering, deduplication, and navigating legal and ethical questions.
  5. Alignment: The process of making a pre-trained "base model" useful and safe. This includes Supervised Fine-Tuning (SFT) to teach the model instruction-following and learning from feedback (using Reinforcement Learning techniques like DPO) to align it with human preferences.

Technical Deep Dive: Tokenization
The lecture concludes with the first technical topic: tokenization, the process of converting raw text into a sequence of integers (tokens) for the model.

The lecture ends by previewing the next session, which will focus on PyTorch and how to precisely account for computational resources (FLOPs).

L2: Pytorch, Resource Accounting

The lecture begins by framing the core theme: efficiency through resource accounting. The instructor poses practical, real-world questions, such as "How long would it take to train a 70B model on 1,024 H100s?" and "What's the largest model you can train on 8 H100s?" The goal is to equip students with the ability to perform these "napkin math" calculations. This mindset is crucial because at scale, compute and memory usage directly translate into millions of dollars. The lecture focuses on the foundational PyTorch primitives and the accounting necessary to understand these costs, deferring a conceptual overview of the Transformer to the next session.
Part 1: Memory Accounting - Tensors and Data Types
The discussion starts with the fundamental building block of deep learning: the tensor. The memory a tensor occupies is determined by its size and its data type.

This leads to the concept of mixed precision training, a strategy where different data types are used for different parts of the model pipeline—for example, using fast bfloat16 for forward passes while keeping the more stable float32 for parameters and optimizer states.
Part 2: Compute Accounting - FLOPs, Matmuls, and MFU
The focus shifts from memory to computation, measured in FLOPs (Floating Point Operations).

Part 3: Gradients and the "Rule of Six"
The analysis of compute cost is extended to the training process, specifically the backward pass for calculating gradients.

Part 4: Building the Model - From Initialization to a Training Loop
The lecture puts all the pieces together by walking through the components of a simple model.

The lecture concludes by reiterating that these fundamental concepts of memory and compute accounting are the essential prerequisites for efficiently building and training any large-scale model, including the Transformer that students will implement in their first assignment.

L3: Architectures, Hyperparameters

The lecture begins by setting a practical, data-driven theme: instead of deriving architectures from first principles, the best way to understand modern language models is to learn from the collective experience of those who have already built them. The approach is an "evolutionary analysis" of dozens of language models released between 2017 and 2025, identifying which architectural choices have survived, which have been discarded, and which have become the consensus.
Part 1: Core Architecture Variations - The "LLaMA-like" Consensus
The lecture systematically breaks down the components of the Transformer, contrasting the original design with the modern consensus that students are implementing.

  1. Normalization (Pre-Norm vs. Post-Norm): This is the first and most critical consensus. The original Transformer used Post-Norm (applying LayerNorm after the attention/MLP blocks), which was found to be unstable and required careful learning rate warm-up. The modern standard is Pre-Norm, where LayerNorm is applied to the input of each block before the operation. This leads to significantly more stable training and smoother gradient flow through the residual stream. A recent innovation noted is the "double norm" (both pre and post) used by models like Grok.
  2. RMS Norm vs. LayerNorm: Nearly all modern models have replaced the standard LayerNorm with RMS Norm. RMS Norm is a simpler variant that omits the mean-centering and bias terms. The key insight is that while the FLOPs savings are negligible, RMS Norm is significantly faster because it reduces memory movement, a major bottleneck in GPU performance.
  3. No Bias Terms: In line with the move to RMS Norm, most modern models omit bias terms from their linear layers. This simplifies the architecture, reduces memory traffic, and has been empirically shown to improve training stability.
  4. Activations (The Rise of GLU): While early models used ReLU or GeLU, the overwhelming modern consensus is to use Gated Linear Units (GLUs), most commonly SwiGLU. Gating is a consistently powerful architectural motif, and GLU variants empirically provide better performance (lower loss) than their non-gated counterparts.
  5. Serial vs. Parallel Layers: The standard Transformer is serial, processing attention and then the MLP block sequentially. Some models (like PaLM) use parallel layers, computing both simultaneously to gain systems-level efficiency. However, the serial approach remains the more dominant choice.

Part 2: Demystifying Hyperparameter Choices
The lecture moves from architectural components to the specific numbers that define them, revealing that many choices follow established rules of thumb.

  1. Feed-Forward (FFN) Dimension: The hidden dimension of the MLP (d_ff) is conventionally set to 4 * d_model for non-gated activations. For GLUs, this is scaled down to 8/3 * d_model to maintain a similar parameter count.
  2. Attention Heads: The standard practice is to set the dimension per head (d_head) such that num_heads * d_head = d_model, keeping the total attention output dimensionality consistent with the residual stream.
  3. Aspect Ratio (Depth vs. Width): There is a clear sweet spot for the ratio of model width to depth (d_model / n_layers), which is consistently around 128. Scaling law studies show this ratio is optimal across different model sizes.
  4. Vocabulary Size: The trend is towards larger vocabularies. While early English-only models used 30k-50k tokens, modern multilingual models use 100k-250k to better handle diverse languages, code, and special characters.
  5. Regularization (A Counter-Intuitive Finding): In the pre-training regime where models rarely see data more than once, dropout has fallen out of fashion as overfitting is not the primary concern. However, weight decay is still widely used, not for regularization, but because it improves optimization dynamics. It interacts with learning rate schedules (like cosine decay) to achieve a better final training loss.

Part 3: Recent Developments - Stability, Inference, and Long Context
The final section covers the latest innovations, driven by the challenges of training at massive scale and improving model usability.

  1. Stability Tricks: As models get larger, training becomes prone to "loss spikes" and instability. The primary culprit is the softmax function. Key interventions include:
    • z-loss: An auxiliary loss that penalizes the softmax normalizer to keep it numerically stable.
    • QK Norm: Applying a LayerNorm to the Query and Key vectors before the dot product to bound the inputs to the softmax.
  2. Inference Optimization (MQA & GQA): During inference, the KV Cache (storing past keys and values) becomes a memory bandwidth bottleneck. To solve this:
    • Multi-Query Attention (MQA): Uses multiple heads for queries but shares a single head for keys and values.
    • Grouped-Query Attention (GQA): A compromise with fewer key/value heads than query heads.
    • Both techniques drastically reduce the size of the KV Cache, significantly speeding up inference.
  3. Long Context: The newest models achieve massive context windows by mixing attention patterns. For instance, they might alternate blocks of local sliding window attention (which is computationally cheaper and uses RoPE) with occasional blocks of full self-attention (often without position embeddings) to allow long-range information to propagate across the entire sequence.

L4: Mixture of experts

The lecture begins by establishing that Mixture of Experts (MoE) is no longer a niche topic but the core architecture behind most modern, high-performance language models, including GPT-4, Grok, and Llama 4. The consensus is that for a given amount of compute (FLOPs), a well-designed MoE model will outperform a traditional dense model.
The Core Idea: More Parameters for the Same FLOPs
The concept of MoE is simple but powerfully effective. Instead of having a single, large feed-forward network (FFN or MLP) in each Transformer block, an MoE architecture replaces it with:

  1. A lightweight router.
  2. Multiple smaller FFNs, called experts.

For each token, the router sparsely activates a small number (top-K, often just 2) of these experts. The key benefit is that the model can have a massive number of total parameters (by having many experts), but the computational cost (FLOPs) for any given forward pass is only determined by the small number of activated parameters. This allows MoEs to achieve better performance for the same training compute budget, a fact supported by numerous empirical studies.
Key Design Questions for Building an MoE
The lecture breaks down the architecture into three main design challenges:

  1. Routing: How are tokens assigned to experts?
  2. Configuration: How many experts should there be, and how large should they be?
  3. Training: How do you train the non-differentiable router effectively?

1. Routing: The Consensus is "Token Choice Top-K"
While early research explored various routing mechanisms (token choice, expert choice, RL, etc.), the modern consensus is token choice top-K routing. The process works as follows:

2. Expert Configuration: Fine-Grained and Shared Experts
A key innovation, pioneered by DeepSeek, was to move beyond simply copying the dense FFN.

3. Training: The Gnarly Challenge of Balancing
The biggest difficulty in training MoEs is their inherent instability. Because the routing decision is discrete (a hard top-K selection), training is a challenging "RL-ish" problem. Without intervention, models quickly suffer from expert collapse, where the router learns to send all tokens to just one or two experts, leaving the others unused and wasting their parameters.
The standard solution is not complex RL but a simple heuristic: the balancing loss. This is an auxiliary loss term added during training that penalizes the model if the distribution of tokens across experts becomes uneven. It encourages the router to spread the load, ensuring all experts are utilized. Modern models use multiple balancing losses to manage load at different levels (per-expert, per-device, and even per-sequence).
Systems, Stability, and Other Tricks

Case Study: The Evolution of the DeepSeek MoE
The lecture concludes by tracing the architectural evolution of the high-performance DeepSeek MoE models, showing how these concepts culminated in a state-of-the-art system.

L5: GPUs

The lecture begins by posing a puzzle: why does the performance of a GPU on matrix multiplication tasks vary in strange, wave-like patterns instead of improving smoothly with problem size? The goal of the session is to demystify this behavior and understand the low-level hardware principles that govern GPU performance. This knowledge is essential for building new, highly efficient algorithms like FlashAttention, which makes long-context transformers possible.
Part 1: GPU Architecture and Execution Model
To understand performance, one must first understand the hardware. The lecture outlines the fundamental differences between CPUs and GPUs and the key components of a GPU.
CPU vs. GPU: CPUs are optimized for low latency (finishing a single task quickly), while GPUs are designed for high throughput (finishing many tasks in aggregate as quickly as possible).
Anatomy: A GPU consists of many Streaming Multiprocessors (SMs), which are like individual cores. Each SM contains many Streaming Processors (SPs) and specialized hardware (like tensor cores for matrix multiplication) that execute instructions in parallel.
The Memory Hierarchy: This is the most critical concept. GPUs have a hierarchy of memory with vastly different speeds:
Fast, On-Chip Memory: Registers and L1/Shared Memory are located directly inside each SM and are extremely fast (e.g., ~20 clock cycles).
Slow, Off-Chip Memory: Global Memory (HBM/DRAM) resides outside the main GPU chip and is much slower to access (e.g., ~200-300 clock cycles). Minimizing access to this slow memory is the single most important goal of GPU programming.
Execution Model: Code on a GPU is executed in a three-level hierarchy:
Blocks: Large groups of threads, with each block assigned to run on a single SM.
Warps: Within a block, threads are executed in groups of 32, called warps.
Threads: All 32 threads in a warp execute the same instruction simultaneously, but on different data (a model called SIMT - Single Instruction, Multi-Thread).
Part 2: Principles of GPU Performance Optimization
With the architecture established, the lecture details the key techniques for writing fast code, framed as a set of "tricks" to overcome the memory bottleneck.
Lower Precision: Using 16-bit formats (like bfloat16) instead of 32-bit float32 halves the amount of data that needs to be moved from slow memory, effectively doubling the memory bandwidth for memory-bound operations.
Operator Fusion: This technique avoids unnecessary round-trips to global memory. Instead of performing one operation, writing the result to global memory, and then reading it back for the next operation, fusion combines multiple sequential operations into a single GPU kernel. This keeps intermediate data in fast, on-chip memory.
Recomputation: A strategy that trades abundant compute for scarce memory bandwidth. Instead of storing an activation during the forward pass (which requires a slow write and read from global memory), it can be cheaper to simply recompute it on the fly during the backward pass.
Memory Coalescing: The GPU's slow DRAM reads data in large, fixed-size chunks ("burst mode"). If the 32 threads in a warp access memory locations that are close together and aligned, the hardware can "coalesce" these into a single, efficient memory transaction. Un-aligned or scattered access patterns break this, dramatically slowing down performance.
Tiling: This is arguably the most important technique. Large problems, like matrix multiplication, are broken down into smaller "tiles" that are small enough to fit entirely within an SM's fast shared memory. The algorithm loads a tile from slow global memory once, performs all possible computations on it within fast memory, and only then moves to the next tile. This drastically reduces the total number of slow global memory accesses.
Part 3: Solving the Mystery Plot
The lecture revisits the initial "mystery plot" of matrix multiplication performance and explains its strange features using the principles just covered:
The initial sharp rise in performance follows the Roofline Model, where performance is limited by memory bandwidth until the problem is large enough to become compute-bound.
The different performance "levels" or "plateaus" are caused by tiling alignment. Matrix sizes that are perfectly divisible by the tile size (e.g., multiples of 32 or 64) are highly efficient. Those that are not lead to underutilized tiles and wasted computation.
The sharp, sudden drops in performance are caused by wave quantization. This occurs when the number of tiles slightly exceeds the number of available SMs on the GPU. This forces the GPU to execute a second, highly inefficient "wave" of computation with very few active SMs, tanking the overall throughput.
Part 4: Case Study - FlashAttention
The lecture concludes by showing how these principles are combined in FlashAttention, a state-of-the-art algorithm for accelerating transformers.
The Goal: The innovation of FlashAttention is not making the computation sub-quadratic in sequence length (it's still O(N²)), but making the number of slow global memory accesses sub-quadratic.
The Problem: The softmax operation in attention is global—it requires summing over an entire row of the attention matrix. This breaks the standard tiling approach, as you can't compute the normalizer for one tile without seeing all the others. The naive approach would require materializing the entire N x N attention matrix in slow global memory, which is the key bottleneck.
The Solution: FlashAttention uses two main tricks:
Online Softmax: It employs a clever algorithm that can compute the softmax incrementally, tile-by-tile, by maintaining and updating running statistics. This allows the entire attention computation to happen without ever writing the full N x N matrix to slow memory.
Recomputation: For the backward pass, it avoids storing the N x N attention matrix by recomputing it on the fly from the inputs, again trading compute for memory.
By combining tiling and recomputation, FlashAttention overcomes the memory bottleneck of the attention mechanism, enabling much longer context lengths.

L6: Kernels, Triton

The lecture begins by setting the stage for Assignment 2, which involves writing high-performance GPU kernels, including a version of FlashAttention-2, using Triton. The primary learning goals are to demystify GPU performance—explaining phenomena like the strange, wave-like performance patterns of matrix multiplies—and to equip students with the skills to accelerate their own algorithms using CUDA.
Part 1: GPU Review and the Importance of Profiling
A brief review of GPU architecture is provided to refresh key concepts:

The most critical high-level takeaway is the necessity of benchmarking and profiling. Benchmarking measures end-to-end runtime, while profiling provides a fine-grained breakdown of where time is spent. The instructor stresses that without profiling, optimization efforts are often misdirected.
Part 2: Practical Benchmarking and Profiling
The lecture walks through practical techniques for measuring performance in PyTorch.

Part 3: Writing High-Performance Kernels
The core of the lecture is a hands-on demonstration of writing and optimizing a kernel for the GeLU activation function. The goal is to match the performance of PyTorch's highly optimized, fused kernel.

  1. Naive PyTorch Implementation: A manual implementation of the GeLU formula using standard PyTorch operations (*, +, tanh) is shown to be extremely slow. Profiling reveals it launches multiple separate CUDA kernels, causing excessive and unnecessary data movement between the SMs and slow global memory.
  2. Handwritten CUDA C++ Kernel: The first attempt at fusion involves writing a raw CUDA C++ kernel. The code demonstrates how to calculate a thread's global index from its block and thread IDs, perform the computation, and write the result. This single, fused kernel is significantly faster than the naive version, getting much closer to the optimized PyTorch implementation.
  3. Triton Kernel: Next, the same kernel is written in Triton, a Python-based DSL for GPU programming. The Triton code is cleaner and operates at a higher level of abstraction (programming at the block level rather than the thread level), letting the Triton compiler handle low-level details like memory coalescing. The performance is comparable to the handwritten CUDA kernel. A dive into the compiled PTX assembly code reveals how Triton intelligently loads data in aligned chunks to maximize memory bandwidth.
  4. : Finally, the lecture demonstrates the power of modern JIT compilers by simply applying torch.compile to the original, naive PyTorch code. The compiler automatically performs kernel fusion, generating a highly optimized Triton kernel under the hood that is even slightly faster than the manually written ones.

The key takeaway is that for simple element-wise operations, torch.compile is often sufficient. However, understanding how to write custom kernels is essential for more complex, novel operations where the compiler may not find the optimal solution (like FlashAttention).
Part 4: A More Complex Kernel - Softmax
The lecture concludes by extending these ideas to a more complex operation: softmax, which involves a reduction (summing across a row). A simple Triton kernel is implemented where each row of the input matrix is assigned to a single thread block. Within the block, the kernel loads the entire row into fast shared memory, computes the max and the sum for normalization, and then writes the final result back to global memory. This approach effectively contains the entire reduction within a single, fused kernel, achieving performance competitive with the native PyTorch and torch.compile versions.

L7: Parallelism 1

The lecture begins by establishing why parallelism across multiple machines is essential. While single GPU performance is scaling rapidly, it's not enough to keep pace with the demands of training state-of-the-art language models. Multi-machine training is necessary for two primary reasons:

  1. Compute Scaling: To train models faster by leveraging the aggregate compute power of a data center.
  2. Memory Scaling: To train models that are too large to fit in the memory of a single GPU.

The lecture introduces the concept of the communication hierarchy: connections within a single machine (e.g., NVLink) are extremely fast, while connections across machines are significantly slower. This hardware reality dictates the strategies for parallelization.
Part 1: Parallelization Strategies
The core of the lecture is a breakdown of the three fundamental types of parallelism, each designed to shard a different component of the training process.
1. Data Parallelism (Sharding the Batch)

2. Model Parallelism (Sharding the Model)
This approach is used when a model is too large to fit on a single GPU, even with FSDP. It involves splitting the model's parameters across GPUs.

3. Activation Parallelism (Sharding the Activations)
Even with model parallelism, activations can become a memory bottleneck, especially with long sequences.

Part 2: Putting It All Together - The Rule of Thumb for 3D Parallelism
The lecture concludes by providing a practical, hierarchical rule of thumb for combining these strategies to train a massive model:

  1. First, use Tensor Parallelism: Scale up to the number of GPUs available within a single machine (typically up to 8). This is the most bandwidth-hungry and is best suited for fast, intra-node connections.
  2. Next, use Pipeline Parallelism or ZeRO-3 (FSDP): If the model still doesn't fit in memory, use one of these strategies to shard the model across machines until it fits.
  3. Finally, use Data Parallelism: Once the model fits, use data parallelism to scale out across all remaining GPUs to maximize compute throughput. Data parallelism is the most tolerant of slower, inter-node network links.

The lecture ends with case studies of recent large models (like Llama 3 and DeepSeek-V3), showing how they apply these principles in practice, while also highlighting the real-world challenges of large-scale training, such as frequent hardware failures and the need for fault-tolerant infrastructure.

L8: Parallelism 2

The lecture begins by framing the challenge of multi-GPU and multi-node training. It recaps the hardware and communication hierarchy, emphasizing the vast differences in speed between on-chip memory (L1 cache), on-GPU memory (HBM), fast intra-node links (NVLink), and slower inter-node connections. The central problem remains the same as in single-GPU optimization: data transfer is the bottleneck. The goal is to design parallel computation strategies that minimize slow communication and keep the powerful GPU compute units saturated.
Part 1: The Building Blocks of Distributed Training
This section covers the fundamental primitives and software layers used for multi-GPU communication.

  1. Collective Operations Review: The lecture revisits the core collective communication primitives (broadcast, scatter, gather, reduce, all_gather, reduce_scatter). These operations are the high-level language used to describe data movement between a group of devices (the "world"). The key identity that all_reduce is equivalent to reduce_scatter followed by all_gather is re-emphasized.
  2. Hardware and Software Stack:
    • Hardware: Modern GPU nodes bypass the slow CPU bus with direct, high-speed interconnects like NVLink (for communication within a node) and NVSwitch (for communication across nodes).
    • NCCL (NVIDIA Collective Communications Library): A low-level NVIDIA library that implements these collective operations as highly optimized CUDA kernels, intelligently using the underlying hardware topology.
    • : The PyTorch library that provides a clean, Python-native interface for these collective operations. It acts as a high-level wrapper around backends like NCCL (for GPUs) or Gloo (for CPUs), making distributed code more portable.
  3. Benchmarking Collectives: The lecture demonstrates how to benchmark these operations in PyTorch, showing how to measure the effective bandwidth of an all_reduce. This highlights that real-world performance depends on many factors and must be measured empirically.

Part 2: Implementing Distributed Training Strategies in Code
The main part of the lecture is a practical, code-driven walkthrough of the three primary parallelism strategies, using a simple multi-layer perceptron (MLP) as the example.

  1. Data Parallelism (DDP):
    • Concept: The model is replicated on each GPU, but the data batch is split or "sharded."
    • Implementation: The code shows that DDP is remarkably simple to implement from scratch. Each process runs a standard training loop on its local slice of data. The only addition is a single line of code after the backward() pass: an all_reduce operation to average the gradients across all GPUs. This ensures that every GPU performs the exact same parameter update, keeping the models synchronized.
  2. Tensor Parallelism:
    • Concept: The data is replicated on each GPU, but the model's weight matrices are split (horizontally).
    • Implementation: The code demonstrates that each GPU holds a "slice" of each layer's weight matrix. After each sharded matrix multiplication, the partial activations must be collected from all GPUs using an all_gather operation to reconstruct the full activation tensor before it can be fed into the next layer. This frequent communication of large activation tensors makes tensor parallelism extremely bandwidth-hungry.
  3. Pipeline Parallelism:
    • Concept: The model is split by layers (vertically). GPU 0 handles layers 1-8, GPU 1 handles layers 9-16, and so on.
    • Implementation: The code uses point-to-point communication (send and receive). Each GPU waits to receive an activation tensor from the previous GPU in the pipeline, processes it through its assigned layers, and then sends the result to the next GPU. The lecture notes that this simple implementation is naive and suffers from large "pipeline bubbles" (idle time), which real-world systems mitigate with complex techniques like micro-batching and overlapping communication and computation.

Conclusion and Higher-Level Abstractions
The lecture concludes by acknowledging that the presented implementations are simplified. Real-world libraries like Megatron-LM and FSDP involve far more complex bookkeeping and optimizations. A brief aside into the Jax ecosystem shows an alternative, more declarative approach where the user specifies the desired sharding strategy, and a powerful compiler automatically generates the complex communication and computation plan.
The final summary reinforces the core idea: all parallelization strategies are different ways of "cutting" the training problem along a specific dimension (data batch, model width, or model depth). They all navigate the fundamental trade-off between computation, memory storage, and communication, a challenge that will persist as long as models continue to push the boundaries of hardware.

L9: Scaling Laws 1

The lecture begins by framing a practical problem: if given a massive amount of compute (e.g., 100,000 H100s for a month), how would you design the best possible language model? Instead of simply copying existing architectures, the goal is to innovate. Scaling laws are presented as the core methodology for making principled, data-driven design decisions. The central idea is to train many small models to learn predictable relationships between resources (compute, data, model size) and performance, and then use these "laws" to extrapolate and optimize the design for a large-scale training run.
Part 1: The History and Context of Scaling Laws
The lecture first grounds the modern concept of scaling laws in a rich history from statistical machine learning theory.

Part 2: The Predictability of Data Scaling
The lecture focuses on the most intuitive type of scaling: data scaling.

Part 3: Model Scaling and Engineering Decisions
The discussion then shifts to using scaling laws to optimize model architecture and hyperparameters.

Part 4: The Chinchilla Scaling Laws - Trading Off Data and Model Size
This section covers the most influential application of scaling laws: determining the optimal balance between model size and the amount of training data for a fixed compute budget.

L10: Inference

The lecture begins by establishing that inference—the process of generating responses from a trained model—is a critical and increasingly dominant workload. It's not just for chatbots; it's essential for model evaluation, test-time compute (e.g., chain-of-thought reasoning), and even parts of training (like reinforcement learning). As AI products scale, the cumulative cost of inference often surpasses the one-time cost of training. Key performance metrics for inference include Time-To-First-Token (TTFT), latency (time per subsequent token), and throughput (total tokens generated per second).
The fundamental challenge of inference is its sequential nature. Unlike training, where computations can be parallelized across the entire sequence, inference must generate tokens one by one, making it difficult to fully utilize the GPU's compute power and rendering it memory-limited.
Part 1: Analyzing the Inference Workload
This section provides a detailed breakdown of the computational and memory costs of inference, using the concept of arithmetic intensity (FLOPs per byte of memory transferred).

  1. Prefill vs. Generation: Inference is a two-stage process:
    • Prefill: The initial processing of the prompt. This is fast, parallelizable, and compute-limited, similar to training.
    • Generation: The sequential, one-by-one generation of new tokens. This is slow and memory-limited.
  2. The KV Cache: To avoid recomputing the entire sequence for each new token, models use a KV Cache, which stores the key and value activations for all previously processed tokens. While this makes generation linear (O(T)) instead of quadratic (O(T²)), the KV Cache itself becomes the primary memory bottleneck.
  3. Arithmetic Intensity of the Transformer:
    • MLP Layers: The arithmetic intensity during generation scales with the batch size (B). To be compute-limited, a large number of concurrent requests (B > ~300 on an H100) is needed.
    • Attention Layers: The situation is worse for attention. Because each sequence has its own unique KV Cache, the arithmetic intensity is always low (approximately 1), regardless of batch size. This means the attention mechanism during generation is always memory-limited.
  4. Latency and Throughput Trade-off: A key trade-off emerges from this analysis. Small batch sizes yield low latency (good for individual user experience), while large batch sizes yield high throughput (good for overall system efficiency). However, the maximum batch size is limited by the GPU's memory, which is dominated by the KV Cache.

Part 2: Making Inference Faster - "Lossy" Shortcuts
The lecture then explores architectural changes that improve inference speed by reducing the size of the KV Cache, even if it means a small drop in accuracy.

  1. Grouped-Query Attention (GQA): Instead of having one key/value head for every query head (Multi-Head Attention), GQA uses fewer key/value heads that are shared across groups of query heads. This significantly shrinks the KV Cache, improving throughput and allowing for larger batch sizes.
  2. Multi-head Latent Attention (MLA): An alternative from DeepSeek that projects the keys and values into a much lower-dimensional "latent" space before caching them.
  3. Cross-Layer Attention (CLA): Shares the same KV projection weights across multiple layers, further reducing the parameter count and cache size.
  4. Local/Sliding Window Attention: Restricts each token to only attend to a small, local window of recent tokens. This keeps the KV Cache size constant, regardless of sequence length, but can harm long-range dependency modeling. The solution is to create hybrid models that interleave local attention layers with a few full, global attention layers.

Part 3: Radical Architectural Changes
This section discusses more fundamental departures from the standard Transformer architecture, driven by the need for faster inference.

  1. State-Space Models (SSMs) and Mamba: SSMs are inspired by control theory and have an RNN-like structure that allows for linear-time inference with a constant-size state (no growing KV Cache). While early versions struggled with language tasks, newer models like Mamba and hybrid architectures have become highly competitive.
  2. Diffusion Models: A completely different paradigm that generates all tokens in parallel and then iteratively refines them. This breaks the sequential bottleneck of autoregressive generation, leading to massive speedups in tokens per second, though their general capabilities are still under investigation.

Part 4: More Tricks - Quantization, Pruning, and Speculative Decoding

  1. Quantization: Reducing the numerical precision of the model's weights (e.g., from 16-bit floats to 8-bit or 4-bit integers). This shrinks the model size, reducing memory transfer time and speeding up inference.
  2. Model Pruning: Removing "unimportant" parts of the model (layers, heads) and then using distillation to repair the accuracy of the smaller, pruned model.
  3. Speculative Decoding: A "lossless" technique that provides a speedup without changing the model's output. It works by:
    • Using a small, fast "draft" model to generate a chunk of candidate tokens.
    • Using the large, slow "target" model to verify all these tokens in a single, parallel prefill step.
    • A statistical method (rejection sampling) is used to accept or reject the drafted tokens, guaranteeing that the final output distribution is identical to that of the large model alone.

Part 5: Systems-Level Optimizations for Live Traffic
The final section addresses the challenges of serving models in a real-world environment with dynamic, heterogeneous traffic.

The lecture concludes by reiterating that the inference "game" is much broader than just systems optimization. The most significant gains are now coming from radical architectural innovations that fundamentally sidestep the bottlenecks of the original Transformer design.

L11: Scaling Laws 2

The lecture begins by posing a practical challenge: given a massive compute budget (e.g., 100,000 H100s), how do you design the best possible language model? Instead of just copying existing models, the goal is to innovate. Scaling laws are the primary tool for this, allowing researchers to train many small models, learn predictable relationships, and then extrapolate these findings to confidently design a large model for a single, expensive training run. The lecture focuses on real-world case studies to answer the question: Does this approach, particularly the Chinchilla recipe, actually work in practice?
Part 1: Case Studies in Scaling
This section examines how different research groups have practically applied scaling laws in their model development.
Case Study 1: Cerebras-GPT

Case Study 2: MiniCPM

Case Study 3: DeepSeek LLM

Common Themes from Recent Models (Llama 3, etc.):

Part 2: A Deep Dive into μP (Maximal Update Parameterization)
The second half of the lecture provides a detailed mathematical derivation and empirical validation of μP.

Conclusion: Best Practices for Scaling in the Wild
The lecture concludes by summarizing the practical recipes for scaling large models. Researchers combine several key ingredients:

These case studies demonstrate that while the exact methods may differ, a principled, data-driven approach based on scaling laws is the universal methodology for successfully training state-of-the-art language models.

L12: Evaluation

The lecture begins by highlighting the current "evaluation crisis" in language models. While there are numerous benchmarks (MMLU, GPQA), leaderboards (Chatbot Arena), and anecdotal "vibes" shared on social media, there is no consensus on what these numbers truly mean or which methods are reliable. The core message is that evaluation is not a simple mechanical process but a profound one that shapes the entire field, as developers inevitably build models to optimize for the metrics being tracked.
Part 1: A Framework for Evaluation
The purpose of evaluation depends on the question being asked. A user making a purchase decision has different needs than a researcher trying to measure scientific progress or a developer trying to debug a model. A robust evaluation framework must consider four key components:

  1. Inputs (Prompts): Where do the prompts come from? Do they cover diverse use cases and difficult "tail" scenarios? Are they static or adapted to the model's responses (as in a chatbot)?
  2. Model Calls: How is the model prompted (zero-shot, few-shot, chain-of-thought)? What tools (e.g., code execution, RAG) does it have access to? Are we evaluating the raw model or an entire agentic system?
  3. Outputs: How are the model's responses assessed? Is there a clean, error-free reference? What metrics are used (e.g., pass@k for code, cost-efficiency), and how is open-ended generation judged?
  4. Interpretation: What does a score of "91" actually mean? Does it reflect true generalization or just train-test contamination? Are we evaluating a specific model, a system, or a research method?

Part 2: A Tour of Evaluation Benchmarks
The lecture then provides a survey of common benchmarks, grouping them by the capabilities they aim to measure.
1. Perplexity and Knowledge-Based Tasks:

2. Instruction-Following and Open-Ended Generation:

3. Agentic Tasks:
These benchmarks evaluate a model's ability to use tools and perform multi-step reasoning to solve complex problems.

4. Safety Benchmarks:

Part 3: The Challenges of Realism and Validity
The lecture concludes by addressing two of the deepest problems in evaluation.

  1. Realism: Most academic benchmarks consist of "quizzing" prompts (where the user knows the answer) and are not representative of real-world "asking" prompts (where a user genuinely needs help). Projects like MedHELM aim to create more realistic benchmarks by sourcing tasks directly from clinicians, but this often creates a tension between realism and data privacy.
  2. Validity (Train-Test Contamination): This is a massive, unsolved problem. Since modern models are trained on huge web scrapes, it is nearly impossible to guarantee that a benchmark's test set was not part of the training data. While decontamination techniques exist, they are imperfect. This issue fundamentally challenges the interpretation of benchmark scores.

The final takeaway is that there is a critical distinction between evaluating a method (which requires a controlled, academic setup) and evaluating a system (which is what users and industry care about). The rules of the game must be clearly defined to understand what a given evaluation is actually measuring.

L13: Data 1

The lecture begins with a bold claim: data is the most important component in building high-performance language models. This is evidenced by the fact that while top labs (like Meta and DeepSeek) openly share their model architectures, they are extremely secretive about their data sources and curation processes. This secrecy stems from both competitive dynamics and legal concerns. The process of creating high-quality data is a massive, parallelizable effort involving curation, cleaning, and filtering, making it a key differentiator.
The Stages of Data in Training
The lecture outlines a multi-stage process for how data is used, moving from vast, low-quality sources to smaller, high-quality, and specialized datasets.

  1. Pre-training: The initial phase where a "base model" is trained on trillions of tokens of raw, largely web-sourced text to learn general knowledge and linguistic patterns.
  2. Mid-training: A more recent, intermediate stage where the model is further trained on a curated, smaller set of high-quality documents to instill specific capabilities like coding, math, or long-context reasoning.
  3. Post-training (Alignment): The final phase where the model is fine-tuned on instruction-following or chat data and aligned for safety and helpfulness using techniques like Reinforcement Learning from Human Feedback (RLHF).

Part 1: Pre-training Data - The Foundation
This section provides a historical tour of the major datasets that have defined the pre-training era.

Part 2: Legal and Ethical Issues - Copyright
A significant portion of the lecture is dedicated to the complex issue of copyright.

Part 3: Mid-training and Post-training Data
The final section covers the smaller, more curated datasets used in later training stages.

The lecture concludes by reiterating that data curation is a messy, heuristic-driven process, but it is the primary lever for improving model capabilities and the main source of competitive advantage in the field.

L14: Data 2

Following the previous lecture's historical overview of pre-training datasets, this session takes a deep dive into the mechanics of data curation, focusing on the algorithmic primitives used for filtering and deduplication. The core challenge is to perform these operations efficiently at the scale of the entire web, which requires fast, scalable algorithms.
Part 1: Filtering Algorithms
The lecture frames filtering as a core task: given a small, high-quality target dataset (T) and a massive raw dataset (R) (like Common Crawl), the goal is to find a subset of R that is "similar" to T. This process must be both generalizable and extremely fast. Three primary methods are discussed:

  1. N-gram Language Models (e.g., KenLM):
    • Concept: Train a simple n-gram model on the target data (T). Then, score each document in the raw data (R) by its perplexity under this model. Documents with low perplexity (i.e., those that "look like" the target data) are kept.
    • Pros/Cons: Extremely fast and simple but crude. It only captures local linguistic patterns and can be easily gamed.
  2. Linear Classifiers (e.g., fastText):
    • Concept: Train a fast, linear classifier to distinguish between positive examples (from the target dataset T) and negative examples (randomly sampled from the raw dataset R). This classifier is then run on the entire raw dataset to filter for high-quality documents.
    • Details: To handle the massive feature space of n-grams, fastText uses the hashing trick, mapping n-grams into a fixed-size vector space. This is a highly scalable and popular method.
  3. Importance Sampling:
    • Concept: A more principled, probabilistic approach. It models the problem as trying to draw samples from the target distribution (P) when we can only sample from the raw distribution (Q).
    • Method: An "importance weight" is calculated for each document in the raw set, defined by the ratio of its probability under the target model versus the raw model (P(doc) / Q(doc)). Documents are then re-sampled with a probability proportional to this weight. To make this computationally feasible, the probabilities P and Q are estimated using simple hash-based n-gram models.

Part 2: Applications of Filtering
The lecture shows how this general filtering framework is applied to solve several key data curation tasks.

Part 3: Deduplication Algorithms
Deduplication is complementary to filtering. It aims to remove redundant documents to improve training efficiency and reduce model memorization. The core algorithmic challenge is to perform this pairwise comparison task in linear, not quadratic, time.

  1. Exact Deduplication:
    • Method: The simplest approach is to hash every document (or paragraph, or sentence span) and keep only one instance per unique hash value.
    • Bloom Filters: A memory-efficient, probabilistic data structure is introduced for this task. A Bloom filter can definitively say if an item is not in a set, but has a small, controllable false positive rate for confirming membership. This is achieved by using multiple hash functions to map each item to several bits in a bit array.
  2. Near-Deduplication (Approximate Matching):
    • Goal: To find documents that are highly similar but not identical (e.g., differing only by a comma or a boilerplate template).
    • Jaccard Similarity: The standard metric for set similarity (intersection over union).
    • MinHash: A clever hashing scheme with a crucial property: the probability that the MinHash of two sets is equal is exactly their Jaccard similarity (P(h(A) = h(B)) = Jaccard(A, B)). This converts a pairwise similarity metric into a unary hash function.
    • Locality Sensitive Hashing (LSH): The final step. To sharpen the probability around a desired similarity threshold, multiple MinHash functions are used. They are grouped into "bands," and two documents are considered a match if they collide in all hash functions within any single band. This "and/or" structure creates a sharp, S-shaped probability curve, ensuring that only documents above the similarity threshold are likely to be flagged as duplicates.

The lecture concludes by emphasizing that while these algorithms provide the necessary tools, building intuition about what makes a good dataset ultimately requires hands-on experience: filtering data, training models, and iterating.

L15: Alignment: SFT/RLHF

The lecture marks a shift in the course from pre-training (building a capable "base model" like GPT-3) to post-training (turning that base model into a useful and safe product like ChatGPT). The goal is to understand the arrow in the diagram GPT-3 -> ChatGPT. This involves instilling two key properties into the model: the ability to follow complex, nuanced instructions and the ability to adhere to safety guardrails. The process involves collecting specific data about desired behaviors and then using algorithms to train the model to exhibit those behaviors.
Part 1: Supervised Fine-Tuning (SFT)
The first and most foundational step in post-training is Supervised Fine-Tuning (SFT). This is an imitation learning process where the model is trained on a dataset of expert-written demonstrations (prompt-response pairs).

Part 2: Reinforcement Learning from Human Feedback (RLHF)
The lecture then transitions to the second, more advanced stage of post-training: RLHF. This moves from a purely generative modeling framework (imitating a reference distribution p*) to a reinforcement learning framework (finding a policy that maximizes a reward).

L16: Alignment: RL 1

The lecture begins by finishing the discussion on Reinforcement Learning from Human Feedback (RLHF), recapping the DPO (Direct Preference Optimization) algorithm. DPO simplifies RLHF by framing it as a supervised learning problem on preference pairs, where the goal is to increase the likelihood of preferred responses and decrease the likelihood of dispreferred ones. While effective, DPO and its variants still suffer from overoptimization, where optimizing the proxy reward model eventually leads to worse performance on true human preferences.
This limitation motivates a shift in perspective. Instead of relying on noisy, gameable human preferences, the lecture proposes moving to domains where rewards are verifiable, objective, and scalable. This approach takes inspiration from classic RL successes like AlphaGo, where the reward is unambiguous (win or lose the game). By applying the powerful machinery of RL to domains like math and coding, where correctness can be automatically verified, we can train models for complex reasoning.
Part 1: The Algorithms - PPO and GRPO
Before diving into case studies, the lecture details the core RL algorithms used for this task.

  1. PPO (Proximal Policy Optimization) Revisited:
    • Concept: PPO is a robust, on-policy RL algorithm. It improves upon the basic policy gradient by using an "advantage" function (a variance-reduced reward) and a "clipping" mechanism to ensure that policy updates are not too large, which stabilizes training.
    • Complexity: While powerful, PPO is notoriously complex to implement, requiring a separate "value model" to estimate advantages and numerous "implementation details" that are critical for performance.
  2. GRPO (Group Relative Preference Optimization):
    • Concept: GRPO is a much simpler algorithm, originally proposed in the DeepSeekMath paper, that has become the standard for reasoning-focused RL. It replaces PPO's complex advantage estimator with a simple, intuitive baseline.
    • The GRPO Advantage: For a given prompt, the model generates a "group" of G possible responses. The advantage for a single response i is its z-score: (reward_i - mean(rewards_in_group)) / std_dev(rewards_in_group). This baseline naturally accounts for problem difficulty and provides effective variance reduction without needing a value model.
    • Subtleties: Later analysis (the "Dr. GRPO" paper) revealed that the standard deviation term can bias the model towards problems that are either too easy or too hard. More importantly, the common practice of normalizing rewards by response length creates a perverse incentive for the model to generate extremely long, incorrect answers, as this minimizes the negative reward.

Part 2: Case Studies in RL from Verifiable Rewards
This section examines three recent, state-of-the-art open models that were trained using these techniques.
Case Study 1: DeepSeek-R1

Case Study 2: Kimi 1.5

Case Study 3: Qwen 3

The lecture concludes that reinforcement learning with simple, outcome-based verifiable rewards has become the dominant and highly effective paradigm for training state-of-the-art reasoning models.

L17: Alignment: RL 2

The lecture begins by finishing the discussion from the previous session on Reinforcement Learning from Human Feedback (RLHF). It recaps the DPO (Direct Preference Optimization) algorithm, which simplifies RLHF by reframing it as a supervised learning problem on preference pairs. Despite its simplicity and effectiveness, DPO and RLHF in general suffer from overoptimization, where maximizing the proxy reward model eventually leads to worse performance on true human preferences.
This limitation motivates a shift to a new paradigm: Reinforcement Learning from Verifiable Rewards. Instead of relying on noisy and gameable human feedback, this approach focuses on domains like math and coding where the reward is objective and can be automatically and reliably verified (e.g., the final answer is either correct or incorrect). This allows for a more robust and scalable application of RL, akin to its successes in games like Go.
Part 1: Reinforcement Learning Fundamentals for Language Models
Before diving into specific algorithms, the lecture grounds the RL framework in the context of language models.

Part 2: Policy Gradient and the Importance of Baselines
This section provides a mechanical deep dive into the core algorithm used in this domain: policy gradient.

  1. The Policy Gradient Theorem: The fundamental equation of policy gradient methods is derived. It shows that the gradient of the expected reward is equivalent to the expectation of the reward multiplied by the gradient of the log-policy (∇log π). Intuitively, this means you increase the probability of actions that lead to high rewards and decrease the probability of those that lead to low rewards.
  2. The Problem of High Variance: The main challenge with naive policy gradient is high variance, especially in sparse reward settings (like math problems, where most attempts are incorrect). If the model rarely succeeds, it gets a reward of 0 most of the time, leading to zero-gradient updates and a failure to learn.
  3. Baselines for Variance Reduction: The key solution is to use a baseline. Instead of updating based on the raw reward R, the model is updated based on R - B(s), where B(s) is a baseline function that depends only on the state (the prompt). This does not change the objective in expectation but can dramatically reduce the variance of the gradient estimates. The optimal baseline is the expected reward for a given state, E[R | s], which is also known as the value function. The quantity R - V(s) is called the advantage.

Part 3: From PPO to GRPO - A Practical Algorithm
The lecture details the evolution of the specific algorithm used for this task, moving from the complex PPO to the simpler GRPO.

The lecture concludes by reiterating that while the optimization part of RL can be challenging, the larger and more active area of investigation is the design of non-hackable rewards and generalizable environments. The systems challenges of scaling RL are also significant due to the complex interplay of inference, multiple model copies (policy, reference model), and environment interaction.

All writing