Interview: Distributed Training

• Data parallelism: full model copied on every GPU, each processes a different data shard, gradients are averaged (all-reduce) - simple but requires the model to fit on one device.
• Tensor/model parallelism: individual layers/matrices are split across GPUs (e.g., splitting a weight matrix column-wise) - needed when a single layer doesn't fit on one device, requires high-bandwidth interconnect due to frequent communication.
• Pipeline parallelism: different layers are placed on different GPUs, and micro-batches flow through them like an assembly line - reduces communication vs tensor parallelism but introduces "bubble" idle time.

ZeRO eliminates the memory redundancy of standard data parallelism (where every GPU holds a full copy of optimizer states, gradients, and parameters) by sharding them across GPUs instead. Stage 1 shards optimizer states, Stage 2 additionally shards gradients, Stage 3 additionally shards the model parameters themselves (this is what FSDP implements) - each stage trades more communication for dramatically lower per-GPU memory, enabling much larger models to be trained on the same hardware.

Gradient accumulation runs several forward/backward passes on smaller micro-batches, summing their gradients before a single optimizer step - simulating a larger effective batch size than would fit in GPU memory at once. It's used when the desired batch size for stable training exceeds available memory.

Instead of storing all intermediate activations for the backward pass, only a subset of activations are kept and the rest are recomputed on-the-fly during backpropagation. This trades extra compute (roughly 20-30% slower) for a large reduction in memory usage, letting you train bigger models or larger batch sizes on the same hardware.

Distributed training requires frequent synchronization (all-reduce for gradients, all-gather for sharded parameters), and as GPU count grows this communication can dominate wall-clock time if bandwidth is low. High-bandwidth interconnects (NVLink within a node, InfiniBand across nodes) keep GPUs compute-bound rather than communication-bound, which is why cluster network topology is as important as raw GPU count for training efficiency.

FSDP (Fully Sharded Data Parallel) is PyTorch's native ZeRO-3 style sharding: parameters, gradients, and optimizer states are sharded across GPUs and all-gathered only when needed. DeepSpeed offers ZeRO stages plus pipeline/tensor parallelism and mature config for very large clusters. Choice often comes down to ecosystem fit (pure PyTorch vs DeepSpeed config) and team familiarity.

Tensor parallelism splits individual weight matrices across GPUs within a layer (high communication, low latency per step). Pipeline parallelism assigns different layers to different GPUs and streams micro-batches (lower communication, but pipeline bubbles). Large models often combine data + tensor + pipeline parallelism (3D parallelism).