Lecture 2: PyTorch and Resource Accounting

What this lecture covers

This lecture builds the PyTorch and "resource accounting" foundation that the rest of the course leans on: given a model and hardware, how much memory does training need, and how many floating-point operations (flops) does it cost? It opens with two motivating napkin-math questions — how long it takes to train a 70B-parameter model on 15T tokens on 1,024 H100s, and the largest model you can train on 8 H100s with AdamW — and then builds up the tools needed to answer them from scratch.

The lecture works bottom-up: floating-point representations and their memory cost, how PyTorch tensors are really views into flat storage, matrix multiplication cost, the einops/jaxtyping style of writing dimension-safe tensor code, forward vs. backward pass flop counts, and finally a full memory and compute budget for a small toy model. By the end you should be able to estimate, before running any code, how much memory a model's parameters, gradients, activations and optimizer state need, and how many flops a training run will consume.

Key ideas

• Floating-point formats trade off range and precision: FP32 (4 bytes) is the safe default; FP16 halves memory but underflows on very small numbers; BF16 keeps FP32's dynamic range in 16 bits by shrinking the fraction instead, making it the usual compute dtype; FP8 (H100-only) is even smaller and less stable.
• Mixed precision: parameters and optimizer state are typically kept in FP32 for stability, while forward/backward compute often runs in BF16 for speed.
• Tensors are views into storage: a PyTorch tensor is metadata (shape and strides) pointing at a flat array; operations like slicing, view(), and transpose often share the same storage rather than copying, so mutating one tensor can mutate another.
• Matrix multiplication flop rule: a matmul with dimensions (m, n) x (n, p) costs 2 * m * n * p flops, and matrix multiplication dominates total compute in deep learning models.
• The "6x" rule: for models dominated by matmuls, the forward pass costs roughly 2 * (number of tokens) * (number of parameters) flops and the backward pass costs about twice that, so total training compute is approximately 6 * tokens * parameters.
• Model FLOPs utilization (MFU): actual achieved flops/second divided by the hardware's promised (marketing) flops/second; above 0.5 is considered good, and MFU should always be measured by benchmarking rather than assumed.
• einops/jaxtyping style: naming tensor dimensions explicitly (e.g. batch sequence hidden) instead of indexing by position (-1, -2) makes tensor code self-documenting and less error-prone.
• Memory budget has four parts: parameters, gradients, optimizer state, and activations, each contributing a term you can compute directly from model shape, batch size and dtype.

Walkthrough

Napkin math and why resource accounting matters (0:05)

The lecture opens with two example calculations: estimating training time for a 70B-parameter model on 1,024 H100s using 6 * parameters * tokens flops divided by hardware throughput at a given MFU, and estimating the largest model trainable on 8 H100s given that AdamW needs about 16 bytes per parameter (parameters, gradients, and optimizer state, in FP32). These calculations are meant to become second nature so that "how much will this run cost" is always answerable before you hit run.

Floating-point representations and memory (5:12)

The lecture explains how tensor memory is simply element count times bytes per element, then compares FP32 (8-bit exponent, 23-bit fraction, the "safe" default), FP16 (which underflows small values like 1e-8), and BF16 (developed in 2018, same 16-bit size as FP16 but with FP32's exponent range and a shorter fraction, at the cost of resolution). FP8, supported on H100, pushes this further but is more unstable. The rule of thumb given is: use FP32 for parameters and optimizer state, and BF16 for most forward/backward compute.

Tensors as views and einops-style dimension naming (14:24)

After covering GPU placement (tensors default to CPU and must be moved explicitly), the lecture explains tensor internals: shape and stride metadata addressing into flat storage. Operations like row/column slicing, view(), and transpose typically create a new view sharing the same underlying storage rather than copying data — but transposed tensors become non-contiguous, so a further reshape may force a copy. The lecture then introduces einops and jaxtyping, showing how naming dimensions (batch sequence hidden) in operations like einsum and reduce avoids the error-prone dim=-1/dim=-2 style and makes intent explicit in code.

Counting flops: matmuls, gradients, and the 6x rule (33:41)

The lecture establishes that a matrix multiplication costs 2 * (product of the three dimensions) flops, and that this dominates deep learning compute since specialized hardware is built for matmuls. Working through a simple two-layer linear model with backpropagation via the chain rule, it shows the forward pass costs 2 * batch * parameters flops while the backward pass, which needs gradients with respect to both weights and activations at each layer, costs roughly 4 * batch * parameters. Summing forward and backward gives the 6 * tokens * parameters estimate used throughout the course. MFU is introduced as actual measured flops/second divided by the hardware's advertised flops/second, with a live benchmark example showing that BF16 can have lower MFU than FP32 despite being faster in absolute terms, because the "promised" peak numbers are optimistic.

Building a toy model and reasoning about randomness (59:14)

The lecture builds a small multi-layer linear model ("cruncher") to make the abstract accounting concrete, including a note on Xavier-style initialization (rescaling by 1/sqrt(input dimension)) to prevent activations from blowing up as hidden size grows. It stresses fixing random seeds separately for each source of randomness (initialization, dropout, data ordering) so bugs are reproducible, and briefly covers memory-mapping large tokenized datasets with numpy.memmap instead of loading everything into RAM.

Implementing an optimizer and totaling the memory budget (1:06:25)

After reviewing the lineage of optimizers (SGD, momentum, Adagrad, RMSProp, Adam), the lecture implements Adagrad from scratch by subclassing PyTorch's optimizer class, showing how per-parameter state (like the running sum of squared gradients) is stored and updated in optimizer.step(). It then totals the toy model's memory: parameters, activations, gradients, and optimizer state, each as a function of model dimensions and batch size, multiplied by bytes-per-value — the same structure students apply to the transformer in assignment one. It closes with practical notes on checkpointing (saving both model and optimizer state periodically) and mixed-precision training, including the observation that low precision is far harder to make stable during training than at inference time, where aggressive quantization is much easier to get away with.

Before you watch

• Watch Lecture 1 (Overview and Tokenization) first, since this lecture assumes the course's efficiency-first framing and references assignment one's tokenizer work.
• Basic familiarity with PyTorch tensors, autograd, and writing a training loop will make the tensor-internals and optimizer sections easier to follow, since the lecture moves quickly through material it expects as partial review.

Check your understanding

1. Why does BF16 keep FP32's dynamic range in half the bits, and why does that matter more for deep learning than FP16's extra fraction precision?
2. Derive the 6 * tokens * parameters flop estimate: where does the factor of 2 for the forward pass and 4 for the backward pass each come from?
3. What is model FLOPs utilization (MFU), and why did the lecture's BF16 benchmark show lower MFU than FP32 despite running faster?
4. Why can slicing or transposing a PyTorch tensor sometimes produce a view that shares storage with the original, and when does an operation force a copy instead?
5. What are the four components of a model's memory budget during training, and which of them scales with batch size rather than model size alone?