Hitchhiker's Guide — Distributed Training & Parallelism

The compressed practitioner tour. If WARMUP.md is the professor, this is the senior who leans over and says "here's what you actually need to remember when the cluster is on fire."

The 30-second mental model

One GPU can't hold a frontier model (~16 B/param → 70B ≈ 1.1 TB). So you split it four ways: data parallel (replicate, all-reduce gradients), tensor parallel (shard a layer's matmul), pipeline parallel (shard the stack of layers), and ZeRO/FSDP (shard the optimizer/grad/param state instead of replicating it). The glue is the ring all-reduce — 2(N-1) steps, ~2·tensor bytes/rank regardless of cluster size. The moment you split anything, the bottleneck becomes the network, not the FLOPs. Real runs compose all four (3D parallelism) and the art is mapping each axis to the right interconnect tier.

The numbers to tattoo on your arm

ThingNumber
Ring all-reduce steps2(N-1)
Ring bytes moved / rank2·tensor·(N-1)/N → ~2·tensor as N grows
All-reduce =reduce-scatter (N-1) + all-gather (N-1)
DP gradientall-reduce sum ÷ N = serial mean = one big batch
TP column-parallelconcat, no comm
TP row-parallelall-reduce partial sums (needs NVLink)
TP placement ruleTP ≤ GPUs per node (8)
Pipeline bubble(P-1)/(M+P-1) → 0 as M→∞
Train memory (Adam, fp16)16N = params 2N + grads 2N + optimizer 12N
ZeRO per-GPUst0 16N · st1 4N+12N/N · st2 2N+14N/N · st3 16N/N
NVLink vs InfiniBand~900 GB/s intra-node vs ~50–400 GB/s inter-node
Activation checkpointing~33% more compute for big activation-memory cut

Back-of-envelope one-liners

# "Will a 70B fit on one 80GB GPU to TRAIN?"
70e9 * 16 = 1.12 TB  → no, not even close;  need ZeRO-3 across many GPUs

# "ZeRO-3 for 70B on 16 GPUs — per-GPU state?"
70e9 * 16 / 16 = 70 GB/GPU  → just fits 80GB (before activations → add checkpointing)

# "Ring all-reduce a 14GB (7B fp16) gradient on 8 GPUs — bytes/rank?"
2 * 14e9 * 7/8 ≈ 24.5 GB moved per rank, constant-ish no matter the cluster size

# "Pipeline P=8 stages — bubble at M=8 vs M=64?"
(8-1)/(8+8-1) = 7/15 ≈ 47%   vs   7/71 ≈ 10%   → more microbatches, not more stages

# "Map 175B onto 8 nodes × 8 GPUs"
TP=8 in-node (NVLink), PP=8 across nodes (IB), DP/ZeRO outermost

The framework one-liners (where these numbers live in real tools)

# Data parallel
model = torch.nn.parallel.DistributedDataParallel(model)        # bucketed, overlapped all-reduce
# Fully sharded (= ZeRO-3)
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)
# DeepSpeed ZeRO
#   ds_config = {"zero_optimization": {"stage": 3}, "gradient_accumulation_steps": M}
# Megatron tensor + pipeline
#   --tensor-model-parallel-size 8  --pipeline-model-parallel-size 8
# The collective itself
torch.distributed.all_reduce(t, op=ReduceOp.SUM)                 # NCCL ring/tree under the hood
# Measure it
#   nccl-tests all_reduce_perf ;  torch.profiler ;  NCCL_DEBUG=INFO ;  NCCL_ALGO=Ring

War stories

  • The 64-GPU slowdown. A DP job flew at 8 GPUs, crawled at 64. The all-reduce had stopped overlapping with backward and the team had TP crossing the node boundary onto InfiniBand. Fix: keep TP ≤ 8 (in NVLink), enable gradient bucketing/overlap. 3× faster, zero new hardware.
  • The optimizer-state OOM. "The 13B model is only 26 GB in fp16, why won't it train on an 80 GB card?" Because training is 16N ≈ 208 GB, not 2N. They were quoting inference weights. Switched to FSDP FULL_SHARD and it fit. Always size training at 16 B/param.
  • The half-idle pipeline. A P=8 pipeline ran at ~50% util because M=8 → 47% bubble. Bumped micro-batches to 64 (with 1F1B so activation memory didn't blow up) → bubble ~10%, util ~90%.
  • The silent divergence. Someone "optimized" the gradient sync and accidentally skipped a rank on uneven batches; the replicas drifted apart and loss got mysteriously worse. The all-reduce ÷ N must include every rank exactly once — that's the correctness contract of DP.

Vocabulary (rapid-fire)

  • All-reduce — element-wise sum across ranks, result on every rank; = reduce-scatter + all-gather.
  • Ring all-reduce — bandwidth-optimal all-reduce, 2(N-1) steps, ~2·tensor bytes/rank.
  • DP / TP / PP — data / tensor / pipeline parallel (split batch / a layer / the stack).
  • ZeRO / FSDP — shard optimizer→grad→param state across DP ranks (stages 1/2/3).
  • Bubble — pipeline idle from fill+drain, (P-1)/(M+P-1).
  • 1F1B — one-forward-one-backward pipeline schedule; GPipe's bubble at ~P activation memory.
  • 3D parallelism — DP × TP × PP composed and mapped to the topology.
  • Activation checkpointing — recompute activations in backward to save memory.
  • NVLink / InfiniBand — fast intra-node / slower inter-node interconnect.
  • MFU — model FLOPs utilization (Phase 00); distributed runs live or die on it.

Beginner mistakes

  • Sizing training at 2N (inference weights) instead of 16N (params+grads+optimizer).
  • Thinking "add GPUs → linearly faster" — communication eventually dominates.
  • Letting tensor parallelism cross the node boundary onto InfiniBand (kills the in-layer all-reduce).
  • Fixing a pipeline bubble with more stages (makes it worse) instead of more micro-batches.
  • Forgetting the gradient all-reduce must include every rank exactly once (÷ N), or replicas diverge.
  • Confusing ZeRO with a separate parallelism axis — it's DP with the redundant state sharded away.
  • Running with no comm/compute overlap and blaming the GPU.
  • Ignoring activation memory (the thing checkpointing exists for) when sizing a run.

The one thing to take away

Before you launch a distributed run, write the plan: what you split (DP/TP/PP/ZeRO), the per-GPU memory it leaves (16N shards), the collective cost (2(N-1) ring), and the mapping to the topology (TP in node, PP across, DP outermost). If you can't, you're hoping; if you can, you're engineering a cluster.