Start with the number that makes the whole topic real. A 70-billion-parameter model trained with Adam in mixed precision needs roughly 16 bytes per parameter just to hold weights, gradients, and optimizer state. That is about 1.1 terabytes before a single activation exists. An H100 has 80 gigabytes. So the first question in any serious AI-infrastructure interview is pure arithmetic: the model does not fit, and you have to say what you do about it.
NVIDIA, Together AI, CoreWeave, and the frontier labs all ask some version of this in the systems or ML-infra round. The prompt is usually concrete. “You have a 70B model and nodes of eight H100s connected by NVLink, with InfiniBand between nodes. Walk me through how you would train it.” What they listen for is whether you understand where memory goes, where bytes travel, and which of those two is your bottleneck.
The memory math they expect you to do out loud
Break the 16 bytes down, because the follow-up is always where that comes from. With bf16 compute you keep a 2-byte weight and a 2-byte gradient per parameter. Adam adds an fp32 master copy of the weights at 4 bytes, plus first and second moment estimates at 4 bytes each. That is 2 + 2 + 4 + 4 + 4 = 16. For a 7B model it works out to about 112 GB of state; for 70B it is roughly 1.1 TB. This is the calculation the ZeRO work out of Microsoft made standard, and interviewers expect you to reproduce it without notes.
Then there are activations, which people forget and which often dominate at long sequence lengths. Activation memory scales with batch size, sequence length, and layer count, and it is the piece you attack with recomputation. Activation checkpointing stores only the layer boundaries and recomputes the rest during the backward pass, buying a large memory reduction for roughly a third more compute. When a candidate brings up checkpointing before I ask, that usually means they have run out of memory on a real job.
Why more GPUs alone does not save you
Data parallelism is the default and the first thing people reach for. Replicate the model on every GPU, give each a different slice of the batch, then average gradients with an all-reduce before the optimizer step. Ring all-reduce moves about 2(N-1)/N times the gradient size per GPU, which stays roughly flat as you add workers, so throughput scales cleanly. The catch is the one the interviewer is fishing for: every replica still holds the full model. If 1.1 TB will not fit on one GPU, a thousand copies of it will not fit either. Data parallelism scales your batch, not your model.
This is where ZeRO and its PyTorch-native cousin FSDP come in. Rather than replicate the optimizer state, gradients, and parameters, you shard them across the data-parallel group. Stage 1 shards optimizer state, stage 2 adds gradients, and stage 3 (equivalent to FSDP’s full shard) adds the parameters themselves, so each GPU only materializes a layer’s weights when it is about to compute that layer, then frees them again. You trade extra communication, gathering the shards on the fly, for a large drop in per-GPU memory. On a fast interconnect that trade is usually worth taking, and it is how many teams train models that would otherwise need exotic hardware.
Splitting the model itself: tensor and pipeline
When a single layer’s matmul is too big for one card, or you want to cut activation memory further, you split the model across devices. There are two ways, and knowing when to use each is most of the interview.
Tensor parallelism splits individual operations. Megatron-LM partitions the attention heads and the MLP weight matrices across GPUs, so each device computes part of every layer. The price is an all-reduce inside every transformer block on both the forward and the backward pass, which is a lot of traffic. That traffic is why you keep tensor parallelism inside a node, where NVLink and NVSwitch give you hundreds of gigabytes per second between GPUs. Push it across the slower InfiniBand fabric between nodes and utilization falls off a cliff. In practice the tensor-parallel degree matches the GPUs per node, often eight.
Pipeline parallelism splits by depth. Assign the first block of layers to one group of GPUs, the next block to the next group, and pass activations forward and gradients back between stages. The problem is the bubble: while stage one works on the first batch, the later stages sit idle. You hide it by chopping the batch into micro-batches so every stage stays busy, and the idle fraction drops to about (p-1)/(m+p-1) for p stages and m micro-batches. The 1F1B schedule, one forward then one backward interleaved, shrinks both the bubble and the peak activation memory compared with the naive GPipe schedule, and being able to explain that difference is a strong signal.
Training runs at frontier scale combine all three. Tensor parallelism inside each node, pipeline parallelism across a handful of nodes, and data parallelism or ZeRO across the resulting replicas. This is the 3D parallelism that Megatron-LM and DeepSpeed made common, and drawing those three axes on the whiteboard is usually where the question was heading.
| Strategy | What it splits across GPUs | Main communication per step | Where it belongs | Main limitation |
|---|---|---|---|---|
| Data parallelism | The input batch; the full model is replicated on every GPU | All-reduce of gradients once per step | Any cluster with decent bandwidth | The whole model must fit on a single GPU |
| ZeRO / FSDP | Optimizer state, gradients, and (stage 3 / full shard) the parameters | All-gather parameters, reduce-scatter gradients | Across the data-parallel group on fast links | Extra gather and scatter traffic; sensitive to interconnect speed |
| Tensor parallelism | Individual weight matrices within each layer | All-reduce twice per transformer block | Inside one node over NVLink or NVSwitch | Too chatty to cross nodes; degree capped near GPUs-per-node |
| Pipeline parallelism | Contiguous groups of layers, assigned as stages | Point-to-point activations between adjacent stages | Across nodes over InfiniBand | Pipeline bubble idles GPUs unless micro-batched well |
The follow-ups that separate people
Once the layout is on the board, the questions get pointier. Mixed precision tends to come up first. bf16 is the default for stability because its wider exponent range avoids the loss-scaling gymnastics fp16 needs. On Hopper and Blackwell hardware, fp8 training through the Transformer Engine has moved from research into production for the large pretraining runs, and a candidate who can say which tensors stay in higher precision, the master weights and the reductions, is ahead of the pack.
Then throughput. Gradient accumulation lets you simulate a large batch when memory caps the micro-batch, at the cost of more forward and backward passes between optimizer updates. Overlapping communication with computation, kicking off one layer’s gradient all-reduce while the next layer still computes, is often the difference between 40 and 55 percent model-FLOPs-utilization, and MFU is the number infra teams actually track. Quote a realistic MFU, low-to-mid 40s is a fair figure for many large runs, and you sound like someone who has watched the dashboards rather than read the paper.
The failure questions are where senior candidates pull ahead. A run across 512 GPUs will lose a GPU. What happens? The next all-reduce hangs, the whole job stalls, and without a plan you have lost everything since the last checkpoint. So you checkpoint on a cadence tuned to your mean time between failures, you shard the checkpoint write so it does not freeze every rank for minutes, and you keep a hot spare or an elastic scheduler that can restart from the last save. Stragglers are the quieter version of the same problem: one slow node drags every synchronous step down to its pace, so teams watch per-rank step times and evict the laggard. None of this shows up in a tutorial, which is exactly why it gets asked.
What a good answer sounds like
The candidates who do well do not recite four strategies in order. They start from the constraint, this model needs 1.1 TB, my GPU holds 80 GB, my fast link lives inside the node, and let that drive the layout. Memory too large for one card points to sharding or model splitting; a matmul too large for one card points to tensor parallelism; too many layers points to pipeline; then data parallelism or ZeRO on top to spend the rest of the cluster. Name the bottleneck out loud at each step and the interviewer can follow your reasoning even when they would have chosen differently.
A few prompts, close to how they land in the room:
- “What is sitting in GPU memory during a training step, and how big is each part for a 7B model?”
- “You have data parallelism working and the model still will not fit. Now what?”
- “Where do you draw the tensor-parallel boundary, and why not run it across nodes?”
- “Loss was smooth, then spiked at step 40k and never recovered. What do you check first?”
- “One rank out of 512 dies four hours into the run. What does the job do, and how do you get back?”
The loss-spike question rewards rehearsal because it reads your instincts in both directions. A weak answer jumps straight to lowering the learning rate. A strong one walks the suspects: a corrupted data shard, an fp8 or fp16 overflow, a gradient norm that blew past the clip threshold, a last-good checkpoint to roll back to before resuming with the bad batches skipped. It is the memory math again, applied under pressure.
If you have never run a multi-node job, the tell is that you will describe the parallelism strategies correctly and miss the operational half completely. So before the interview, pull a real training config from Megatron-LM or an open reproduction, find the lines that set the tensor-parallel and pipeline-parallel degree, and work out why they picked those numbers for that cluster. That one exercise teaches more than any diagram.
