CS2680 Modern AI Systems: Agents and Systems Optimizations
Lecture 2 — Modern ML basics I: transformers and the training loop, from a systems view

Lecture 1 used 6ND on faith. Today we earn it. We build a transformer out of matrix multiplications, count its parameters exactly, count the FLOPs a token costs in each direction, and then count the bytes — which turn out to be the number that decides what hardware you need. No machine learning background is assumed; systems background is, and it is the right background, because every question we ask today is an accounting question. By the end you should be able to read a model configuration file and say what it costs to store, to run, and to train.

Date: Tuesday, September 8, 2026 · 11:15am – 12:30pm · SEC LL2.221

Readings — all optional. Skim before class if you can; none is required.

Where this sits

Lecture 1 drew the stack and did one calculation: serving compute overtakes training compute at T = 3D generated tokens. That calculation borrowed C ≈ 6ND without justification. Today supplies it, along with the reference model configuration that Part I reuses. Thursday (Sep 10) takes today's parameter count and today's FLOP rules, runs the model forwards only, and shows that the cost of serving is governed by bytes rather than by arithmetic — which is the fact the rest of the course is organized around.

Instructor notes — Timing plan

75-minute class (Tue/Thu, 11:15am – 12:30pm, SEC LL2.221).

TimeSegmentNotes
0–4Recap and the reference 7BWrite the config on a side board and leave it up all term. Do not derive yet.
4–92.1 Tokens and embeddingsFast. This is the I/O boundary, not a topic.
9–202.2 What attention computesThe miniature with real numbers is the section. Do it by hand.
20–272.3 GEMM and not-GEMMClassify, then the 0.19%-of-FLOPs / 15×-the-bytes contrast.
27–342.4 Parameter accountingDerive N live, digit by digit. They will use it all term.
34–432.5 Where the FLOPs go2N from one multiply and one add. Then the crossover at 24,704.
43–522.6 The training loop and 6NDTwo backward GEMMs on the board is the whole factor of 3.
52–612.7 Four memory consumers16 bytes/param, then activations dwarfing everything.
61–672.8 The budget, and the punchline107.8 GB against 80. Say nothing for three seconds after.
67–722.9 What a framework doesLaunch counting, then the PyTorch/TensorFlow argument.
72–752.10 Mixed precisionbf16 versus fp16, the master copy, done.

If running long: compress 2.9 to the launch-count arithmetic and one sentence on eager versus graph, and cut 2.10 to "bf16 has fp32's exponent range, which is why it needs no loss scaling." Protect 2.4, 2.6, and 2.8 — the parameter count, the factor of three, and the 107.8 GB. If those three land, Thursday works.

Learning objectives

By the end of this class you should be able to:

  1. Compute the exact parameter count of a decoder-only transformer from L, d_model, d_ff, V, and the head configuration, and say which term dominates.
  2. Derive the ≈2N FLOPs-per-token forward rule and the ≈4N backward rule, and explain the factor of two from the two backward GEMMs.
  3. Classify every operation in a transformer block as GEMM or elementwise, and predict which ones are limited by arithmetic and which by memory traffic.
  4. Compute the persistent training-memory footprint of a model under mixed-precision Adam in bytes per parameter, and separate it from weight memory.
  5. Estimate activation memory for a given batch size and sequence length under a stated counting convention, and say whether a model fits on one 80 GB GPU.
  6. Find the sequence length at which attention's own matrix multiplies rival the linear layers, for a given configuration.

2.1 From text to numbers: tokens and embeddings

A language model does not see text. It sees a sequence of integers, each an index into a fixed vocabulary of V symbols. Tokenization — the rule that maps characters to those indices — is a preprocessing decision made once, before training, and it fixes the unit that every cost in this course is denominated in. For English prose a rough working figure is about four characters per token, which is an order-of-magnitude illustration rather than a constant; code, non-Latin scripts, and long numbers all tokenize much worse.

That matters more than it sounds. Latency is quoted per token. Throughput is tokens per second. Cost is dollars per million tokens. Context limits are token counts. If you change the tokenizer you change the denominator of every number in the system, and a comparison across two models with different tokenizers is not a comparison at all until you normalize.

The first layer of the network turns those integers into vectors. The embedding matrix has shape V × d_model, and the operation is a gather: look up row i, get a vector of d_model floating-point numbers. It is not a matrix multiply — it performs no arithmetic worth counting — but it does occupy V · d_model parameters, and it is read with terrible locality, since consecutive tokens in a batch index arbitrary rows.

At the other end the network produces, for each position, a vector of d_model numbers that has to become a probability over the whole vocabulary. That is a genuine matrix multiply against a d_model × V matrix, followed by a softmax. Some models tie the two matrices — the same weights used for input lookup and output projection — and some keep them separate. Our reference configuration keeps them separate, which costs an extra V · d_model parameters and matters when we count.

Everything between those two boundaries is the machine we spend the rest of the class accounting for.

Instructor notes

Minutes: 5. Board: Three boxes left to right: [ids][embed: V × d_model][... L blocks ...][unembed: d_model × V][logits]. Write V = 32000 and d_model = 4096 under the first and last. Ask the room: "Model A quotes 20 ms per token, model B quotes 18 ms. Which is faster?" Answer: unknown, until you know whether their tokenizers agree on how many tokens the same paragraph is. Expect confusion: Tokens get treated as words. Say: "A token is whatever the tokenizer says it is. Roughly four characters of English, and much worse on code."

2.2 What attention actually computes

Here is attention with no prerequisites. Every token in the sequence produces three vectors from its own representation: a query, a key, and a value. Think of the key as a label the token advertises, the value as the content it offers, and the query as the request the current token is making. To compute the output at position i, take that position's query, take the dot product against the key of every position up to and including i, turn those scores into weights that sum to one, and return the weighted average of the values.

That is a lookup table, made continuous. A hash table returns the one value whose key matches exactly. Attention returns a blend of all values, weighted by how well each key matches — and the blend is differentiable, which is why it can be learned. The dot product is the match score, and the softmax that normalizes the scores is what makes them weights.

Work one by hand. Take d_head = 2 and three tokens, and compute the output at position 3.

Attention at position 3, three tokens, d_head = 2

Query at position 3: q = [1, 0] Keys: k₁ = [1, 0], k₂ = [0, 1], k₃ = [1, 1] Values: v₁ = [1, 0], v₂ = [0, 1], v₃ = [1, 1]

Scores, q · kⱼ: 1, 0, 1 Scaled by 1/√d_head = 1/√2 = 0.7071: 0.7071, 0, 0.7071 Exponentiate: e^0.7071 = 2.028, e^0 = 1.000, e^0.7071 = 2.028 — sum 5.056 Weights: 2.028/5.056 = 0.401, 1.000/5.056 = 0.198, 0.401 — sum 1.000

Output = 0.401·[1,0] + 0.198·[0,1] + 0.401·[1,1] = [0.802, 0.599]

The query matched positions 1 and 3 and not position 2, so the output is mostly their values blended, with a fifth of v₂ mixed in. Nothing was selected; everything was weighted.

Two details that look like decoration and are not. The scaling by 1/√d_head exists because dot products of d_head-dimensional vectors grow with d_head, and large scores drive the softmax toward a one-hot distribution where gradients vanish. And the sum runs only up to position i, never past it — the causal mask — because a language model predicts the next token and must not see it. That mask is what makes generation sequential, and it is the reason Lecture 3's decode phase exists as a separate thing.

Real models run this many times in parallel. With h heads, each token produces h separate query/key/value triples of width d_head, attention runs independently per head, and the h outputs are concatenated back to width h · d_head = d_model. Different heads learn to attend to different things; from a systems point of view heads are just a batch dimension.

Now the reframing that matters for this course. Strip the interpretation and attention is: three matrix multiplies to produce Q, K and V from the input; one matrix multiply Q · Kᵀ to produce all scores at once; a softmax; one matrix multiply of the weights against V; and one more matrix multiply to project the result back to d_model. Five matrix multiplies and one softmax. The softmax is the only part that is not a GEMM, and §2.3 is about why that distinction is the one that decides performance.

Instructor notes

Minutes: 11. This is the section a student with no ML background either gets or does not. Board: Do the miniature by hand, digit by digit, including the exponentials. Do not display it finished. Then write the five matmuls in a column and circle "softmax" as the odd one out. Ask the room: "What happens to the output if I make all three scores equal?" Answer: a uniform average of the values — attention that has decided nothing. Useful for showing the softmax is doing selection. Expect confusion: Q, K, V get read as three different kinds of data. They are three linear projections of the same input vector, computed by three learned matrices. Say that sentence exactly. Common wrong answer: "The mask means we throw away half the work." In prefill the masked-out half is often computed and discarded anyway; the mask is a correctness device, not a performance one. Oct 1 makes it a performance one.

2.3 A transformer block is a small pile of matrix multiplies

A decoder-only transformer is L identical blocks stacked, each containing an attention sublayer and a feedforward (MLP) sublayer, each sublayer wrapped in a normalization and a residual add. Enumerate every operation in one block, for one token, and classify it.

OperationGEMM?FLOPs per token per layer
Normalization (×2)no2 · 4 · d_model = 32,768
Q, K, V projectionsyes3 · 2 · d_model² = 100,663,296
Attention scores Q·Kᵀ and weights·Vyes4 · d_model · S = 16,384·S
Softmax over scoresno≈ 5 · h · S = 160·S
Output projectionyes2 · d_model² = 33,554,432
Gate and up projectionsyes2 · 2 · d_model · d_ff = 180,355,072
Activation function and gating multiplyno≈ 5 · d_ff = 55,040
Down projectionyes2 · d_model · d_ff = 90,177,536
Residual adds (×2)no2 · d_model = 8,192

The GEMM rows sum to 404,750,336 FLOPs per token per layer, plus the sequence-dependent attention term. The non-GEMM rows, at S = 4096, sum to about 751,000 — 0.19% of the layer's arithmetic.

So the elementwise work is free. Except it is not, and the reason is bytes. A normalization reads d_model values and writes d_model values while doing about four operations per value: at bf16 that is 4 · 4096 = 16,384 FLOPs against 2 · 4096 · 2 = 16,384 bytes, an arithmetic intensity of exactly 1 FLOP per byte. Thursday will give you the H100's ridge point of 295 FLOP/byte; a kernel at intensity 1 runs at 1/295 of the machine's arithmetic peak, which is to say it is doing nothing but moving memory.

Count the traffic for a realistic step. Under a conservative accounting — each elementwise operation reads its inputs and writes its output once, nothing fused — the non-GEMM operations in one layer touch about 192,000 bytes per token. At a training microbatch of 32,768 tokens that is 6.29 GB per layer, against 0.405 GB of weights in that same layer. The cheap operations move fifteen times more bytes than the expensive ones. Across 32 layers it is 201 GB of traffic, about 60 ms on an H100's 3.35 TB/s, against roughly 1.1 s of forward-pass arithmetic at 40% of dense BF16 peak — five percent of the forward pass spent on 0.19% of the FLOPs.

That asymmetry is the entire reason kernel fusion exists. If the normalization, the activation, and the residual add can be executed while the data is still in registers or shared memory, the intermediate round trips to HBM disappear and the elementwise work costs almost nothing. How that is done — tiling, fusion, and the IO-aware attention kernel that applies the same idea to the softmax — is Sep 29 and Oct 1. Today the point is only that you can predict which operations will benefit, before writing any code, from a FLOP count and a byte count.

Instructor notes

Minutes: 7. Board: Two columns, "GEMM" and "not GEMM". Fill them from the block diagram, then write "0.19% of FLOPs" under the right column and "15× the bytes" under it. The gap between those two lines is the section. Ask the room: "Which of these operations would you optimize first?" Most say the big matmuls. Then show the byte count. Expect confusion: FLOP count is treated as a proxy for time. Say: "Time is max of arithmetic and traffic, not arithmetic alone. A kernel at 1 FLOP/byte is a memcpy with opinions." If short on time: Keep the intensity-of-1 calculation and the 15× line; drop the 5%-of-forward estimate.

2.4 Parameter accounting

Fix a configuration and use it for the rest of Part I. Call it the reference 7B: L = 32 layers, d_model = 4096, h = 32 attention heads, d_head = 128, d_ff = 11008, V = 32000, no grouped-query attention, untied input and output embeddings, no biases. It is deliberately ordinary — an open 7B-class decoder — and Lecture 3 reuses it without change.

Attention contributes four square matrices per layer: Q, K, V, and the output projection, each d_model × d_model. Note h · d_head = 32 · 128 = 4096 = d_model, so the per-head split costs nothing extra. The MLP contributes three matrices in the gated style now standard: a gate and an up projection, each d_model × d_ff, and a down projection d_ff × d_model.

Parameter count, reference 7B

Attention per layer: 4 · d_model² = 4 · 4096² = 4 · 16,777,216 = 67,108,864 MLP per layer: 3 · d_model · d_ff = 3 · 4096 · 11008 = 3 · 45,088,768 = 135,266,304 Per layer: 67,108,864 + 135,266,304 = 202,375,168 All layers: 32 · 202,375,168 = 6,476,005,376 Embeddings, untied: 2 · V · d_model = 2 · 32000 · 4096 = 262,144,000

N = 6,476,005,376 + 262,144,000 = 6,738,149,376 ≈ 6.74B parameters At b = 2 bytes (bf16): 13,476,298,752 bytes = 13.5 GB of weights

Thirteen and a half gigabytes is what the model costs to hold, before it does anything.

Three observations about where the mass sits. The MLP is 135,266,304 of the 202,375,168 per-layer parameters — two-thirds of the model is the feedforward network, not attention, which surprises people whose mental model of a transformer is the attention diagram. The embeddings are 262,144,000 of 6,738,149,376, or 3.9%; they are negligible here and would not be for a small model with a large vocabulary. And the normalization parameters — one vector of d_model per normalization, 65 of them — total 266,240, about 0.004% of N. We omit them and say so; that is the right kind of omission because it is four orders of magnitude below the leading term.

Now the sensitivity. Per-layer parameters are d_model · (4·d_model + 3·d_ff), and with d_ff ≈ 2.6875 · d_model that is 12.06 · d_model², so the body of the model is L · 12.06 · d_model². Parameters grow linearly in depth and quadratically in width. Doubling L to 64 gives 12.95B in the body; doubling d_model to 8192 gives 25.9B. If you want a bigger model, width is the aggressive lever.

Width is also the cheaper lever per parameter, for two systems reasons. Wider matrices mean larger GEMMs, which reach a higher fraction of peak; deeper models mean more sequential kernel launches, more layers on the critical path, and under tensor parallelism more collectives per token. And the KV cache — Thursday's subject — scales as L · n_kv · d_head, so quadrupling parameters by doubling d_model doubles the cache, while quadrupling them by quadrupling L quadruples it. Depth is not free, and the reasons it is not free are almost all systems reasons. Depth still buys something real about what the model can represent, which is why nobody builds a two-layer 7B, but the trade is being made against a real cost.

Instructor notes

Minutes: 7. Board: Derive all five lines live. Make them do 4 · 4096² in their heads and check against 67,108,864. Box N = 6,738,149,376 and 13.5 GB — both go up on the permanent side board. Ask the room: "Attention or MLP — which has more parameters?" Most say attention. It is two to one the other way. Expect confusion: "7B" is read as an exact 7,000,000,000. Say: "It is a marketing rounding of 6.74 billion. Every number we derive this term uses 6.738, not 7." If short on time: The count must survive intact. Cut the depth-versus-width paragraph to its first sentence.

2.5 Where the FLOPs go

Take one linear layer, Y = XW, with W of shape d_in × d_out, and one token of input. Each output element is a dot product of length d_in: d_in multiplies and d_in adds. There are d_out outputs, so the cost is 2 · d_in · d_out FLOPs — and d_in · d_out is exactly the number of parameters in W. Every parameter that participates in a matrix multiply costs two FLOPs per token: one multiply and one add.

forward FLOPs per token ≈ 2 · N

For the reference 7B that is 2 · 6.738e9 = 13.48 GFLOP per token, forward. Two honesty notes. The rule slightly over-counts, because the embedding lookup is a gather rather than a matmul, and treating its 131,072,000 parameters as if they multiplied adds 1.9% that is not there. It also under-counts, because it omits the attention score and value matmuls, whose cost depends on parameters not at all — and that omission is the interesting one.

Attention's own matmuls scale with sequence length. For a token at position S, per layer, the scores cost 2 · h · d_head · S FLOPs and the weighted sum of values costs the same, giving 4 · d_model · S = 16,384 · S FLOPs per token per layer. There are no parameters in that expression. It is pure data movement through arithmetic, and it grows without bound as the context does, which is why the linear-layer accounting that dominates at short context stops being the whole story at long context.

Find where they meet. Per layer per token, the linear layers cost 2 · 202,375,168 = 404,750,336 FLOPs and attention costs 16,384 · S. Setting them equal:

S_crossover = (per-layer parameters) ÷ (2 · d_model) = 202,375,168 ÷ 8192 = 24,704

At about 24,700 tokens of context, attention's matmuls cost as much as every weight matrix in the layer. Below that they are a correction; above it they are the dominant term, and because the cost is linear in S per token — quadratic in S for a whole prefill — it keeps growing. At S = 4096 attention is 16.6% of the layer's arithmetic. At S = 100,000 it is four times the linear layers. This is the precise sense in which long context stops being free, and it is worth noticing that it happens at a context length products now advertise routinely.

The closed form is worth keeping: the crossover is the per-layer parameter count divided by 2 · d_model, which for the standard d_ff = 4·d_model shape works out to exactly 2^15 = 32,768. Fatter MLPs push the crossover out; wider models push it out proportionally to their parameter growth. Note also what the crossover is not: it says nothing about memory. The bytes attention consumes — the KV cache — start hurting long before 24,704 tokens, and that is Thursday's §3.5.

Instructor notes

Minutes: 9. Board: Y = XW, one row of X, one column of W, count one multiply and one add per weight. Then 2N boxed. Then 4·d_model·S beside it, and solve for S in front of them. Ask the room: "How many FLOPs does the attention score matmul spend per parameter?" It has no parameters. That is the point, and the silence before someone says it is productive. Expect confusion: "Attention is quadratic" is remembered as a claim about the whole model. Say: "Quadratic in S for one term, out of a model whose other terms are linear. Which one wins depends on S, and here the answer is about 25,000." Common wrong answer: Students compute the crossover using 2N for the whole model instead of per layer. It lands at 25,704 — close enough that they will not notice the error, so name it.

2.6 The training loop, and where the factor of three comes from

Training is a loop: run the model forward on a batch, compare its predictions to the actual next tokens, compute how each parameter should change, change them, repeat. In PyTorch it is about ten lines.

for batch in loader:
    ids = batch["input_ids"].to(dev)      # H2D copy: B·S int64, trivial bytes
    logits = model(ids)                   # ≈2N FLOP/token; allocates the activation tape
    loss = cross_entropy(                 # logits are B·S·V — 2.10 GB at bf16, 4.19 GB at fp32
        logits[:, :-1].reshape(-1, V),
        ids[:, 1:].reshape(-1))
    loss.backward()                       # ≈4N FLOP/token; fills .grad, frees the tape
    optimizer.step()                      # no matmuls; ~28 bytes/param of traffic → 189 GB
    optimizer.zero_grad(set_to_none=True) # releases 2 bytes/param

Read that as a resource trace rather than as code. The forward pass costs 2N FLOPs per token and, more importantly, allocates every intermediate tensor the backward pass will need — that allocation is §2.7's subject and it is the line that decides your batch size. The loss line materializes a tensor of shape B · S · V; at a 32,768-token microbatch and V = 32000 that is 2.10 GB in bf16 and 4.19 GB if the softmax is done in fp32, which is routinely the single largest tensor in the step and routinely a surprise. The backward pass costs about twice the forward. The optimizer step does no matrix multiplies at all but streams roughly 28 bytes per parameter — 189 GB for the reference 7B, about 56 ms on an H100 — which is invisible at large batch and is not at small batch.

Now the factor of two. Consider Y = XW again, and suppose the backward pass has arrived carrying dY, the gradient of the loss with respect to Y. Two things are needed. The gradient with respect to the input, so the chain rule can continue down the network:

dX = dY · Wᵀ

and the gradient with respect to the weights, which is what the optimizer will consume:

dW = Xᵀ · dY

Both are matrix multiplies of the same shape class as the forward's single one, each costing 2 · d_in · d_out FLOPs per token. The backward pass does two GEMMs where the forward did one. That is the entire factor: ≈2N FLOPs per token forward, ≈4N backward, ≈6N for the pair.

C ≈ 6 · N · D FLOPs

That is Lecture 1's rule, and now it is derived rather than asserted. Lecture 1 used it to show that cumulative serving compute overtakes training compute at T = 3D generated tokens — a token generated costs 2N while a token trained on costs 6N, so three of the former equal one of the latter — and that result stands; we do not re-derive it here.

Where 6ND is wrong, and in which direction. It counts only the parameter-bearing matmuls, so it omits the attention score and value terms of §2.5, which matters above S ≈ 25,000 and is a rounding error below. It ignores activation recomputation, which is a deliberate trade of extra forward FLOPs for less memory and can add 30% or more to the true cost. It ignores everything a real project spends outside the successful run: failed runs, restarts from checkpoints, hyperparameter search, data pipeline work. And it says nothing about utilization6ND is a FLOP count, and converting it to wall clock requires an efficiency assumption that is usually somewhere between 30% and 50% and that you should never take from a paper without asking how it was measured. Every one of those corrections pushes the real cost up, none down.

Instructor notes

Minutes: 9. Board: The two backward equations, one above the other, with "same shape as the forward" written beside each. Then 2N + 4N = 6N and box C ≈ 6ND. This is the most reusable ten seconds of the lecture. Ask the room: "Why exactly two, and not three or one and a half?" Because there are exactly two things downstream needs: the gradient flowing further back, and the gradient of this layer's own weights. Expect confusion: Backward is imagined as a second forward pass in reverse. Say: "It is a different computation with a different FLOP count, and the count is two GEMMs per forward GEMM." If short on time: The two equations and the boxed 6ND. The caveat list can be assigned as reading.

2.7 Memory: four consumers, and the one that surprises people

Ask a student how much memory it takes to train a 7B model and most say fourteen gigabytes, because the weights are 13.5 GB. That answer is wrong by roughly an order of magnitude, and the gap is the most consequential thing in this lecture.

Four things occupy memory during training. Weights, the parameters themselves. Gradients, one number per parameter, produced by the backward pass. Optimizer state, which for Adam is two additional numbers per parameter — a running mean of the gradient and a running mean of its square — plus, under mixed precision, a full-precision master copy of the weights. And activations, the intermediate tensors saved during the forward pass because the backward pass needs them.

The first three are proportional to N and are the same every step. Under standard mixed-precision Adam:

Consumerbytes per parameter
bf16 weights (used in the matmuls)2
fp32 master weights4
Adam first moment, fp324
Adam second moment, fp324
persistent subtotal14
bf16 gradients2
working figure16

The subtotal is 2 + 4 + 4 + 4 = 14 bytes per parameter. Whether you carry 14 or 16 depends on whether your framework frees the gradient buffer inside the optimizer step; we carry 16 bytes per parameter and state that we are counting a bf16 gradient. Be explicit about the convention whenever you quote a figure like this, because the classic error in this area is not arithmetic — it is comparing someone else's 14 to your 16 without noticing that they counted different things.

For the reference 7B: 6,738,149,376 · 16 = 107,810,390,016 bytes = 107.8 GB, of which the weights are 13.5 GB. Weight memory is 12.5% of persistent training memory. Turn the same arithmetic around and the largest model whose persistent state alone fits in 80 GB is 80e9 ÷ 16 = 5 billion parameters — before a single activation, before a single token of data.

Activations are the fourth consumer and they behave differently: they scale with the number of tokens in the microbatch, not with N. Counting them requires a convention, because frameworks disagree about what they keep. Ours, stated so it can be attacked: per layer per token, in bf16, we save the residual stream entering the block, the normalized input to the QKV projections, Q, K and V, the attention output before the output projection, the residual stream entering the MLP, the normalized MLP input, the gate and up projections, and the gated product feeding the down projection.

Activation memory, reference 7B, one reasonable convention

Per layer per token: 6 · d_model + 3 · d_ff = 6 · 4096 + 3 · 11008 = 24,576 + 33,024 = 57,600 elements In bf16: 57,600 · 2 = 115,200 bytes = 112.5 KiB per layer per token All 32 layers: 115,200 · 32 = 3,686,400 bytes = 3.69 MB per token

At B = 1, S = 4096 (4,096 tokens): 3,686,400 · 4,096 = 15.1 GB At B = 8, S = 4096 (32,768 tokens): 3,686,400 · 32,768 = 120.8 GB

A single 4,096-token sequence already stores more activation bytes than the model has weight bytes.

This is one reasonable convention and not the only one. A framework with attention fused end to end saves less; one that keeps a separate copy of every normalization input and output saves more; one that materializes the S × S attention matrix saves catastrophically more. Treat 3.69 MB per token as the right order of magnitude for this configuration and re-derive it for whatever stack you actually measure.

The practical consequence is that activation memory is the term you control, because B and S are yours to choose and N is not. Microbatching splits the batch you want into pieces small enough to fit; gradient accumulation runs several microbatches, summing gradients into the same buffers, and calls the optimizer once, so the optimizer sees the large batch the training recipe asked for while memory only ever holds one microbatch of activations. The cost is that the weights and optimizer state are read once per accumulation step rather than once per large batch, which at small microbatch sizes starts to matter. Activation recomputation goes further: discard most activations during the forward pass and recompute them from checkpoints during the backward, trading roughly an extra forward pass of FLOPs for a large reduction in memory. That trade, and the parallelism strategies that make the persistent 107.8 GB somebody else's problem, are Oct 8 and Oct 6.

Instructor notes

Minutes: 9. Board: The bytes-per-parameter table, built one row at a time with the question "what else does Adam need?" between rows. Write 107.8 GB and 13.5 GB one above the other and draw the ratio. Ask the room: "How much memory to train a 7B model?" Take the answers before showing the table. Almost everyone says 14 GB. Expect confusion: The master weights are read as a redundant copy that a careful implementation could drop. Say: "It is the only copy with enough mantissa bits to accumulate a small update. §2.10 shows the arithmetic." Common wrong answer: "Activations are small because each one is small." Each one is a few kilobytes and there are L per token and tens of thousands of tokens.

2.8 A worked memory budget, and the punchline

Put it together and answer the question a practitioner actually asks: does the reference 7B train on one 80 GB GPU?

Training footprint, reference 7B, mixed-precision Adam, one H100 80GB

bf16 weights: 6,738,149,376 · 2 = 13.48 GB bf16 gradients: 6,738,149,376 · 2 = 13.48 GB fp32 master weights: 6,738,149,376 · 4 = 26.95 GB Adam first moment: 26.95 GB Adam second moment: 26.95 GB Persistent subtotal: 107.81 GB

Activations at the smallest useful microbatch, B = 1, S = 4096: 15.10 GB Logits tensor, B·S·V at bf16: 4,096 · 32,000 · 2 = 0.26 GB

Total ≈ 123.2 GB against 80 GB of HBM.

It does not fit — and it does not fit by 28 GB even if you set the batch to zero.

That last clause is the punchline, so sit with it. Activations are not the problem here. The persistent state alone, 107.8 GB, exceeds the card by 35%, and no batch size, no sequence length, no recomputation strategy, and no fusion changes that number by a byte. It is N · 16, and the only levers on it are a smaller model, a smaller optimizer, or more GPUs.

The choice the field made is more GPUs, and the observation that gets you there is that the 94 GB of gradients and optimizer state is not needed in full on every device. Each of eight GPUs could hold an eighth of the optimizer state and an eighth of the gradients, exchanging what is needed when it is needed — trading interconnect bandwidth for memory capacity. That is ZeRO's idea, and together with the orthogonal question of splitting the model itself across devices (Megatron-LM) it is the optional-content track on the course page, not a lecture. Do not go looking for those mechanisms today. Take away only that distributed training is not a scaling optimization; it is a correctness requirement, imposed by a single division. Nobody chose to make training distributed. Sixteen bytes per parameter did.

One arithmetic check that is worth running yourself. If persistent state is 16 bytes per parameter, then a 70B model needs 1.12 TB before activations — fourteen 80 GB cards' worth of memory, held by a model whose weights are 140 GB. The ratio never improves with scale; it is a constant.

Instructor notes

Minutes: 6. Board: Stack the five persistent rows, sum them, write "80" underneath, and draw the line. Say nothing for three seconds. Ask the room: "What do you cut?" Let them propose smaller batches, then point at the subtotal. The realization that batch size is irrelevant here is the moment. Expect confusion: Students assume gradient checkpointing or a smaller sequence length rescues this. Say: "Those touch the 15 GB, not the 108." If short on time: Nothing in this section is cuttable. Cut §2.9 instead.

2.9 What a framework actually does for you

You wrote ten lines in §2.6 and never described the backward pass. Something computed it, and that something is worth understanding, because its design decisions show up as performance.

The mechanism is autograd as a tape. As the forward pass executes, the framework records each operation and the tensors it consumed onto a directed graph. Calling .backward() walks that graph in reverse, and at each node applies the local derivative rule — the two GEMMs of §2.6 for a linear layer, and an analogous rule for every other operation — chaining gradients from the loss back to every parameter. Two consequences follow directly. The tape holds references to the forward's intermediate tensors, which is why activations are a memory consumer at all rather than transient scratch. And the tape is built from the control flow that actually ran, so a Python if in the model is recorded as whichever branch it took.

The second thing a framework does is launch kernels, and launches are not free. Each one costs the host CPU some microseconds of dispatch work to enqueue onto the GPU stream. Treat that as order 5 µs in eager mode — a clearly-flagged order of magnitude, not a measurement, and one you should measure on your own stack before relying on it.

Count them. A reasonable eager implementation of one transformer layer forward issues on the order of 15 kernels: two normalizations, three projection GEMMs, a rotary embedding application to Q and K, a fused attention call, an output GEMM, two residual adds, two MLP GEMMs, an activation, a gating multiply, and a down GEMM. Backward roughly doubles that, so call it 45 per layer per step. Across 32 layers that is 1,440, plus perhaps 60 more for embeddings, the final normalization, the logits GEMM, and the loss — about 1,500 launches per training step, before the optimizer. A foreach-style fused Adam adds a handful; a naive per-tensor implementation issuing ten elementwise kernels for each of ~291 parameter tensors adds nearly 3,000.

Now compare against the step's compute. At B = 8, S = 4096 the step processes 32,768 tokens, so C = 6 · 6.738e9 · 32,768 = 1.32e15 FLOPs, and at 40% of the H100's 989 TFLOP/s dense BF16 peak that is 3.35 s. Against 3.35 seconds, 1,500 launches at 5 µs is 7.6 ms — 0.2%, invisible.

Shrink the step and the picture inverts. Setting 1,500 · 5 µs equal to 6 · N · D_step ÷ 3.956e14 gives D_step75 tokens. Below roughly 75 tokens per step, this model spends more time launching kernels than executing them. That regime is not hypothetical: it is exactly Thursday's decode phase, where each step processes one token per sequence. A forward-only eager pass over 32 layers is on the order of 490 launches, about 2.45 ms of dispatch, against Lecture 3's 4.0 ms memory-bandwidth floor for the same step. Launch overhead alone can be half the budget of a decode step. The answers — CUDA graphs, which record a launch sequence once and replay it as a unit, and compilation, which fuses many small kernels into few large ones — are Sep 29 and Oct 15 material.

Which brings us to the two assigned system papers, and the argument between them. TensorFlow asked the user to declare the computation as a graph, then compiled and optimized that graph before running it. Given a whole graph a system can fuse operations, plan memory, place work across devices, and eliminate exactly the launch overhead just counted. PyTorch executed operations immediately as Python called them, giving up the whole-program view in exchange for a program you can debug with a print statement and a stack trace and a Python if.

PyTorch won, and it is worth being precise about why, because the reason is a systems lesson rather than an ML one. It did not win on throughput; the graph-first design had every structural advantage on throughput. It won because researchers iterate, and the cost of iteration — write, run, misread the error, fix, run again — dominated the cost of execution for the work that was actually being done. The system optimized for the workload's real bottleneck, which was the human. The PyTorch paper is candid about this being a deliberate trade rather than a free lunch, and its engineering sections are about how much performance you can recover inside an eager design.

The postscript is that graphs came back. torch.compile, CUDA graphs, and the tracing compilers now standard in serving stacks all recover the whole-program view — but as an opt-in applied to an already-working eager program, rather than as the price of entry. The winning design was eager first, graph second, and that ordering is the thing to remember.

Instructor notes

Minutes: 5. Board: 1,500 launches × 5 µs = 7.6 ms above step compute = 3.35 s, then erase the second and write 75 tokens under it. The inversion is the section. Ask the room: "TensorFlow could fuse across the whole graph and PyTorch could not. Why did PyTorch win?" Push past "it was easier" to "iteration speed was the binding constraint, and they optimized the right resource." Expect confusion: Eager versus graph is heard as a settled historical question. Say: "It settled as eager-first with graphs bolted back on. Both halves of that sentence matter." If short on time: Keep the 75-token crossover and one sentence on the design bet.

2.10 Mixed precision, only as far as the accounting needs

We have been assuming b = 2 bytes per parameter for the weights that participate in matmuls, and b = 4 for the master copy and the optimizer moments. Two questions follow: which 16-bit format, and why keep a 32-bit copy at all.

Both fp16 and bf16 use 16 bits and split them differently. fp16 spends 5 bits on the exponent and 10 on the mantissa, giving good precision over a narrow range that tops out around 65,504 and, more dangerously, underflows to zero for small values. bf16 spends 8 bits on the exponent — the same as fp32 — and 7 on the mantissa, so it covers fp32's full dynamic range with about three significant decimal digits. Gradients in a deep transformer span many orders of magnitude, and in fp16 the small ones silently become zero. The historical workaround is loss scaling: multiply the loss by a large constant before the backward pass to lift gradients into fp16's representable range, unscale before the optimizer step, and back off when an overflow is detected. It works, it is fiddly, and it is a source of silent divergence. bf16 needs none of it, because it never had the range problem, and that is essentially the whole reason it displaced fp16 for training.

The master copy answers the second question. bf16 has 8 bits of mantissa precision, so the gap between representable numbers near a value w is about w · 2⁻⁸ ≈ 0.39% of w. An Adam update has magnitude roughly the learning rate, since the moment ratio is normalized to order one. Early in training, with a learning rate of 3e-4 against weights of order 0.02, each update is about 1.5% of the weight — comfortably above the rounding threshold. But learning-rate schedules decay by one to two orders of magnitude, and a late update of 0.015% of the weight, applied to a bf16 number, rounds to no change at all. Every step of the tail of training would be a no-op. The fp32 master copy, with relative resolution near 6e-8, accumulates those updates; the bf16 copy used in the matmuls is derived from it after each step. Four bytes per parameter is what it costs to keep the last third of a training run from doing nothing.

This is worth separating cleanly from a topic that sounds identical. Quantizing a model for inference — Oct 20 and Oct 22 — takes a finished set of weights and asks how few bits can represent them while the outputs stay acceptable. There is no accumulation of small updates, no optimizer, and no schedule; the only question is output quality. Training precision is about whether a number can absorb a small change, inference precision is about whether a number is close enough. They share a datatype table and nothing else.

The last observation is the one that pays off in Lecture 3. Precision is a knob on bytes, and bytes have shown up in every section today: 16 bytes per parameter of persistent state, 3.69 MB per token of activations, 189 GB of optimizer traffic, 201 GB of elementwise traffic. On Thursday the same knob turns the weight-streaming floor and the KV-cache footprint. Once you have internalized that almost every quantity in this course is a byte count divided by a bandwidth, the rest of the term is bookkeeping with high stakes.

When you read a TFLOP/s figure on a datasheet, check whether it is the dense number or the 2:4 structured-sparse number. NVIDIA quotes sparse figures prominently and they are exactly double. Dense is what a transformer gets: 312 TFLOP/s BF16 on an A100 80GB SXM, 989 on an H100 SXM. Every calculation in this course uses the dense row.

Instructor notes

Minutes: 3. Board: bf16: 8 exponent bits, 7 mantissa over fp16: 5 exponent, 10 mantissa, then 2⁻⁸ ≈ 0.39% and lr/|w| ≈ 0.015% late in training. Two lines, then stop. Ask the room: "If bf16 is less precise than fp16, why did it win?" Range beats precision when your gradients span ten orders of magnitude. Expect confusion: Mixed precision is assumed to halve training memory. It does not touch the 16 bytes per parameter — Exercise 2 makes them prove it. It halves activations and raises arithmetic throughput by 2× against TF32 tensor and 15× against non-tensor FP32.

Key takeaways

  • A transformer is five matrix multiplies and a softmax per attention sublayer, three matrix multiplies per MLP, and a handful of elementwise operations. The elementwise work is 0.19% of the FLOPs and moves fifteen times more bytes than the weights, which is why fusion exists.
  • The reference 7B is L = 32, d_model = 4096, d_ff = 11008, V = 32000: N = 6,738,149,376 parameters and 13.5 GB of bf16 weights. Two-thirds of the parameters are in the MLP, not attention.
  • Every parameter in a matmul costs two FLOPs per token forward. Backward needs two GEMMs where forward needed one, so the total is ≈6N per token and C ≈ 6ND for a training run.
  • Attention's own matmuls carry no parameters and scale with S. For this configuration they equal the linear layers at S = 24,704 — the point where long context stops being a rounding error.
  • Mixed-precision Adam costs 16 bytes per parameter of persistent state, of which the weights are 2. The reference 7B needs 107.8 GB before a single activation, which is why distributed training is a requirement rather than an optimization.
  • Activations scale with tokens per step, not with N: about 3.69 MB per token here, so one 4,096-token sequence already outweighs the model's weights. That is the term microbatching, gradient accumulation, and recomputation exist to control.

Numbers worth memorizing

QuantityValueSource
Reference 7B parameters6,738,149,376 ≈ 6.74B32 · 202,375,168 + 262,144,000
Reference 7B bf16 weights13.5 GBN · 2 bytes
MLP share of parameters2/3 (135.3M of 202.4M per layer)d_model·d_ff vs 4·d_model²
Forward FLOPs per token2N = 13.5 GFLOPone multiply + one add per weight
Training computeC ≈ 6ND2 forward + 4 backward
Attention/linear crossoverS = 24,704202,375,168 ÷ (2 · 4096)
Persistent training state, Adam mixed precision16 bytes/param (14 without the gradient)2 + 2 + 4 + 4 + 4
Reference 7B persistent footprint107.8 GB, against 80 GB of HBMN · 16
Activations, reference 7B≈ 3.69 MB per token(6·d_model + 3·d_ff) · 2 · L

Self-check

  1. Why is the backward pass about twice the forward, rather than equal to it or three times it?Because each forward GEMM Y = XW requires two backward GEMMs of the same shape: dX = dY·Wᵀ to continue the chain rule downward, and dW = Xᵀ·dY to give the optimizer something to consume. Two, because those are exactly the two quantities anything downstream needs.
  2. A 13B model under mixed-precision Adam. Persistent state in bytes, and how many 80 GB GPUs does it need before activations?13e9 · 16 = 208 GB, so three 80 GB cards hold it with 32 GB to spare for activations, workspace, and fragmentation — and that spare is thinner than it looks. Weights are only 26 GB of the 208.
  3. Double d_model from 4096 to 8192, holding L, d_head, and the d_ff/d_model ratio fixed. What happens to N, to FLOPs per token, and to KV-cache bytes per token?N in the body quadruples (per-layer parameters go as d_model²) and the embeddings double, so N goes from 6.74B to about 26.4B. FLOPs per token track N, so they roughly quadruple. KV bytes per token are 2·L·n_kv·d_head·b, and doubling d_model at fixed d_head doubles the head count, so the cache only doubles. Width buys parameters more cheaply than it buys cache.
  4. If the matmuls run in bf16, why keep an fp32 copy of the weights?bf16 resolves changes of about 0.39% of a value. Late in training, after the learning rate has decayed one or two orders of magnitude, a typical Adam update is smaller than that and would round away entirely. The fp32 master copy accumulates them; the bf16 copy is regenerated from it each step.
  5. At a microbatch of 32,768 tokens, which is larger for the reference 7B — the weights or the activations, and by how much?Activations: 3.69 MB/token · 32,768 = 120.8 GB against 13.5 GB of weights, a factor of about 9. Weights are fixed; activations are the term you control by choosing B and S.
  6. Is the output projection to logits a GEMM? Is the embedding lookup?The output projection is a real GEMM, d_model × V, costing 2 · 32000 · 4096 = 262 MFLOP per token — about 65% of one transformer layer. The embedding lookup is a gather with no arithmetic, which is why the 2N rule over-counts by about 1.9% here.

Exercises

  1. Grouped-query attention, recounted. Redo §2.4's parameter count with n_kv = 8 key/value heads instead of 32, keeping 32 query heads and d_head = 128. Report the new per-layer count, the new N, the bf16 weight bytes, and the percentage reduction. Then state precisely what does not change: the query head count, the attention score and value FLOPs of §2.5, and the crossover sequence length — and say what does change dramatically, with the number. Solution sketch: Q and the output projection stay d_model² = 16,777,216 each; K and V shrink to d_model · (n_kv · d_head) = 4096 · 1024 = 4,194,304 each. Attention per layer = 2·16,777,216 + 2·4,194,304 = 41,943,040, down from 67,108,864. Per layer 41,943,040 + 135,266,304 = 177,209,344; × 32 = 5,670,699,008; + 262,144,000 = N = 5,932,843,008 ≈ 5.93B, or 11.87 GB in bf16 — 12.0% fewer parameters (805,306,368 saved). Unchanged: all 32 query heads still attend over all S positions, so 4·d_model·S per layer per token is identical, and the crossover moves only because per-layer parameters fell — 177,209,344 ÷ 8192 = 21,632. What changes dramatically is the KV cache: 2·32·8·128·2 = 131,072 bytes per token, 128 KiB instead of 512 KiB, a 4× reduction, which is the entire reason the architecture exists and is Lecture 3's §3.5.
  2. fp32 throughout. Redo §2.8's budget with no mixed precision: fp32 weights, fp32 gradients, fp32 Adam moments, fp32 activations. Report the persistent bytes per parameter, the total persistent footprint, activations at B = 1 / S = 4096, and the peak arithmetic throughput available. Then answer the question the exercise is really asking: what does mixed precision actually save? Solution sketch: Pure fp32 needs weights 4 + gradients 4 + first moment 4 + second moment 4 = 16 bytes per parameter, with no master copy required — identical to mixed precision's 16, so persistent state is 107.8 GB either way. Activations double: 57,600 · 4 · 32 = 7.37 MB per token, so 4,096 tokens cost 30.2 GB instead of 15.1 GB, and the total rises from 123.2 GB to about 138.5 GB. The real saving is elsewhere: arithmetic drops from 989 TFLOP/s dense BF16 to 495 TFLOP/s TF32 tensor (2.0×) or 67 TFLOP/s FP32 non-tensor (14.8×), and every activation-carrying byte of memory traffic doubles. Mixed precision is a throughput-and-activations optimization, not an optimizer-state optimization — a result most people guess backwards.
  3. Is the step launch-bound? Using §2.9's conventions — 15 kernels per layer forward, 2× that backward, 32 layers, ~60 extra launches for embeddings and loss, a fused optimizer, and 5 µs per launch — compute total launches and launch time per step. Then compute the step's compute time at 40% of the H100's dense BF16 peak for B·S = 32,768 tokens and for B·S = 512 tokens, and give the launch-overhead fraction in each case. Finally, find the tokens-per-step at which the two are equal. Solution sketch: Launches = 32 · (15 + 30) + 60 = 1,500; at 5 µs that is 7.5 ms. Compute: C = 6 · 6.738e9 · 32,768 = 1.325e15 FLOP ÷ (0.4 · 989e12) = 3.35 s, so launches are 0.22%. At 512 tokens: C = 2.07e13 ÷ 3.956e14 = 52.3 ms, so launches are 14% — the same model, the same code, an order of magnitude worse. Equality at 6 · 6.738e9 · D = 7.5e-3 · 3.956e14 → D75 tokens per step. An unfused per-tensor Adam adds roughly 2,900 launches and moves that threshold to about 220 tokens.
  4. Move the crossover. Recompute §2.5's attention/linear crossover for d_ff = 4 · d_model = 16384, holding everything else fixed. Report the new per-layer parameter count, the new N, and the new crossover. Then explain why the crossover moved by almost exactly the same factor as N. Solution sketch: Per layer = 4·4096² + 3·4096·16384 = 67,108,864 + 201,326,592 = 268,435,456; × 32 = 8,589,934,592; + 262,144,000 = N = 8,852,078,592 ≈ 8.85B. Crossover = 268,435,456 ÷ (2 · 4096) = 32,768 tokens, up from 24,704 — a factor of 1.327, against N's factor of 1.314. They track because the crossover is per-layer parameters ÷ (2·d_model), and d_model is fixed, so the crossover is proportional to per-layer parameters; the tiny discrepancy is the fixed 262M of embeddings diluting N's growth. Fattening the MLP pushes the point where attention starts to matter further out, in exact proportion to the parameters you added.
  5. From 6ND to break-even. Train the reference 7B at the Chinchilla ratio D ≈ 20N. Compute D, the training FLOPs, the H100-hours at 40% utilization, and the board-level energy at the 700 W limit. Then state the break-even output from Lecture 1's T = 3D and check it against the figure Lecture 3 quotes. Solution sketch: D = 20 · 6.738e9 = 1.348e11 tokens (135B). C = 6 · 6.738e9 · 1.348e11 = 5.45e21 FLOP. At 0.4 · 989e12 = 3.956e14 FLOP/s: 5.45e21 ÷ 3.956e14 = 1.377e7 s = 3,826 H100-hours, which at an illustrative $2/GPU-hour — substitute your own rate — is order $7,700, and at 700 W is 3,826 · 0.7 = 2,678 kWh of board energy, boards only. Break-even is T = 3D = 4.04e11 generated tokens, about 404 billion, matching Lecture 3's ≈405 billion. The check that matters: serving those tokens costs 2 · 6.738e9 · 4.04e11 = 5.45e21 FLOP, equal to the training run, as T = 3D requires.

Reading guide

Attention Is All You Need — optional. Read §3, the model architecture, and nothing else on the first pass. You want the shapes: what Q, K, and V are, why the scores are divided by √d_head, and how multi-head attention concatenates. Figure 2 is the one to study — the left panel is §2.2's miniature drawn as a dataflow. Skip the machine translation results, the training details, and the positional-encoding discussion, all of which have been superseded. Hold this question: the paper's model is an encoder-decoder and ours is decoder-only, so which parts of §3 does a modern language model actually keep?

PyTorch and TensorFlow — optional, and read as a pair. Read TensorFlow's §2 and §3 for what a dataflow graph buys you — placement, fusion, and whole-program optimization — then read PyTorch's §2 and §3 for the argument that giving all of that up was the right call. In the PyTorch paper the sections that matter to us are the ones on how much performance an eager design can recover: the caching allocator and the multiprocessing and autograd internals. Skip both papers' benchmark tables; the hardware is a decade old and the numbers do not transfer. Hold this question across both: what was each system optimizing, and was it the resource that was actually scarce? That is the question How to Read a Paper calls the first pass, and it is what the paper discussion guide will ask presenters for all term.

Looking ahead

Thursday, September 10 takes today's N, today's 13.5 GB, and today's 2N per token and runs the model forwards only, splitting inference into prefill and decode, introducing the roofline, and showing that the binding constraint is memory bandwidth rather than arithmetic. Everything deferred today has a date: fusion and the memory hierarchy on Sep 24, Sep 29, and Oct 1, where §2.3's 15× byte ratio becomes an actual kernel; parallelism and ZeRO, which are the answer to §2.8's 107.8 GB, in optional content rather than in lecture; batching and scheduling on Oct 6 and Oct 8, where §2.9's launch counting reappears as CUDA graphs; and quantization on Oct 20 and Oct 22, which is §2.10's question asked about a finished model instead of a training run. Between now and then, Sep 15 and Sep 17 look at agents, whose token bills you can already price with 2N — and Assignment 1 goes out on Sep 15 (see assignments).