Data, tensor and pipeline parallelism and collectives

Why it matters

When a model or its batch no longer fits on one GPU, or one GPU is too slow, the work must be split. The split fixes per-GPU memory, how much data crosses which links, and how long devices sit idle. Interviewers expect you to choose a layout from memory and communication arithmetic and to explain what goes wrong at scale: stragglers, bubbles and unhidden communication.

Core ideas

Collectives. All-reduce leaves every rank with the sum of all inputs; reduce-scatter leaves each rank with the sum of one shard; all-gather leaves every rank with all shards; all-reduce equals reduce-scatter followed by all-gather. In ring all-reduce over n ranks on S bytes, each rank sends and receives about 2(n−1)/n × S bytes, nearly independent of n, so it is bandwidth-optimal; but it takes 2(n−1) sequential steps, so latency grows with n. Tree and hierarchical algorithms cut latency-bound steps for small messages and many ranks; NCCL picks algorithms by message size and topology. Benchmarks report bus bandwidth, which rescales by 2(n−1)/n so the number is comparable to link bandwidth.

Data parallelism. Each rank holds a full replica, processes different data, and averages gradients with all-reduce before the optimizer step. PyTorch DDP groups gradients into buckets and starts a bucket's all-reduce as soon as its gradients are ready, overlapping communication with the rest of backward. Per-GPU memory does not shrink, and communication per step is proportional to model size, independent of batch size.

Sharded data parallelism (ZeRO, FSDP). Remove replicated model states: stage 1 shards optimizer state, stage 2 also gradients (reduce-scatter replaces all-reduce), stage 3 or full sharding as in FSDP also parameters, which are all-gathered layer by layer just before use in forward and backward and freed afterwards. Per-GPU model state drops roughly by the data-parallel degree. The ZeRO analysis shows stages 1 and 2 move about the same volume as plain data parallelism, while sharding parameters adds about half again because they are gathered in both passes.

Tensor parallelism. Split individual weight matrices. In a Megatron-style MLP the first linear layer is split by columns and the second by rows, so each block needs one all-reduce in forward and one in backward; attention is split by heads. These collectives sit inside every layer on the critical path, so tensor parallelism normally stays within a node on NVLink or NVSwitch.

Pipeline parallelism. Assign consecutive layers to stages on different devices and split the batch into m micro-batches. With p stages, a simple schedule leaves each device idle for about p − 1 micro-batch slots out of m + p − 1, so the bubble shrinks as m grows relative to p. One-forward-one-backward schedules keep that bubble but bound in-flight activations to about p micro-batches; interleaved schedules shrink the bubble further at the cost of more communication. Stages exchange only boundary activations, so pipelines can span nodes, but stages must be balanced.

Topology and layout. Bandwidth inside a node is much higher than across nodes. A common layout puts tensor parallelism inside a node, pipeline stages across nodes, and data parallelism or FSDP outermost. Sequence/context parallelism and expert parallelism add dimensions for long contexts and mixture-of-experts models.

Stragglers. Synchronous training waits for the slowest rank at every collective. A slow data loader, thermal throttling, a degraded link or an unbalanced stage stretches every step for everyone. Compare per-rank step times and time spent waiting inside collectives before tuning anything else.

Common misconceptions

• "Data parallelism reduces per-GPU memory." Plain DDP replicates full model states; only sharding or model parallelism reduces them.
• "Ring all-reduce sends more data per GPU as GPUs are added." Per-GPU volume stays near 2S; the step count and exposure to the slowest link grow.
• "FSDP is model parallelism." It is data parallelism with sharded storage; each rank still computes whole layers on its own data after gathering parameters.
• "Pipeline parallelism is free with enough layers." Bubbles, stage imbalance and stored activations cost throughput and memory, and it needs many micro-batches.
• "Tensor parallelism scales across nodes like data parallelism." Its per-layer collectives on the critical path make slow inter-node links the bottleneck.

How to approach it

Answer layout questions in a fixed order: memory per GPU for each option, communication volume per step and on which links, idle time (bubbles, waiting), then operational concerns such as checkpointing and failures. At graduate level, explain DDP and the collectives precisely. At mid level, compute ZeRO memory and the pipeline bubble. At senior level, design a layout for a given cluster and say how you would verify it with collective benchmarks and traces.