Distributed training (DDP, ZeRO, FSDP, TP, PP)

Why this matters

Above ~7B parameters you cannot fit a model on a single GPU. Above ~70B you cannot even fit it on a single node. Frontier models train on thousands of GPUs in coordinated parallelism. Knowing the names — DDP, ZeRO, FSDP, TP (Tensor Parallel), PP (Pipeline Parallel), SP (Sequence Parallel), EP (Expert Parallel) — and which to combine, is the entire expertise of a "large-model training engineer." It is among the highest-paid specialisations in AI ($350k–$700k+ in 2026).

You will not run a 1000-GPU job. But you will discuss it in interviews and will use a 2-8 GPU subset of these techniques.

Learning objectives

1. Distinguish data, tensor, pipeline, and expert parallelism.
2. Use PyTorch DDP and FSDP with confidence.
3. Read a torchrun command and explain every flag.
4. Estimate which parallelism strategy you need for a given model + GPU count.
5. Recognise common bottlenecks (NCCL, slow networks, gradient checkpointing).

1. The 5 axes of parallelism

1.1 Data parallelism (DP)

Replicate the model on every GPU; shard the batch. After backward, all-reduce gradients across GPUs.

• Simple; near-linear scaling on 8-64 GPUs.
• Memory: each GPU has the full model, full optimiser state.
• Bottleneck for big models.

PyTorch DDP is the default:

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = MyModel().to(rank)
model = DDP(model, device_ids=[rank])

1.2 Tensor / Model parallelism (TP)

Split individual matrices across GPUs. The Q projection of size (d, d) becomes (d, d/N) per GPU; the output is gathered.

• Used for matrices too large for a single GPU.
• Heavy intra-layer all-reduce; needs fast NVLink (i.e., within a single node).
• Megatron-LM is the canonical implementation.

1.3 Pipeline parallelism (PP)

Split layers across GPUs. GPU0 has layers 1-20, GPU1 has 21-40, etc. Mini-batches flow through the pipeline.

• "Bubble" overhead at start/end of each step (GPUs idle while pipeline fills).
• Mitigated by micro-batches (1F1B schedule, GPipe, PipeDream).
• Used across nodes (slower interconnect is OK because layer outputs are tiny vs gradient all-reduces).

1.4 Sequence parallelism (SP)

Split along the sequence dimension. Useful when sequences are very long. Needed in tandem with TP for very long-context training.

1.5 Expert parallelism (EP)

For MoE: place different experts on different GPUs. Token routing → all-to-all communication.

Combining axes (3D / 4D parallelism)

A real frontier training run uses 3-4 axes simultaneously:

total GPUs = DP × TP × PP × EP
e.g. 4096 = 128 × 8 × 4 × 1

Megatron-LM, DeepSpeed, and torchtitan provide configurations for this.

2. The two main "easy mode" strategies

For 95% of jobs (≤ 1024 GPUs, dense models up to 70B), choose between ZeRO-3 / FSDP and TP + PP. Both work; FSDP is simpler.

ZeRO (DeepSpeed) — Stages 1, 2, 3

ZeRO shards memory across DP ranks instead of replicating:

• Stage 1: shard optimiser states across DP ranks.
• Stage 2: + shard gradients.
• Stage 3: + shard parameters (gather on demand for forward/backward).

Memory savings:

• Stage 1: ~4× (Adam state is ~12 bytes/param vs 4 for params).
• Stage 3: ~N× (where N is the number of DP ranks).

FSDP (PyTorch native)

FullyShardedDataParallel is PyTorch's clean reimplementation of ZeRO-3, integrated with torch.distributed.

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, MixedPrecision
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

mp = MixedPrecision(param_dtype=torch.bfloat16,
                    reduce_dtype=torch.bfloat16,
                    buffer_dtype=torch.bfloat16)

wrap_policy = functools.partial(transformer_auto_wrap_policy,
                                transformer_layer_cls={Block})

model = FSDP(MyTransformer().cuda(),
             auto_wrap_policy=wrap_policy,
             mixed_precision=mp,
             device_id=rank)

Key practices:

• Wrap each transformer block as its own FSDP unit (so memory cycling matches the natural granularity).
• Use MixedPrecision to keep params/gradients in bf16; full state-dict for checkpointing.
• Combine with activation_checkpointing to trade compute for memory.

torchtune and torchtitan give you FSDP-2 + TP recipes ready to use.

3. Activation checkpointing (a.k.a. gradient checkpointing)

Activations dominate memory in long sequences. Activation checkpointing recomputes them during backward instead of storing them — ~25% extra compute, but unlocks large batch sizes.

from torch.utils.checkpoint import checkpoint
def block_forward(x): return block(x)
x = checkpoint(block_forward, x, use_reentrant=False)

Almost every modern recipe enables it.

4. NCCL — the communication backbone

NVIDIA Collective Communications Library handles all_reduce, all_gather, reduce_scatter across GPUs. Critical performance knobs:

• NVLink (intra-node): ~600-900 GB/s on H100/B200.
• InfiniBand (inter-node): 200-400 Gb/s — orders of magnitude slower than NVLink.
• Strategy implication: keep TP within a node, DP/PP across nodes.

Common environment vars:

NCCL_DEBUG=INFO
NCCL_IB_DISABLE=0
NCCL_SOCKET_IFNAME=eth0

When training stalls, NCCL is usually the suspect. Check nvidia-smi, network throughput, and whether all ranks reached the same step.

5. torchrun — the launcher

torchrun \
  --nnodes=4 \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:29500 \
  --rdzv_id=run42 \
  train.py --config llama3_8b.yaml

What each flag does:

• nnodes / nproc_per_node: total ranks = nnodes × nproc_per_node.
• rdzv_*: rendezvous backend; c10d is built into PyTorch.
• rdzv_endpoint: the master node's address.

Inside train.py:

import os, torch.distributed as dist
dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

6. Picking a strategy (quick decision tree)

Model fits on one GPU (<= ~6B params, full bf16)?
  └─ DDP. Done.

Model fits on one node (≤ 8 GPUs, ~70B with FSDP)?
  └─ FSDP with bf16 + activation checkpointing. Done.

Model spans multiple nodes (200B+)?
  └─ FSDP across nodes for DP, +TP within node, possibly +PP across nodes.

Mixture-of-Experts model?
  └─ FSDP + Expert Parallel; route experts to different GPUs.

The frontier compute mix in 2026 (e.g., DeepSeek-V3, Llama 4):

• DP via FSDP (or ZeRO-3).
• TP=8 within a node.
• PP across racks.
• EP for MoE layers.
• bf16 weights, fp8 matmuls (Hopper / Blackwell), activation checkpointing.

7. Fault tolerance

Real training runs fail — GPU dies, network blip, host kernel panic. You need:

• Periodic checkpointing (e.g., every 1k steps).
• Async checkpointing (write to NVMe / GCS in background).
• Auto-restart from latest checkpoint.
• Skip-on-NaN: throw away the spike step and rewind a few steps.

Tools: torch.distributed.checkpoint (DCP), torchsnapshot, internal frameworks at every lab.

Hands-on lab (4 hours, 2-GPU machine or Colab Pro)

distributed_lab.ipynb + train_ddp.py:

1. Take your nano-GPT from Lesson 2.4. Wrap it in DDP. Verify same loss curve as single-GPU. Use torchrun --nproc_per_node=2.
2. Replace DDP with FSDP. Wrap each Block. Print memory per rank — should drop ~2×.
3. Add bf16 mixed precision via FSDP MixedPrecision.
4. Add activation checkpointing on each block. Increase block_size (sequence length) until memory is again the limit; show how much further you got.
5. Save and reload an FSDP checkpoint correctly (use state_dict_type=FullStateDictConfig).
6. Bonus: pretend torchrun --nnodes=2 with two physical machines on a LAN; run the rendezvous and verify.

Common pitfalls

1. Module ordering — model = MyModel(); model.cuda(); model = DDP(model) is correct; the reverse breaks.
2. Per-rank random seed — set torch.manual_seed(seed + rank) for the data loader, NOT for the model init (model init must be identical across ranks).
3. Forgetting dist.barrier() before checkpointing → ranks save inconsistent states.
4. Calling .zero_grad() before all ranks finished backward — DDP will hang.
5. Mismatched FSDP wrap policy — wrapping the entire model as one unit defeats sharding benefits.
6. Print from all ranks — flooded logs. Use if rank == 0: print(...).

Self-check

1. What does ZeRO-3 shard?
2. Why is TP usually intra-node and DP inter-node?
3. What is the bubble in pipeline parallelism?
4. When would you turn on activation checkpointing?
5. What is torchrun --nproc_per_node?

References

• Rajbhandari et al. (2020), "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models."
• Korthikanti et al. (2022), "Reducing Activation Recomputation in Large Transformer Models."
• Smith et al. (2022), "Megatron-Turing NLG 530B."
• Shoeybi et al. (2019), "Megatron-LM."
• PyTorch FSDP getting started.
• torchtitan — Meta's reference distributed-training framework (FSDP-2 + TP + PP).
• DeepSpeed docs.