Back to Learning
AI Next Labs / Learning
Data parallelism and ZeRO
Copy the model, split the batch: ring all-reduce, overlapping communication, the batch-size recipe and ZeRO stages 1 to 3.
Published
Written by Jaydeep Raijada.
The simplest way to use more GPUs: give every GPU a full copy of the model and a different slice of the data, then make them agree on one gradient. This lesson covers how that agreement is reached cheaply, how to size the batch, and how ZeRO removes the memory wasted by keeping N identical copies.
Part 1 of 8 in the 5D Parallelism series. Previous: Foundations. Next: Tensor parallelism.
Work through the interactive lesson below: step through each diagram, try the controls, and check yourself with the questions at the end.
References
- PyTorch Distributed: Experiences on Accelerating Data Parallel Training (Paper). Li et al., 2020. DDP design: gradient bucketing, overlap with backward, hooks.
- PyTorch Distributed Data Parallel design note (Docs). How DDP buckets gradients and launches all-reduce during backward.
- NCCL collective operations (Docs). All-reduce, reduce-scatter, all-gather and broadcast, with diagrams.
- An Empirical Model of Large-Batch Training (Paper). McCandlish et al., 2018. The gradient noise scale for choosing batch size.
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Paper). Rajbhandari et al., 2019. Stages 1–3 and their communication analysis.
- PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (Paper). Zhao et al., 2023. ZeRO-3-style sharding in PyTorch, with prefetching.
- ZeRO-Offload: Democratizing Billion-Scale Model Training (Paper). Ren et al., 2021. Optimizer states and the Adam step on the CPU.
- ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning (Paper). Rajbhandari et al., 2021. Offloading parameters, gradients and optimizer states to CPU and NVMe.
- ZeRO++: Extremely Efficient Collective Communication for Giant Model Training (Paper). Wang et al., 2023. Quantized weights and gradients, and a hierarchical parameter partition.
- DeepSpeed configuration: ZeRO optimizations (Docs). Every zero_optimization key used on this page.
- PyTorch FSDP2: fully_shard (Docs). The per-module sharding API and reshardafterforward.
- The Ultra-Scale Playbook (Guide). Hugging Face. The data parallelism and ZeRO chapters follow the same path as this page.