11 ML Engineering Distributed Training Interview Questions and Answers

Short answer: Data parallelism — when the model fits on one GPU: a replica on every device, the batch is split, gradients are combined with all-reduce. Tensor parallelism splits individual matmuls across GPUs — needs NVLink-class links, intra-node. Pipeline parallelism splits layers into stages; micro-batches fight the bubble. The largest models combine all three — 3D parallelism.

In depth:

1. Order of choice — start with data parallelism; a layer doesn't fit → tensor parallelism within a node; the whole model doesn't fit → pipeline across nodes.
2. The bubble — stages sit idle at the start and end of each step; mitigated by slicing the batch into micro-batches.

⚠️ Common mistake: conflating tensor and pipeline parallelism: TP splits matrices inside a layer (all GPUs compute the same layer together), PP splits the model by layers into sequential stages.

Short answer: DDP runs one process per GPU, each holding a full model replica. Backward hooks gather gradients into buckets and launch asynchronous all-reduce overlapped with the ongoing backprop; after the sync every replica takes an identical optimizer step. DataParallel is a single process: the GIL, scatter/gather through GPU-0, which becomes the bottleneck.

In depth:

1. One process = one GPU — no GIL; launched via torchrun.
2. Gradient buckets — gradients are grouped (~25 MB by default); as soon as a bucket is full, its all-reduce starts without waiting for backward to finish.
3. Overlap — communication runs in parallel with computing gradients for earlier layers; the step pays almost nothing for synchronization.
4. Identical step — everyone has the same averaged gradients → replicas never diverge; a weight broadcast is needed only at startup.

backward:  layer N ─► bucket full ─► async all-reduce ───┐
           layer N−1, N−2 … gradients keep computing     │ overlap
optimizer.step() waits only for the last bucket ◄────────┘

⚠️ Common mistake: thinking gradients are uploaded to a central server. All-reduce is a peer-to-peer collective: every GPU ends up holding the same sum, there is no dedicated aggregator.

Short answer: All-reduce is a collective operation: every worker contributes its gradient tensor, and every worker receives the elementwise sum across all of them. Ring all-reduce moves data around a ring in chunks: each GPU transfers ≈2M bytes regardless of the number of workers — which is why it is bandwidth-optimal. NCCL implements it.

In depth:

1. Scatter-reduce — the tensor is cut into N chunks; over N−1 ring steps each GPU accumulates the complete sum of one chunk.
2. All-gather — another N−1 steps, and the finished chunks travel to everyone.
3. The arithmetic — 2(N−1) steps of M/N bytes: per-GPU traffic = 2(N−1)/N × M ≈ 2M, independent of N.
4. Versus a parameter server — the server's inbound link scales as O(N); the ring loads all links evenly.

GPU0 ─► GPU1 ─► GPU2 ─► GPU3 ─► GPU0    (chunks of M/N)
scatter-reduce: N−1 steps → each GPU owns the sum of 1 chunk
all-gather:     N−1 steps → everyone has the full sum
traffic/GPU: 2·(N−1)/N · M ≈ 2M — does not grow with GPU count

⚠️ Common mistake: "everything is collected on a master and broadcast back" — that is exactly the parameter server, whose link is the bottleneck; the ring exists precisely to avoid it.

Short answer: Weights + gradients (same size) + Adam states (m and v — two more copies) + activations + buffers. In fp32 with Adam that is 16 bytes per parameter — four times the weights themselves: a 1 GB model needs 3–5 GB before activations.

In depth:

fp32 + Adam, per parameter:
  weights     4 B
  gradients   4 B
  Adam m      4 B
  Adam v      4 B
  ─────────────────
  total      16 B/param (×4 the weights)

1 GB of weights (~250M params) → ~4 GB (3–5 GB) before activations

Mixed precision: 2 (fp16 weights) + 2 (grads) + 12 (fp32 master + m + v) ≈ 16 B
7B params → ≈112 GB — full fine-tuning does not fit on a single GPU

1. Activations — on top of that; they grow with batch size and sequence length, and are exactly what gradient checkpointing cuts.
2. Buffers — temporary CUDA tensors, allocator fragmentation, NCCL buffers.
3. Mixed precision doesn't rescue the optimizer — fp32 master weights and moments stay, so it is still ≈16 B per parameter; the win is in activations and speed.

⚠️ Common mistake: counting only the weights. The interviewer expects the ×4 multiplier for fp32+Adam and the conclusion: full fine-tuning of a 7B model requires multiple GPUs.

Short answer: The speedup comes from tensor cores plus half the memory and bandwidth for weights and activations. fp16 has only 5 exponent bits — a tiny range, small gradients underflow: you need dynamic loss scaling and fp32 master weights. bf16 has fp32's 8 exponent bits — the same range, no loss scaling; the mantissa is shorter, but training tolerates it.

In depth:

1. Where the speed comes from — tensor cores multiply half-precision matrices several times faster; half the bytes → less pressure on memory and interconnect.
2. fp16's pain — small gradients get zeroed (underflow); the fix is loss scaling: multiply the loss, divide the gradients — extra machinery and a source of NaNs.
3. Why bf16 won — fp32's range with no scaling; the lost mantissa is absorbed by SGD's stochasticity. The standard on A100/H100 and TPUs.

⚠️ Common mistake: "bf16 is more precise than fp16." The opposite: bf16 has a shorter mantissa and is coarser in precision — it wins on range, and range is what training stability needs.

Short answer: Accumulation simulates a big batch: K micro-batches accumulate gradients, one optimizer step — it saves activation memory at the cost of extra steps, but does NOT reduce optimizer-state memory. Checkpointing drops activations in the forward pass and recomputes them in backward — you pay ~30% extra compute for large activation-memory savings.

In depth:

1. Divide the loss by K — mandatory, otherwise the effective learning rate gets multiplied by K.
2. They compose — in practice you turn on both at once, plus mixed precision.
3. The classic CIS follow-up — accumulation breaks BatchNorm: statistics are computed over the small physical batch, not the effective one → SyncBatchNorm or LayerNorm.

⚠️ Common mistake: claiming accumulation saves optimizer memory. Adam moments exist one per parameter regardless of batch size — only ZeRO/FSDP shards those.

Short answer: Regular BatchNorm under DDP computes mean/var only over its own shard of the batch on each GPU. With a small per-GPU batch the statistics are noisy and differ across replicas — quality drops. SyncBatchNorm all-reduces the mean and variance across GPUs: statistics as if over the full global batch, at the cost of an extra sync per BN layer.

In depth:

1. What happens without it — DDP synchronizes gradients but not batch statistics: each replica has its own mean/var and its own running statistics.
2. When it hurts — detection/segmentation with a batch of 1–4 per GPU: normalizing over a couple of samples is nearly noise.
3. When you don't need it — with 32+ per GPU local statistics are fine; in transformers LayerNorm doesn't depend on the batch at all.
4. Enabling it — nn.SyncBatchNorm.convert_sync_batchnorm(model) before wrapping in DDP.

⚠️ Common mistake: assuming DDP "synchronizes everything." It only reduces gradients — BN statistics stay local and differ per replica, and that is exactly what interviewers probe.

Short answer: ZeRO distributes across data-parallel workers not the batch, but the training state. Stage 1 shards optimizer states, stage 2 adds gradients, stage 3 (FSDP full shard) adds the parameters themselves: a layer is assembled via all-gather just before its computation and freed right after. Memory drops from ~16P to ~16P/N bytes at the cost of extra communication.

In depth:

1. Just-in-time assembly — in stage 3 the full layer exists only for the duration of its computation; afterwards the shard is freed.
2. It is still data parallelism — each GPU pushes its own slice of the batch through the (assembled) full model; storage is divided, not computation.
3. The cost — the higher the stage, the more traffic: stage 3 moves parameters every step, so it wants a fast interconnect.

⚠️ Common mistake: calling FSDP "model parallelism." In tensor/pipeline parallelism GPUs compute different parts of the model; FSDP shards only storage — the computation stays data-parallel.

Short answer: Work through a list instead of guessing: dataloader starvation → communication-to-compute ratio → stragglers → too-small per-GPU batch. And reach for a profiler first (torch.profiler, nsys) — the trace shows where the gaps are.

In depth:

Check in order:
1. Dataloader: does GPU util look like a sawtooth in nvidia-smi?
   → num_workers, pin_memory, prefetch, data format/cache
2. Communication vs compute: large gradients, no NVLink,
   tiny buckets → bucket_cap_mb, sync less often
3. Stragglers: one slow GPU or node stalls the whole all-reduce
   → compare step times across ranks
4. Micro-batch too small: kernel-launch overhead eats the gains
   → bigger per-GPU batch, torch.compile / CUDA graphs

1. Measure first — the profiler trace shows what the GPU is waiting for: data, network, or a neighbor; answering "probably the network" without a trace loses points in an interview.
2. All-reduce can hide — with good overlap communication is buried under backward; if it sticks out in the trace, that is your missing 3x.

⚠️ Common mistake: blaming the network right away. In practice the dataloader starves most often — GPU-utilization gaps between steps are visible even in nvidia-smi.

Short answer: Periodic checkpoints with the full state: weights + optimizer state + LR scheduler + step counter + RNG states + dataloader position; plus a loop that can resume from them, atomic checkpoint writes, and keeping the last k checkpoints.

In depth:

ckpt = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),   # Adam moments!
    "scheduler": scheduler.state_dict(),
    "step": step, "epoch": epoch,
    "rng": {"torch": ..., "cuda": ..., "numpy": ...},
    "sampler": sampler.state_dict(),       # position in the data
}
torch.save(ckpt, "ckpt.tmp")
os.replace("ckpt.tmp", "ckpt.pt")          # atomic write

1. Atomicity — write to a temp file and rename: dying mid-torch.save must not corrupt the last good checkpoint.
2. Keep-last-k — not a single file (it can get corrupted), and not all of them (disk is finite).
3. Frequency — a tradeoff between lost hours and write cost; for large models writes are made asynchronous.

⚠️ Common mistake: checkpointing weights only. Without Adam moments the optimizer restarts cold on resume — the accumulated statistics are gone, and the loss spikes.

Short answer: Full fine-tuning — no: ≈16 bytes per parameter × 7B ≈ 112 GB. You can make it fit by cutting the terms one by one: LoRA (optimizer states only for ~0.1–1% of parameters), a quantized base (QLoRA: 4-bit weights + LoRA), gradient checkpointing, batch 1 with accumulation, bf16.

In depth:

Full FT: 7B × 16 B ≈ 112 GB  ≫ 24 GB → does not fit
Cut the terms:
  LoRA:  gradients and Adam only for adapters (~0.1–1%)
         frozen base: 7B × 2 B (bf16) = 14 GB + adapters
  QLoRA: base in 4-bit: 7B × 0.5 B ≈ 3.5 GB
         + LoRA + checkpointing → fits with room to spare

The menu, in order:

1. LoRA — removes the dominant term: the 12 B/param optimizer cost remains only for the tiny adapters.
2. Quantized base (QLoRA) — frozen weights in nf4: ×4 smaller than bf16.
3. Gradient checkpointing — cuts activations for ~30% extra compute.
4. Batch 1 + gradient accumulation — minimal activations while keeping the effective batch.
5. bf16 — for everything that isn't quantized.

⚠️ Common mistake: listing techniques without the arithmetic. A strong answer first shows 112 GB versus 24, then explains which term each trick removes.