TL;DR: Optimizer states, not weights alone, drive training memory: under mixed precision with Adam it works out to about 16 bytes per parameter (2 weights + 2 gradients + 4+4+4 FP32 master/momentum/variance). Data parallelism copies the model and divides the batch; tensor and pipeline parallelism divide the model itself; FSDP/ZeRO spreads the optimizer states, gradients, and weights across data-parallel ranks to lower per-GPU memory. Production training blends all three.
How to approach it
Open with the memory math, since it shows why sharding is needed at all, then present the parallelism types as answers to a single question: what gets split, the data or the model? Call out each one's communication cost, because that is the true tradeoff. The trap is treating data parallelism as a way to make a model fit when it only boosts 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.