Skip to content

The distributed training handbook

An inference role still has to understand the parallelism on the training side: interviews ask about ZeRO and Megatron, reinforcement-learning training ties an inference engine to a training framework, and the tensor, expert and context parallelism used in inference all came from training. This handbook starts from what will not fit on one card, implements data parallelism, ZeRO, tensor parallelism, pipeline parallelism, context parallelism and expert parallelism from scratch one by one, and lines each of them up against a single process on the CPU using multiple processes.

16 chapters14 exercisesMulti-process (gloo) on CPU, matched against a single process
Start readingRoadmapExercises

Who this handbook is for

  • You have never trained a language model by hand: read Train a Small Model first, walk the whole pipeline through on your own computer, and then come back for multiple GPUs.
  • You know the Transformer's structure and can write a training loop in PyTorch, but have never done multi-GPU training.
  • Or you have used DeepSpeed, Megatron-LM or FSDP, but cannot say what each step communicates or where the memory is saved.
  • You work on inference systems and need to understand the training side's parallelism (reinforcement-learning training, weight synchronisation, the tensor-parallel and expert-parallel implementations shared with inference).

After reading and working through this handbook you should be able to:

  • Compute the memory budget per card (parameters, gradients, optimizer states, activations) and the communication volume per step for any model and cluster configuration.
  • Say clearly what data parallelism, ZeRO-1/2/3, tensor parallelism, sequence parallelism, pipeline parallelism, context parallelism and expert parallelism each partition, what they communicate and what they cost.
  • Write DDP from scratch (bucketing with overlapped communication), ZeRO-1, a tensor-parallel MLP including the backward pass, Ulysses attention and an expert-parallel mixture-of-experts layer, and verify each against the single-card result.
  • Choose a sensible combination of parallelism for a specific model and cluster, and explain why.

The route through

Stage Chapters What you can do afterwards Suggested time
1. Basics Overview: the memory budget and the time model · Collective primitives work out the memory and the communication, and know where the bottleneck is 2 to 3 days
2. Data parallelism DDP · ZeRO and FSDP write DDP with bucketed overlap, and ZeRO-1 2 to 3 days
3. Model parallelism Tensor and sequence parallelism · Pipeline parallelism · Context parallelism · Mixture of experts and expert parallelism explain each kind's partitioning and communication, and write a minimal implementation of it 1 week
4. Precision and strategy Mixed precision and FP8 · Combining 3D and 5D parallelism · Training frameworks and reinforcement-learning systems pick a configuration for a specific case, and read Megatron, DeepSpeed and verl 3 to 4 days
5. Optimizers, stability and RL algorithms AdamW, Muon and distributed optimizers · Training stability · Reinforcement-learning algorithms in depth explain the newer optimizers, the root causes of a loss spike and what to do about them, and the reinforcement-learning algorithms after GRPO 3 to 4 days

The chapter-by-chapter route through every handbook is in the roadmap, and the week-by-week schedule is in the plan; this book is week 9 of the plan, studied alongside distributed inference in the inference-systems handbook.

How it was verified

  • Every script with a file name on it actually runs on the CPU build of PyTorch, and the output on the page matches the run line by line.
  • The multi-process scripts are launched with torchrun --nproc-per-node N (this book's checker passes an explicit --master-addr 127.0.0.1 and a free port rather than --standalone, which hangs on some Macs while reverse-resolving the hostname) and the communication backend is gloo (CPU); moving to a GPU only takes changing the backend to nccl and putting the tensors on cuda.
  • Every kind of parallelism is lined up item by item against the single-process computation: the forward output, the loss, and the gradients.
# the CPU build of PyTorch is required (2.4 or newer)
pip install torch --index-url https://download.pytorch.org/whl/cpu
torchrun --standalone --nproc-per-node 4 collectives.py
# if --standalone hangs (common on a Mac), give it an explicit address and port:
# torchrun --nproc-per-node 4 --master-addr 127.0.0.1 --master-port 29500 collectives.py

Found a mistake or a better way to say it? Edit this page on GitHubReport an issue