Senior interview questionAsked at Anthropic · interview write-up
Estimate an FFN layer's FLOPs, parameter and activation bytes, and the communication from sharding it. Which limit binds: compute, HBM or network?
Model answer
Write the three costs down, then compare times.
# d model, f hidden (often 4d), T tokens, bf16
flops = 2 * (2 * T * d * f) # two matmuls
w_bytes = 2 * d * f * 2 # both weights
act_bytes = T * (2 * d + f) * 2 # in, out, hidden
# tensor parallel over n GPUs: one all-reduce of T*d
comm = 2 * (n - 1) / n * T * d * 2 # bytes per GPU
t_compute = flops / n / peak_flops
t_hbm = (w_bytes / n + act_bytes) / hbm_bw
t_net = comm / link_bw
Scroll sideways to see the whole diagram
Every forward pass has to bring the model's weights from memory to the arithmetic units: about 498 MB for GPT-2 small in 32-bit floats.
Text version
- Every forward pass has to bring the model's weights from memory to the arithmetic units: about 498 MB for GPT-2 small in 32-bit floats.
- A decode step does that whole read to produce one token for one sequence, plus that sequence's KV cache. About 2 operations per weight, 4 bytes per weight: half an operation per byte.
- Decode 8 sequences together and the same read serves all 8. Only the KV caches grow with the batch.
- Prefill pushes every prompt token through in one pass, so one read of the weights serves all 960.
- Operations per byte decides the limit. Few, and the chip waits on memory; many, and it waits on arithmetic. The next pages measure both on one laptop.
GPT-2 small (124,439,808 weights, float32); KV cache from the Module 01 formula at a 128-token context.
The largest time binds. With few tokens, as in decode, intensity is roughly T FLOPs per byte, below a GPU's FLOPs-to-bandwidth ratio, so HBM binds; big prefill batches turn compute-bound. Sharding divides compute and weights by n but not the all-reduce, so over slow inter-node links the network binds.
Practise the rest
10 Where the time goes questions, 492 in all, each with a model answer. Free.
More from Anthropic interview questions.
More Where the time goes interview questions
- What's the difference between compute-bound and memory bandwidth-bound work, and which one is decode?
- How would you show, on a given machine, that decode is bandwidth-bound and prefill compute-bound?
- Why does decoding a batch of 8 cost far less than 8 separate decode steps, and why can't the batch grow forever?
- You're sizing an 80 GB accelerator for Llama 2 13B in 16-bit with 2,048-token conversations. How many fit, and what would you change if that's not enough?
- Your chat app's replies range from 8 to 256 tokens. Which batching strategy would you pick, and what did the simulation show?
- Explain the optimisations an ML compiler makes, such as kernel fusion, and how they map to the hardware.