TL;DR: Training memory is dominated by optimizer states, not just weights: with mixed precision and Adam it is roughly 16 bytes per parameter (2 weights + 2 gradients + 4+4+4 FP32 master/momentum/variance). Data parallelism replicates the model and splits the batch; tensor and pipeline parallelism split the model itself; FSDP/ZeRO shards the optimizer states, gradients, and weights across data-parallel ranks to cut per-GPU memory. Real training combines all three.
How to approach it
Start with the memory math, because it explains why you shard at all, then lay out the parallelism types as answers to one question: what do you split, the data or the model? Name the communication cost of each, since that is the real tradeoff. The trap is treating data parallelism as a fit solution when it only helps throughput.
A strong answer
Memory math (mixed precision + Adam). Per parameter you store: BF16 weights (2 bytes) + BF16 gradients (2) + FP32 master weights (4) + FP32 momentum (4) + FP32 variance (4) = 16 bytes/param. So a 7B model needs about 112GB just for model state, before activations, which already exceeds an 80GB GPU. A 70B model is about 1.1TB. That arithmetic is why large-model training must shard.
Parallelism taxonomy:
- Data parallelism (DDP): replicate the full model on each GPU, split the batch, all-reduce gradients each step. Simple and communication-light per step, but every GPU holds the entire 16-bytes/param state, so it does nothing when the model itself does not fit.
- Tensor parallelism: split individual layer matrices (column/row-wise) across GPUs; each does part of the matmul and they all-reduce within the layer. Heavy, fine-grained communication, so keep it inside a fast NVLink domain (a single node).
- Pipeline parallelism: put different layers on different GPUs and stream micro-batches through, like an assembly line. Communication is only at stage boundaries, but you fight pipeline "bubbles" (idle time); micro-batching shrinks them.
- FSDP / ZeRO: the key memory lever. Instead of replicating the 16-bytes/param state on every data-parallel rank, shard it across them. ZeRO-1 shards optimizer states, ZeRO-2 adds gradients, ZeRO-3 (FSDP) shards parameters too, gathering each layer's weights just-in-time for its forward/backward then releasing them. Per-GPU model-state memory drops roughly to 16N/G bytes across G ranks, at the cost of extra all-gather and reduce-scatter traffic.
Real frontier training is 3D parallelism, tuned to the interconnect:
Key takeaways
- Training state is about 16 bytes/param with mixed precision + Adam; counting only weights underestimates memory roughly 8x.
- Data parallelism helps throughput, not fit; model sharding (tensor, pipeline, FSDP/ZeRO) is what makes a too-big model trainable.
- Tensor parallelism is chatty: keep it intra-node; use pipeline parallelism to span nodes.
- FSDP/ZeRO-3 cuts per-GPU model state to ~16N/G bytes by trading memory for all-gather/reduce-scatter communication.
What interviewers probe next
- "Activation memory?" Grows with batch × sequence × layers and can rival weights; cut it with activation/gradient checkpointing (recompute in backward) and sequence/context parallelism.
- "NCCL collectives?" All-reduce (DDP gradients), all-gather plus reduce-scatter (FSDP). Profile them; interconnect bandwidth (NVLink vs PCIe vs InfiniBand) often bounds scaling.
- "When tensor vs pipeline parallelism?" Tensor for intra-node (needs high bandwidth), pipeline to span nodes (tolerates slower links); combine them.
- "Inference memory vs training?" Inference needs only weights plus KV cache (no gradients or optimizer states), which is why a model that needs 16 bytes/param to train serves in about 2 bytes/param, or less when quantized.
Common mistakes
- Counting only weights and forgetting gradients and FP32 optimizer states, underestimating memory roughly 8x.
- Reaching for tensor parallelism across nodes, where its chatty communication tanks throughput.
- Confusing data parallelism (helps throughput, not fit) with model sharding (helps fit).
- Ignoring activation memory, then OOM-ing despite the weights fitting.