PyTorch (einops), Resource Accounting

Napkin math you can do before writing any code

The course’s framing question is fixed: given finite compute and memory, train the best model you can. Before you can maximize efficiency, you have to be able to account for it, and Percy opens with two questions you should be able to answer by the end of the lecture.

How long does it take to train a 70B-parameter model on 15T tokens with 1,024 H100s? Total work is $6 \cdot 70\times10^9 \cdot 15\times10^{12}$ FLOPs (the $6ND$ formula this lecture derives). An H100 promises 1,979 teraFLOP/s — but that’s the sparsity number, so divide by 2 — and you’ll realize about half of that in practice (MFU of 0.5). Chain it together:

total_flops = 6 * 70e9 * 15e12
h100_flop_per_sec = 1979e12 / 2
flops_per_day = h100_flop_per_sec * 0.5 * 1024 * 60 * 60 * 24
days = total_flops / flops_per_day   # ≈ 143 days

What’s the largest model you can train on 8 H100s with AdamW? Each parameter costs 2 bytes (bf16 weight) + 2 (gradient) + 4 + 4 (fp32 Adam moments) = 12 bytes, so $8 \cdot 80\text{GB} / 12 \approx 53$B parameters — an upper bound, since activations (which scale with batch size and sequence length) aren’t counted.

Tensors and what precision buys you

Everything is a tensor: parameters, gradients, activations, optimizer state, data. Memory is simply the number of elements times the bytes per element, and the bytes per element is a design choice with real consequences.

x = torch.zeros(4, 8)              # fp32 by default
x.numel() * x.element_size()       # 32 * 4 = 128 bytes

One feedforward matrix in GPT-3 is already 2.3 GB at fp32 — these get big.

The practical recipe is mixed precision: bf16 for parameters, activations, and gradients; fp32 for optimizer state, where you’re accumulating averages over many steps and sloppiness compounds. PyTorch’s AMP library automates the safe casts (matmuls yes, exponentials no). And remember the deep-learning inversion of scientific computing: “we’re going the other way, because even 32 is a lot.”

einops: stop transposing, start naming

Percy’s motivation is honest: x @ y.transpose(-2, -1) makes you reason about what -2 and -1 are, and “I always get confused by transposes.” einops names the dimensions instead.

There’s no speed cost — it compiles to the same primitives, “you can think about it as just sugar” — but the transposes disappear into the naming, and code stops caring how many batch dimensions ride along.

Counting FLOPs (and the two acronyms that sound identical)

A FLOP is one floating-point addition or multiplication. FLOPs counts work done; FLOP/s measures hardware speed. GPT-3 took about $3.14\times10^{23}$ FLOPs; GPT-4 is speculated around $2\times10^{25}$.

For the matmul at the heart of everything — $x$ of shape $B \times D$ times $w$ of shape $D \times K$ — there is one multiply and one add per $(i, j, k)$ triple:

$\text{FLOPs} = 2 \cdot B \cdot D \cdot K$

Read it as 2 × (number of data points) × (number of parameters) — the form that generalizes to Transformers. Elementwise operations are linear in the tensor size; nothing you’ll meet is as expensive as the matmuls, so matmuls are what we count.

Arithmetic intensity: why MFU isn’t 1

A computation moves inputs from memory to the accelerator, computes, and moves results back. Two hardware numbers govern it: compute speed (H100: $989\times10^{12}$ FLOP/s dense bf16) and memory bandwidth (H100: $3.35\times10^{12}$ bytes/s). Assuming perfect overlap, runtime is the max of communication time and computation time — whichever dominates names the regime, memory-bound or compute-bound.

Percy walks the ladder for a million-element bf16 vector:

The full training step: 6ND, and where the bytes go

For gradients, zoom in on one layer $h_2 = h_1 w_2$ of a deep linear network. The forward pass costs $2BD^2$. The backward pass must produce two gradients — one flowing to the input ($\partial\text{loss}/\partial h_1$) and one for the parameters ($\partial\text{loss}/\partial w_2$), each another matmul against the incoming gradient:

h1_grad = einsum(h2.grad, w2, "batch out, in out -> batch in")
w2_grad = einsum(h2.grad, h1, "batch out, batch in -> in out")

So backward is twice forward, and per training step:

Memory per parameter under mixed precision with an Adam-family optimizer: 2 bytes (bf16 parameter) + 2 (gradient) + 8 (fp32 first and second moments) — the 12 bytes from the opening question. AdaGrad-style optimizers with one accumulator need 4 instead of 8. Activations add $2BDL$ bytes on top, which motivates the two closing tricks:

• Gradient accumulation — compute gradients on micro-batches and defer the optimizer step, shrinking activation memory by the micro-batch ratio while keeping the effective batch size.
• Activation checkpointing — keep activations at only $\sqrt{L}$ of the layers and recompute the rest during backward: memory drops from $O(L)$ to $O(\sqrt{L})$ for one extra $O(L)$ recompute. “Trade memory for compute.”

Check yourself

Exercises in the lecture’s napkin-math spirit — do them on paper with a calculator, nothing else.

scores = einsum(Q, K, "b h sq d, b h sk d -> b h sq sk")   # d summed out
out = rearrange(out, "b h s d -> b s (h d)")               # merge heads