Unit 7.6 · about 4 h
219 / 324Parameter, FLOP and memory accounting for attention
Attention scales quadratically in sequence length. The budget must be derived before launch, because the out-of-memory error arrives after the design is fixed.
- derive
- measure
You can now
- Derive attention and MLP parameter, FLOP, score-memory and activation costs across batch, heads, sequence and width.
- Measure score materialization for a stated shape and report it in correct binary units.
- Explain which term dominates in short- versus long-sequence regimes and what the math-backend upcast changes.
Covers
- Parameter and FLOP accounting
- Score memory, activation storage and data movement
- Quadratic score materialization
Study, in this order
3 steps
- Complexity, sequential ops and path lengthsD2L - attention and transformer implementation route
Comparing CNNs, RNNs, and Self-Attention: \(O(k n d^2)\) with \(O(1)\) ops and \(O(n/k)\) path for CNNs, \(O(n d^2)\) with \(O(n)\) ops and path for RNNs, \(O(n^2 d)\) with \(O(1)\) ops and path for self-attention
- Quadratic attention and model sizesUnderstanding Deep Learning - Transformers
Section 12.2.1 (p. 209) for \(N^2\) attention weights independent of width; Section 12.9 (pp. 227 to 228) for roughly half the interactions in decoders and the sparsification gallery; BERT and GPT-3 parameter listings (pp. 219 and 224 to 225) as worked sizes
Depth · Treat the sparsification gallery as orientation only; sparse kernels are a later topic.
- Backends and precision behaviorPyTorch - scaled dot product attention
The three-implementation Note with backend selection functions, and the Numerical accuracy paragraph: math backend keeps
fp32intermediates forhalfandbf16inputs, and outputs may differ by backendDepth · Skip grouped-query attention; it belongs to decoder specialization.
Focus on
- Count with multiply-add as 2 FLOP: QKVO projections \(8 n d^2\), scores plus weighting \(4 n^2 d\), width-\(4d\) MLP \(16 n d^2\). At \(n\) 512 and \(d\) 1024 that is about 5.37 GFLOP for attention and 8.59 GFLOP for the MLP.
- Score memory is heads times \(n^2\) entries in the working dtype: 16 heads at \(n\) 512 in
fp32is exactly 16 MiB, 1 MiB per head. At \(n\) 4096 the same block needs 1 GiB. - Causal masking keeps \(n(n+1)/2\) pairs, so the necessary score FLOPs roughly halve, as UDL notes. A dense matmul followed by a mask still computes all \(n^2\) scores; only a kernel that skips masked tiles saves the work, and only a fused kernel avoids materializing the scores.
- The math backend upcasts
halfandbf16intermediates tofp32, so measured bytes can exceed a naive dtype-times-entries count. Fused backends can change numerics without changing the asymptotic count.
Reproduce closed-book
fp32: attention parameters 4.19M, MLP parameters 8.39M, score memory 16 MiB, attention forward about 5.37 GFLOP. Then scale scores to n 4096 and show 1 GiB.Traps
- Mixing MACs with FLOPs. One multiply-add is 2 FLOP; state the convention or totals disagree by exactly 2x.
- Quoting GB for GiB quantities. 16 times \(4096^2\) times 4 bytes is \(2^{30}\) bytes: 1 GiB, not 1 GB.
- Trusting fused-kernel numerics to match the math path bit for bit. The docs warn outputs may differ by backend; test within tolerance, not equality.
Practice
D2L Attention Scoring Functions: Exercise 3 (cost scaling)
Exercises: cost versus key, query and value dimensionality and count, plus the memory-bandwidth follow-up
Answer with the derived \(8nd^2\) plus \(4n^2d\) model and the 512-case numbers, stating the multiply-add convention.
TensorTonic: Transformer Training FLOP Estimator
The estimator problem in the PyTorch and einops group of the free language-modeling-from-scratch plan
It estimates one training step from forward matrix multiplications plus a supplied forward attention cost. Match its forward matmul terms to your \(8nd^2\) and \(16nd^2\) counts, then document how it counts the backward pass and every term it includes or leaves out.
Self-test
-
How many score bytes do 8 heads at \(n\) 1024 in
fp32need?8 times \(1024^2\) entries times 4 bytes equals 33,554,432 bytes, which is 32 MiB. -
Why does decoder attention need only about half the encoder score FLOPs, and when is the saving real?
Causality keeps the lower triangle: \(n(n+1)/2\) pairs instead of \(n^2\), roughly half for long sequences. A dense matmul plus mask still spends all \(n^2\); the saving is real only in a kernel that skips masked blocks. -
How many bias-free attention parameters does one BERT-width layer hold?
4 times \(1024^2\) equals 4,194,304, about 4.19M. Biases would add 4 times 1024 on top. -
Sequence length grows from 512 to 4096 at fixed width. What scales by 64, and what becomes the main cost?
Score entries, bytes and score FLOPs all scale by \(8^2\) equals 64, while projection FLOPs scale by 8 and weights stay fixed. The score tensor grows from 16 MiB to 1 GiB per layer and example, so it is usually the first memory limit. Since \(n\) is now above \(2d\), score FLOPs also exceed the projection FLOPs.
Completion criterion
Stop when
Your hand table matches measured score bytes and a profiler FLOP count within stated conventions, and you can name the dominant attention term on either side of \(n\) equals \(2d\) (\(6d\) once the MLP is included).
Mark passed, then open 7.7 Only mark this when you can do it without notes.Leads to
- 7.7 Architecture trade-offs and attention correctness tests Turn the cost model into architecture choices and tests.
- 8.3 GPU memory: allocated, reserved, peak and memory models Reconcile the hand count with measured allocator categories.
Comes after
- 7.3 Self-attention, cross-attention and multi-head shapesCount the projections and scores you just shaped.
- 7.5 Transformer blocks: residual stream, normalization and feed-forwardCount the block you just built.
Resources for this unit
- Understanding Deep Learning - Transformers Use Chapter 12.1-12.11 and the matching notebooks for dot-product self-attention, extension to transformers, transformer layers and vision transformers.
- D2L - attention and transformer implementation route Use queries-keys-values, attention-scoring-functions, multihead-attention, self-attention-and-positional-encoding, transformer and vision-transformer. Implement masked single-head attention, multi-head attention and one pre-normalized block before comparing with PyTorch.
- PyTorch - scaled dot product attention Use the official function documentation as a differential oracle for shapes, masks, dropout and backend behavior after the manual version passes intervention tests.
Also in this section
- Deep-ML candidate - attention problems Use 490 Scaled Dot-Product Attention as the fifth platform admission sample. If the judge catches a seeded masking error, use at most one of 490, 491, 53, 904 or 905 as a changed-prompt drill.
- TrenTorch - attention and block questions Solve the multi-head split, sinusoidal positional encoding, causal mask and pre-norm versus post-norm questions the units name, then carry the same checks into your own block.