PyTorch, resource accounting

Transcript: cleaned auto-captions with timestamps

This is the course's least glamorous lecture and quietly its most load-bearing. CS336's whole premise is that language-model research is now an efficiency discipline: you are not searching a space of ideas, you are searching a space of ideas you can afford. That claim only becomes operational if you can convert a design decision into a number of GPU-days before you spend them. So before any Transformer appears — Tatsunori Hashimoto covers architecture next lecture — Percy spends 79 minutes installing a habit: whenever you write a line of tensor code, know what it costs in bytes and what it costs in FLOPs. The mechanics are ordinary PyTorch. The mindset is the deliverable.

Outline, with timestamps

• 00:05 — Why this lecture exists: the primitives, bottom-up, with the meter running
• 00:39 — Two napkin questions: 144 days, and 40B parameters
• 05:03 — Memory accounting: the tensor, and the float zoo (fp32 / fp16 / bf16 / fp8)
• 14:08 — Compute accounting begins: where your tensor actually lives, CPU vs GPU
• 17:35 — A tensor is a pointer plus a stride table: storage, views, contiguity
• 23:09 — Elementwise ops, triu and the causal mask, batched matmul
• 25:30 — einops: naming dimensions, einsum, reduce, rearrange
• 33:29 — FLOPs vs FLOP/s, training-run intuitions, and the sparsity asterisk
• 37:16 — The linear model: 2·m·n·p, and forward = 2 × tokens × parameters
• 42:24 — Wall clock, promised FLOP/s, and model FLOPs utilization
• 49:46 — The backward pass: why it is 4×, and where the 6 comes from
• 59:18 — nn.Parameter, initialization that survives width, and a toy model
• 66:36 — Optimizers as a lattice, and the full memory ledger
• 74:30 — Checkpointing, mixed precision, and the low-precision frontier

Two questions you should be able to answer on a napkin

Percy opens with two questions and answers both in under two minutes, which is the point: these are not research problems, they are arithmetic you should be embarrassed not to have done.

How long to train a 70B-parameter dense Transformer on 15T tokens using 1,024 H100s? Total compute is 6 × 70e9 × 15e12 = 6.3e24 FLOPs. An H100 is advertised at 1,979 TFLOP/s of bf16 tensor-core throughput, but that headline number assumes structured sparsity; the dense number is exactly half, 989.5 TFLOP/s. Assume you extract half of that in practice (an MFU of 0.5, defined below) and you get ~494.75 TFLOP/s per GPU. Times 1,024 GPUs, times 86,400 seconds, is 4.38e22 FLOPs per day. Divide: ~144 days. Note what just happened — a number that sounds like it requires a cluster, a team and a quarter of planning fell out of four multiplications.

What is the largest model you can train on 8 H100s with AdamW, naively? Each H100 has 80 GB of HBM, so 640 GB total. Naive fp32 costs 16 bytes per parameter: 4 for the parameter, 4 for its gradient, and 8 for AdamW's two moment estimates. 640e9 / 16 = ~40B parameters. Two caveats Percy flags immediately. First, this ignores activations entirely, which depend on batch size and sequence length and which Assignment 1 will make you account for. Second, the obvious "save memory by going bf16" move does not actually save memory here: bf16 parameters and gradients (2 + 2) plus an fp32 master copy of the parameters (4) plus 8 bytes of moments is still 16 — it buys speed, not space. Getting the number down requires sharding the optimizer state across devices, which is what ZeRO does and what lecture 8 covers.

"Mindset is about resource accounting, and it's not hard. It's just — you just have to do it."— Percy Liang, 04:30

Precision is a dynamic-range problem, not a size problem

The naive story about low precision is "fewer bits, less memory, faster." The real story is that the bits are split between an exponent, which buys you range, and a mantissa, which buys you resolution — and deep learning cares far more about the first than the second.

fp16 shrinks the exponent from 8 bits to 5, and that is the whole problem: torch.tensor([1e-8], dtype=torch.float16) is exactly zero. On a small model you may never notice; on a large one, underflow in a gradient somewhere becomes an instability you will spend a week bisecting. bfloat16, developed at Google Brain in 2018, fixes this by refusing the trade — it keeps fp32's 8 exponent bits and pays for them out of the mantissa. Same 2 bytes, same dynamic range as fp32, coarser steps between representable values. Deep learning turns out to be extremely tolerant of coarse steps and extremely intolerant of overflow, so bf16 won.

"If you're talking to a machine-learning person, float32 is the max you'll ever probably need to go, because deep learning is kind of sloppy like that."— Percy Liang, 06:47

The working rule the lecture lands on, and returns to at the end: keep parameters and optimizer state in fp32 because those accumulate over the whole run and errors compound there; treat bf16 as transitory — cast down, run the forward pass and the matmuls, come back up to accumulate. That is what mixed-precision training means, and PyTorch's AMP automates most of the bookkeeping because precision is a concern that cuts across a cleanly-factored model. Some places, like attention softmax accumulation, people keep in fp32 by hand.

A tensor is a pointer plus a stride table

A PyTorch tensor is not a block of numbers. It is a pointer into a flat allocation plus metadata that says how to walk it: one stride per dimension. For a 4×4 float matrix, stride(0) == 4 (skip a row) and stride(1) == 1 (skip a column), so element (1,2) lives at offset 1·4 + 2·1 = 6. Once you see it that way, the memory behaviour of the operations you use every day becomes obvious rather than folkloric.

Slicing a row, slicing a column, .view(), .transpose() — all of these produce a new tensor object with new strides pointing at the same storage. They cost zero bytes and essentially zero time. Two consequences. The good one: name intermediate views freely, because readable code is free here. The dangerous one: mutating x mutates y, silently, if y was a view of it.

Contiguity is where the free lunch ends. A transposed tensor is not contiguous — walking it in index order means skipping around in storage — and PyTorch will refuse to .view() it into another shape, because a view is only definable when the walk is a simple stride. Calling .contiguous() (or .reshape(), which is roughly "contiguous-then-view") makes it work by copying. So in ledger terms: views are a rounding error, and every .contiguous() is a full duplicate of the tensor in both bytes and bandwidth. If a hot loop keeps calling .reshape() on a transposed activation, you have found a real cost, not a style issue.

The same lens explains the elementwise operations. Every one of them allocates a new tensor of the same shape, because it has to put the answer somewhere; each costs O(mn) FLOPs, which is to say, nothing. torch.ones(3,3).triu() gets a special mention because that upper-triangular matrix is exactly the causal attention mask you will build in Assignment 1.

Naming your dimensions. Percy takes a detour into einops and it is worth taking seriously, because the failure mode it prevents is a correctness bug rather than a performance one. Standard PyTorch attention code contains x @ y.transpose(-2, -1), and six weeks later nobody remembers what -2 was. The fix is to name dimensions instead of counting backwards from the end. jaxtyping annotations like Float[torch.Tensor, "batch seq heads hidden"] turn the comment into a type — documentation only, with no runtime enforcement unless you add a checker, but documentation that lives where it can be read. einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2") then expresses the same matmul with the contraction visible: every dimension named in the inputs but absent from the output is summed over. reduce does the same for aggregations, and rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2) splits and re-merges the flattened head dimension that multi-head attention spends its life juggling. Asked whether this compiles down to something efficient, the answer is yes — the library picks a contraction order, and under torch.compile it does that once and reuses the plan.

2·m·n·p, and the 6 that follows from it

Here is the arithmetic the rest of the course rests on. Multiplying an m×n by an n×p matrix requires, for every (i,j,k) triple, one multiply and one add. Two operations, m·n·p triples: 2·m·n·p FLOPs. Everything else you do in a deep network — elementwise ops, additions, normalizations — is linear in the tensor size and, for matrices large enough to be worth putting on a GPU, rounds to zero next to the matmuls.

Percy notes the circularity honestly, and it is a good thing to notice early: it is not a law of nature that matmuls dominate. It is that hardware is built for matmuls, so models whose cost is not dominated by matmuls are models nobody trains. The regime is self-selecting. That is precisely why napkin math works at all in this field.

Now reinterpret the shapes for a linear model mapping B points of dimension D to K outputs. The cost is 2·B·D·K, where B is the number of data points (read: tokens) and D·K is the number of parameters. So:

Forward pass ≈ 2 × (tokens) × (parameters). And this generalizes to Transformers to first order — with an asterisk for the sequence-length-dependent attention term, which is small as long as your context is not enormous relative to your width.

The backward pass is where the interesting factor comes from, and Percy derives it rather than asserting it. Take a two-layer linear net, x --W1--> h1 --W2--> h2 --> loss, and focus on W2. You need two different gradients there, not one. The first is the one you actually want: W2.grad[j,k] = Σ_i h1[i,j] · h2.grad[i,k], which is a contraction over (i,j,k) and therefore 2·B·D·K FLOPs. The second is the one you need in order to keep going: h1.grad[i,j] = Σ_k W2[j,k] · h2.grad[i,k], another 2·B·D·K, and without it you cannot backpropagate into W1 at all. Two matmuls per layer instead of one.

Backward pass ≈ 4 × (tokens) × (parameters). Total ≈ 6 × tokens × parameters. That is the 6 in the opening napkin question, and it is the constant that turns "we want a 70B model on 15T tokens" into "that is 6.3e24 FLOPs, so budget accordingly." It breaks in the obvious edge cases — heavy parameter sharing means one parameter can drive many FLOPs — but real language models sit comfortably in the regime where every parameter is touched about once per token.

MFU, and why you benchmark instead of trusting the brochure

FLOPs are not seconds. To close that gap you need two more numbers: what the hardware promises, and what you actually got. Model FLOPs utilization is their ratio — the FLOPs your model logically required, divided by the elapsed time, divided by the vendor's peak FLOP/s for that dtype.

(A vocabulary note Percy makes a point of: FLOPs with a lowercase s is a count of operations; FLOP/s is a rate. They are pronounced identically and mean completely different things, so the course writes the rate with a slash and never as "FLOPS".)

The live benchmark: a 16384 × 32768 by 32768 × 8192 matmul is 8.8e12 FLOPs. In fp32 on an H100 it takes about 0.16 s, so ~5.5e13 FLOP/s against a promised 67.5 TFLOP/s — an MFU near 0.8, which is very good, as it should be for a single enormous matmul with no communication. Recast the same matmul in bf16 and it runs in about 0.03 s: ~2.9e14 FLOP/s, roughly five times faster in absolute terms. But the promised bf16 number is 989.5 TFLOP/s, so the MFU is now around 0.3. Faster in wall clock, worse against the brochure. Percy's read is that the bf16 headline is simply optimistic.

"Always benchmark your code, and don't just assume that you're going to get certain levels of performance."— Percy Liang, 47:28

Three things to keep straight about MFU. It is model FLOPs, not hardware FLOPs: it counts the operations your model logically needs, so you are not penalized for being clever, and conversely, activation recomputation shows up as extra hardware work while leaving model FLOPs unchanged. Rules of thumb: ≥ 0.5 is respectable, single digits means something is badly wrong, and you will never approach 1.0 because communication and every kind of overhead are excluded from the numerator by construction. And watch the asterisk on the datasheet — NVIDIA's 1,979 TFLOP/s for the H100 assumes 2:4 structured sparsity, which essentially nobody uses for language-model training. The number you should compare against is half of it.

The full ledger: parameters, activations, gradients, optimizer state

The last third of the lecture assembles everything into one accounting exercise. Four buckets consume memory during training:

• Parameters — for a toy deep-linear model of width D with L square layers plus a scalar head, D²·L + D.
• Activations — B·D·L, one per layer per example per dimension. You store them because the gradient at layer i depends on the activation at layer i; you can trade them back for compute via activation checkpointing, which the course returns to later.
• Gradients — one per parameter, so exactly the parameter count.
• Optimizer state — one extra copy per parameter for AdaGrad's accumulated squared gradient; two for Adam or AdamW.

Worked in the lecture with D=4, L=2, B=2: 36 parameters, 16 activations, 36 gradients, 36 optimizer slots, times 4 bytes for fp32 = 496 bytes. Compute for one step is 6·B·(parameters). The numbers are absurd; the shape of the calculation is exactly the one Assignment 1 asks you to redo for a real Transformer, where the only change is that there are more matrices to enumerate and attention to handle.

The optimizer family is presented as a lattice worth memorizing, because it explains the memory column directly. SGD takes a step along the gradient. Momentum adds an exponential moving average of the gradient. AdaGrad instead scales the step by the accumulated sum of squared gradients. RMSProp is AdaGrad with an exponential average instead of a flat sum. Adam is RMSProp plus momentum — which is why it carries two state tensors per parameter, and why the naive 16-bytes-per-parameter figure from the opening napkin question is what it is. Writing an optimizer in PyTorch is unglamorous: subclass torch.optim.Optimizer, implement step(), keep whatever you need in the per-parameter self.state dict. You will write AdamW yourself in Assignment 1.

Three operational notes get folded in here and all three will bite you if you skip them. Seeds: randomness enters through initialization, dropout and data ordering, and there are three separate generators to set — torch.manual_seed, np.random.seed, random.seed. Set all three, and prefer a distinct seed per source so you can hold initialization fixed while varying data order. Data loading: tokenized data is just an array of integers; LLaMA's corpus is 2.8 TB, so you never load it, you np.memmap it and let the OS page in the slices you index. Pinning the CPU batch with .pin_memory() lets .to(device, non_blocking=True) overlap the host-to-device copy with GPU compute. Checkpointing: long runs crash. Save the model state dict, the optimizer state dict, and the iteration number — all three, or resuming silently restarts your optimizer.

What you build with this

This lecture is the toolkit for Assignment 1: Basics (handout PDF, leaderboard), whose second half — everything after the BPE tokenizer from lecture 1 — is exactly the material here. Concretely, it feeds: the nn.Module / nn.Parameter scaffolding and truncated-normal initialization you use for every Transformer submodule; the causal mask, which is triu plus a fill; the einsum/rearrange idioms the handout recommends for multi-head attention's head-splitting; your hand-written AdamW step(); the memmap-backed data loader and its random-offset batch sampler; and save/load checkpointing of model, optimizer and step. Most directly, the handout contains resource-accounting problems that ask you to derive the parameter count, activation memory and FLOP count of a real Transformer configuration — the 496-byte toy ledger above is a rehearsal for exactly that, and the MFU discipline is what you will need on the leaderboard, where the ranking is quality under a fixed compute budget. Nothing here asks you to differentiate by hand; autograd does that. The counting is the assignment.

Supporting materials, verified

• Executable lecture: lecture_02 trace — Percy Liang (2025) · the lecture is this program; the source is lecture_02.py, and reading it end-to-end takes twenty minutes and settles every number on this page.
• torch.Tensor — PyTorch docs — PyTorch · the dtype table and element_size(); paired with Tensor.stride for the storage/view model.
• bfloat16 floating-point format — Wikipedia · the bit layout Percy shows on the slide; the fastest way to internalize why bf16 has fp32's range.
• FP8 Formats for Deep Learning — Micikevicius et al., NVIDIA/Arm/Intel (2022) · the source of E4M3 and E5M2 and the ±448 / ±57344 ranges quoted in the lecture.
• FP8 primer — NVIDIA Transformer Engine — NVIDIA · what fp8 actually looks like in a training stack, including the per-tensor scaling the formats need to be usable.
• Mixed Precision Training — Micikevicius et al. (2017) · the paper behind the "fp32 master weights, low precision everywhere else" recipe the lecture closes on.
• torch.amp — automatic mixed precision — PyTorch · the library that keeps precision from contaminating your module definitions.
• FP8-LM: Training FP8 Large Language Models — Peng, Wu et al., Microsoft (2023) · the "people are pushing the envelope" claim at the end, made concrete: fp8 through gradients and optimizer state, not just the forward pass.
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al., Microsoft (2019) · cited under the 40B napkin answer; the way out of paying 16 bytes per parameter on every device.
• Adaptive Subgradient Methods for Online Learning and Stochastic Optimization — Duchi, Hazan & Singer, JMLR (2011) · AdaGrad, the optimizer Percy implements live because Adam is your homework.
• Understanding the difficulty of training deep feedforward neural networks — Glorot & Bengio, AISTATS (2010) · Xavier initialization; the 1/√(input_dim) rescaling that makes activations invariant to width.
• Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation — Rogozhnikov, ICLR (2022) · the argument for named dimensions; the tutorial is the practical entry point Percy recommends.
• NVIDIA H100 Tensor Core GPU datasheet — NVIDIA · where 1,979 TFLOP/s and its sparsity asterisk come from; also the A100 datasheet for the 312 TFLOP/s comparison.
• The FLOPs Calculus of Language Model Training — Dzmitry Bahdanau · the animated derivation of forward-2 / backward-4 that Percy plays on the slide.
• Transformer FLOPs — Adam Casson (2023) · carries the 6ND approximation up to a full Transformer, including where the attention term starts to matter.
• Transformer Math 101 — Quentin Anthony, Stella Biderman & Hailey Schoelkopf, EleutherAI (2023) · the memory-side companion this lecture wants: parameter, gradient, optimizer-state and activation ledgers for real Transformer configs, plus the inference-side accounting. (The lecture's own memory link, erees.dev/transformer-memory, no longer resolves.)
• How to Optimize Data Transfers in CUDA C/C++ — NVIDIA developer blog · why pinned memory makes non_blocking=True mean anything.
• The Transformer Model in Equations — John Thickstun · one of the three Transformer expositions Percy points at instead of covering it himself; the other is Jay Alammar's Illustrated Transformer.
• Decoupled Weight Decay Regularization — Loshchilov & Hutter (2017) · field map extra. The lecture names Adam but Assignment 1 asks for AdamW; this is the one-page difference, and it changes the optimizer-state accounting not at all.
• Reducing Activation Recomputation in Large Transformer Models — Korthikanti et al., NVIDIA (2022) · field map extra. The technique Percy defers when a student asks why activations must be stored — and the reason MFU is defined on model FLOPs rather than hardware FLOPs.
• Scaling Laws for Neural Language Models — Kaplan et al., OpenAI (2020) · field map extra. Where C ≈ 6ND became a convention, and the reason this lecture's arithmetic is the entry point to the scaling-laws arc later in the course.

Exercises

1. Redo both napkin questions with your own hardware — pick a cluster you could plausibly get (8 A100s? 64 H100s?), a model size and a token budget, and compute wall-clock days and peak memory. Do it before looking at anything. A good answer states the assumed MFU explicitly, says whether the peak FLOP/s you used was the dense or the sparse number, and names what the memory estimate omits (activations, fragmentation, the fp32 master copy). Then sanity-check: does your answer say a 1B-parameter model on 20B tokens is an afternoon or a month?
2. Measure your own MFU code — (1) allocate x of shape (B,D) and w of shape (D,K) on GPU, large enough that the kernel dominates launch overhead; (2) time x @ w over several trials with torch.cuda.synchronize() on both sides of the timed region — omit it and you will time an async launch, not the matmul; (3) compute 2·B·D·K and divide by the mean time; (4) look up your GPU's peak for that dtype and divide again; (5) repeat in bf16 and fp16. A good answer reports MFU for at least two dtypes, notes that the absolute FLOP/s went up while the MFU went down, and explains the discrepancy in terms of the datasheet rather than the code.
3. Build the memory ledger for a real config code — (1) instantiate a Transformer of a size you care about (GPT-2 small is fine); (2) compute the four buckets by hand from the config — parameters, gradients, AdamW state, activations for a given batch and sequence length; (3) multiply by bytes per element under an explicit dtype plan; (4) actually run one training step and read torch.cuda.max_memory_allocated(); (5) reconcile. A good answer gets within ~20% and, more importantly, names the specific things that account for the gap — the fp32 master copy, autograd's saved intermediates inside fused ops, the allocator's caching.
4. Break contiguity on purpose code — (1) build a small tensor, take a transpose, and confirm with .untyped_storage().data_ptr() that it shares storage; (2) confirm .view() raises and .reshape() does not; (3) check data_ptr() again after .contiguous(); (4) mutate the original and observe which of your aliases changed. A good answer explains, in stride terms, why the view is impossible after a transpose — and states the cost, in bytes and in bandwidth, of the .reshape() that silently rescued you.