How would you optimize large-scale training/inference?

In this article

1. Before you start
2. Why interviewers ask this
3. Large-scale optimization from first principles
4. The Staff-level answer
5. The rubric
6. Three variants of this question
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page sits inside the large language model (LLM) inference optimization guide, and it assumes you have met the parts below.

Why interviewers ask this

An answer that lists techniques without an order is the usual way this question goes wrong. A candidate who says "use ZeRO, FlashAttention, quantization and batching" has named tools but not shown when each one helps, and a tool applied to the wrong limit changes nothing. Interviewers use the question to separate people who memorized a list from people who debug a slow system by measuring it.

A weak answer is a list of buzzwords. A strong answer splits the problem into training and serving, names the limit in each, and pairs each limit with one fix. A Staff-level answer adds the cost of every fix, such as extra communication or lower numeric precision, and says how it would confirm the gain. Candidates report this question in system design and machine learning rounds; no public source ties it to a named company.

Large-scale optimization from first principles

Model states

Model states are the three things a training run must keep for every parameter: the parameter itself, its gradient, and the optimizer's extra running values, which together usually take far more memory than the parameters alone.

Model FLOPs utilization (MFU)

Model FLOPs utilization (MFU) is the throughput a training run achieves, counted as the floating-point operations the model itself needs per token (excluding recomputation), divided by the hardware's theoretical peak floating-point operations per second (FLOP/s); the PaLM paper (Chowdhery et al., 2022) introduced it, and a low value means the GPUs spend time waiting or redoing work rather than advancing the model.

In plain words

Picture a restaurant kitchen. Orders pile up. The cause might be a small fridge, a slow oven, cooks waiting on each other, or a narrow pass-through window to the dining room. Buying a bigger oven fixes only the oven case.

Large models work the same way. A GPU can run out of memory, run out of arithmetic speed, wait on other GPUs, or wait on its own memory. The analogy breaks in one place: in a GPU the limit can move after each fix, so you measure again.

In one precise sentence, you find which of memory capacity, compute, communication or memory bandwidth limits the run, apply the fix for that limit, and re-measure.

Step 1: Find the bottleneck before changing anything

The bottleneck decides the fix, so the first action is a measurement, not a technique. The authors of the Megatron language model (LM) paper state the two training limits plainly: GPU memory is too small to fit large models, and the number of operations makes training times unrealistically long. Add the communication between GPUs and the memory traffic inside one GPU, and you have the four limits.

Four instruments tell the limits apart. A profiler trace splits each step into compute, communication and idle time. MFU, as defined in the PaLM paper, shows how busy the GPUs are. Achieved memory bandwidth against the hardware peak shows memory-bound kernels. Arithmetic intensity, the operations done per byte moved, compared with the roofline ridge point (see GPU fundamentals for ML engineers) says whether a kernel is compute-bound or bandwidth-bound.

An out-of-memory error means capacity is the limit. Low MFU with a trace full of waiting means communication. Slow attention with bandwidth near peak means memory traffic: FlashAttention's authors note (§1) that on modern GPUs compute has outpaced memory speed, so most transformer operations wait on memory. You now know which of the next four steps applies.

Step 2: Make the model fit with partitioning

When memory capacity is the limit, partition the model states instead of copying them. Data parallelism is a scheme that gives every GPU a full copy of the model and a different slice of the data. It does not shrink per-GPU memory: the ZeRO authors report that with plain data parallelism a model of more than 1.4 billion parameters runs out of memory on 32 GB GPUs (§1).

ZeRO is a family of methods that splits the model states across the data-parallel GPUs so each one stores only a share. With mixed-precision training and the Adam optimizer, the paper counts 16 bytes of model state per parameter. For an illustrative 70-billion-parameter model that is 1,120 GB, and splitting it across 64 GPUs leaves 17.5 GB each.¹

In its third stage, memory falls in proportion to the number of GPUs, at the cost of a 50% rise in communication (ZeRO §5). Activations, the values kept for the backward pass, are the other big memory use; recomputing them and splitting them along the sequence cut it (Korthikanti et al., 2022), and gradient accumulation keeps the effective batch while each micro-batch stays small enough to fit.

Text description of the diagram

• Plain data parallelism keeps 1,120 GB of model states on every GPU for a 70B model.
• ZeRO-3 across 64 GPUs leaves 17.5 GB of model states per GPU.

Step 3: Cut communication between GPUs

When GPUs wait on each other, reduce cross-node traffic. Tensor parallelism is a scheme that splits the matrices inside each layer across GPUs. Pipeline parallelism is a scheme that gives consecutive groups of layers to different GPUs. Narayanan et al. (2021) keep tensor parallelism inside a server, use pipeline parallelism across servers, and add data parallelism on top.

They reached 502 petaFLOP/s on 3,072 GPUs for a one-trillion-parameter model, 52% of theoretical peak. Per GPU that is 163 teraFLOP/s (Table 1), and 163 times 3,072 gives 500.7 petaFLOP/s, which matches the paper within rounding.¹ The same paper reports that its approach beats ZeRO-3 by 70% at 175 and 530 billion parameters because of less cross-node communication (§5). Overlapping communication with computation, sending one chunk while calculating the next, hides part of the cost that remains.

Step 4: Cut memory traffic inside a GPU

FlashAttention is an exact attention algorithm that uses tiling to cut reads and writes between the main GPU memory and on-chip static random-access memory (SRAM). Its authors report a 7.6 times speedup on the attention computation alone for Generative Pre-trained Transformer 2 (generative pre-trained transformer (GPT)-2) (Figure 1) and memory that grows linearly with sequence length (§3). The method does more arithmetic through recomputation, yet it runs faster because it moves less data. For an 8,192-token sequence the full attention matrix for one head holds 67 million 2-byte entries, 0.134 GB; at 65,536 tokens it is 8.59 GB, which is the traffic that tiling avoids.¹

Step 5: Raise the serving batch

For serving, throughput depends on batch size, and the KV cache caps the batch. The vLLM authors state that serving throughput requires batching many requests, but each request's KV cache is huge and changes size. Their PagedAttention divides the cache into fixed-size blocks, as an operating system pages memory, and they report 2 to 4 times higher throughput at the same latency.

The arithmetic, with an illustrative 70-billion-parameter model of 80 layers, 8 KV heads, head size 128 and 2-byte numbers, is 327,680 bytes of cache per token: 1.34 GB per 4,096-token request and 42.95 GB for 32.¹ With 140 GB of 16-bit weights that is 183 GB, so the model spans at least three 80 GB GPUs, and 2,000 GB per second is an illustrative aggregate bandwidth.¹

Prefill (processing the prompt) is compute-bound, while decode (generating each token) is bandwidth-bound, which is why batching helps decode. Reading the weights takes 70 milliseconds per step: one request gets 14.29 tokens per second, and 32 sharing each read reach about 350.¹ 8-bit weights take 70 GB and 4-bit weights 35 GB; check accuracy on your own tests.¹

1 · Find the bottleneck2 · Make it fit3 · Cut communication4 · Cut memory traffic5 · Raise the batchReplay

Step-by-step description

1. Measure which limit binds.
2. Partition the model states.
3. Split the model across GPUs wisely.
4. Tile attention inside the GPU.
5. Manage the KV cache in blocks.

• Find the bottleneck before changing anything.
• Partitioning model states makes the model fit.
• Raising the serving batch depends on the KV cache.

The Staff-level answer

What makes it Staff-level: it names the measurement first, ties each fix to one limit, and states the cost of each fix, the extra communication of ZeRO and the accuracy check after quantization.

The rubric

Three variants of this question

How would you train a model that does not fit on one GPU?

Partition the model states with ZeRO or split the model with tensor and pipeline parallelism; plain data parallelism copies the model and does not help. Pick by communication cost.

How would you reduce the cost of serving a language model?

Raise the batch size by managing the KV cache better, then store weights in fewer bits to free memory. The model-serving walkthrough for a 70B model applies this to a latency target.

What if training is slow but memory is fine?

Then compute or communication is the limit, so look at model FLOPs utilization and at attention kernels. Reduced-precision training, covered in 8-bit floating point (FP8) training, is one option.

Common mistakes

• "Data parallelism lets me train any size of model." (wrong) It copies the model on every GPU, so per-GPU memory does not fall. (right) "I partition the model states with ZeRO or split the model across GPUs."
• "More GPUs always means proportionally faster." (wrong) Naive use of parallelism leads to scaling issues at thousands of GPUs. (right) "I measure model FLOPs utilization and reduce cross-node communication as I scale."
• "FlashAttention approximates attention to save arithmetic." (wrong) It is exact and cuts memory traffic, not arithmetic. (right) "It tiles the computation so the large attention matrix never goes to main memory."
• "Inference is limited by compute." (wrong) Prefill is compute-bound, but decode, which generates each token, is bandwidth-bound, and the KV cache caps the batch. (right) "I check memory and bandwidth first, then batch."

Follow-up questions

What is the difference between data, tensor and pipeline parallelism?

Data parallelism copies the model and splits the data. Tensor parallelism splits the matrices inside each layer. Pipeline parallelism splits the layers into consecutive groups. Large runs compose all three.

Why is attention memory-bound, and what does FlashAttention change?

Standard attention writes a large sequence-by-sequence matrix to main memory, and reading it costs more time than the arithmetic. FlashAttention tiles the work to avoid that traffic. See high bandwidth memory for why that memory is the limit.

How does PagedAttention raise serving throughput?

It stores the KV cache in fixed-size blocks, so little memory is lost to fragmentation and duplication. A larger batch then fits on the same GPU, and throughput rises. See continuous batching, paged attention and LLM quantization.

What if the cluster grows from hundreds to thousands of GPUs?

Communication becomes the limit sooner. The ZeRO authors trained over 100 billion parameters on 400 GPUs, while the Megatron-LM work needed composed parallelism to scale to 3,072 GPUs, so expect to re-plan the split.

How would you cut serving cost without changing the model?

Use speculative decoding, better batching and caching of repeated work, then confirm output quality is unchanged.

Sources

1. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arXiv 1910.02054) — arXiv (Rajbhandari et al.; SC20), 4 Oct 2019
2. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (arXiv 2104.04473) — arXiv (Narayanan et al.; SC21), 9 Apr 2021
3. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv 2205.14135) — arXiv (Dao et al.; NeurIPS 2022), 27 May 2022
4. Efficient Memory Management for Large Language Model Serving with PagedAttention (arXiv 2309.06180) — arXiv (Kwon et al.; SOSP 2023), 12 Sep 2023
5. PaLM: Scaling Language Modeling with Pathways (arXiv 2204.02311) — arXiv (Chowdhery et al.), 5 Apr 2022
6. Reducing Activation Recomputation in Large Transformer Models (arXiv 2205.05198) — arXiv (Korthikanti et al.), 10 May 2022