Stanford CS336: Language Models from Scratch (Simplified Notes)¶

Welcome to the simplified notes for Stanford's CS336 course! This course focuses on the mechanics, mindset, and intuitions behind building Large Language Models (LLMs) from scratch. These notes simplify the key concepts from the first five lectures to help you build a solid foundation.

Lecture 1: Course Overview and Tokenization¶

Why Build from Scratch?¶

Historically, researchers trained their own models. Today, most just prompt APIs (like GPT-4 or Claude). To do fundamental research, you need to understand the underlying mechanics. The philosophy is understanding via building.

The "Bitter Lesson" and Efficiency¶

The biggest takeaway in AI development is that algorithms that scale well are what matter. Efficiency is everything at scale: accuracy = efficiency x resources. Squeezing the most out of your given compute and data budget is the primary goal.

Tokenization¶

Language models don't process text directly; they operate on discrete units called tokens. - The Problem: Characters are too granular (sequences are too long), and words are too sparse (vocabulary gets too big). - The Solution: Byte-Pair Encoding (BPE). - Starts with individual bytes. - Iteratively merges the most frequent adjacent pairs. - Result: Common words become single tokens, while rare words are split into multiple smaller chunks. - Why it matters: Shorter token sequences mean less computation (since attention is quadratic in sequence length).

Lecture 2: Resource Accounting & Systems¶

To maximize computational efficiency, you need to measure it.

Tensors and Data Types¶

Everything in ML is a tensor. - fp32 (32-bit float): Precise but uses a lot of memory. - fp16 / bf16 (16-bit): Saves memory. bf16 is preferred because it matches the dynamic range of fp32, preventing training instability. - Mixed Precision: We train models using bf16 for weights and activations, but keep optimizer states (like Adam's moving averages) in fp32 for stability.

FLOPs (Floating Point Operations)¶

Forward Pass: Takes ~2 FLOPs per parameter per token.
Backward Pass: Takes ~4 FLOPs per parameter per token.
Total: ~6 FLOPs per parameter per token.

# Interactive CS336 Parameter & FLOPs Calculator
def calculate_training_flops(num_params_billions, num_tokens_billions):
    # Standard formula: ~6 FLOPs per parameter per token
    flops = 6 * (num_params_billions * 1e9) * (num_tokens_billions * 1e9)
    return flops

params_b = 7.0    # 7B model (e.g. LLaMA-7B)
tokens_b = 2000.0 # 2 Trillion tokens
flops = calculate_training_flops(params_b, tokens_b)

print(f"Model Parameters: {params_b} Billion")
print(f"Tokens Trained:   {tokens_b} Billion")
print(f"Total FLOPs:      {flops:.2e}")
print(f"Estimated A100 GPU hours (@ 150 TFLOPs effective): {flops / (150e12 * 3600):,.1f} hours")

The Roofline Model: Memory vs. Compute Bound¶

When running an operation on a GPU, it takes time to compute, but it also takes time to move data from memory to the compute units. - Memory-bound: Waiting for data to arrive (e.g., Activation functions like ReLU, generating tokens one-by-one). - Compute-bound: Data arrives fast, waiting for the math to finish (e.g., large Matrix Multiplications during training).

Tricks to Save Memory¶

Gradient Accumulation: Doing many small batches and adding up their gradients before updating the model.
Activation Checkpointing (Rematerialization): Discarding intermediate activations during the forward pass to save memory, and simply recalculating them during the backward pass. Trading compute for memory!

Lecture 3: Architecture and Hyperparameters¶

Modern LLMs share a very similar "LLaMA-like" architecture, differing slightly from the original 2017 Transformer.

Key Architectural Shifts¶

Pre-Norm vs. Post-Norm:
The original Transformer normalized data after adding the residual connection (Post-norm).
Modern models normalize before the operation (Pre-norm). This is much more stable for deep networks and prevents gradient spikes.
LayerNorm vs. RMSNorm:
Modern models use RMSNorm. It’s simpler, faster, and just as effective. It scales the variance but doesn't subtract the mean, removing the need to store bias terms.
Positional Embeddings:
Replaced old sinusoidal embeddings with RoPE (Rotary Position Embeddings), which encode relative positions by rotating the attention vectors.
Activations (FFN):
Moved from ReLU to SwiGLU, which introduces a gating mechanism into the feed-forward network, increasing expressivity.

Lecture 4: Attention Alternatives and Mixture of Experts¶

Standard Attention is quadratic \(O(N^2)\), which becomes incredibly expensive as context windows grow (e.g., 100k+ tokens).

Linear Attention & State Space Models¶

Linear Attention: By changing the order of matrix multiplication, we can reduce complexity to \(O(N)\). Interestingly, this makes it look mathematically like a Recurrent Neural Network (RNN).
Mamba & Gated Delta Nets: These architectures generalize linear attention by adding data-dependent gates. They allow the model to "selectively forget" past information.
Hybrids: Many recent models use a mix of standard attention layers and linear/state-space layers (e.g., 1 standard attention layer for every 7 linear layers) to get the best of both worlds.

Mixture of Experts (MoE)¶

Instead of one massive Feed-Forward Network (FFN), MoE uses several smaller FFNs ("experts"). - A router decides which experts process each token. - Why it’s popular: You can drastically increase the parameter count (making the model smarter) without increasing the FLOPs (compute cost) per token, because only a subset of experts are active at any time.

Lecture 5: GPUs and FlashAttention¶

To make models fast, you must understand the hardware.

GPU Anatomy¶

CPUs have a few powerful cores optimized for latency (fast single-thread performance).
GPUs have thousands of weak cores optimized for throughput (massively parallel operations).
Tensor Cores: Specialized circuits built purely for multiplying matrices incredibly fast.

GPU Execution Hierarchy¶

Threads: Execute identical instructions on different data points.
Warps: A group of 32 threads running in lockstep.
Blocks: A larger group of threads sharing local memory.
Streaming Multiprocessors (SM): The physical hardware unit that executes blocks.

The Memory Hierarchy (The Real Bottleneck)¶

The closer memory is to the compute cores, the faster it is—but there’s less of it. - Global Memory (HBM/DRAM): Massive (e.g., 80GB), but incredibly slow to read/write. - Shared Memory (SRAM): Tiny, but blisteringly fast.

FlashAttention¶

Standard attention algorithms naively read and write huge \(N \times N\) matrices to Global Memory, causing massive bottlenecks. FlashAttention is a hardware-aware algorithm that calculates attention by carefully breaking the matrices into tiles that fit perfectly into the ultra-fast Shared Memory (SRAM), drastically reducing slow memory reads/writes.

[!TIP] Key Takeaway: Across all these lectures, a central theme emerges—modern AI engineering is a constant balancing act between algorithmic expressivity, training stability, and hardware efficiency (managing memory and compute).