Optimiseurs ZeRO et Sharded

ZeRO, the Zero Redundancy Optimizer, is a way to train neural networks that are too big for one GPU's memory while keeping the simple programming model of data parallelism. In ordinary data parallelism, every GPU holds a full copy of the model, processes a different slice of the training batch, and the GPUs average their gradients after each step. That is easy to use but wasteful: with 64 GPUs, there are 64 identical copies of everything. ZeRO removes that duplication by giving each GPU responsibility for only a fraction of the training state, and fetching the rest from its peers just when it is needed. A sharded optimizer is the general name for this pattern; PyTorch's Fully Sharded Data Parallel (FSDP) is the best-known implementation outside Microsoft's DeepSpeed library.

The reason duplication hurts is that the weights are not the biggest item. The ZeRO paper does the accounting for mixed-precision training with the Adam optimizer, where each parameter needs 2 bytes for a 16-bit copy of the weight, 2 bytes for its 16-bit gradient, and 12 bytes of 32-bit optimizer state: a master copy of the weight plus Adam's two running averages, momentum and variance. That is 16 bytes per parameter, so a 7.5-billion-parameter model needs about 120 GB just for model state, before counting activations, the intermediate values saved for the backward pass. The optimizer state alone is three quarters of the total, which is why sharding it first gives the biggest win.

ZeRO works in three cumulative stages. Stage 1 partitions the optimizer states: each of N GPUs keeps Adam's state for only 1/N of the parameters and updates only that slice, after which the updated weights are shared with everyone. Stage 2 also partitions the gradients: instead of every GPU receiving the full averaged gradient, a reduce-scatter operation sums gradients across GPUs and delivers each GPU only the slice it will update. Stage 3 partitions the 16-bit parameters too: each GPU stores only its shard of the weights, and before each layer runs in the forward or backward pass, the GPUs perform an all-gather to temporarily rebuild that layer's full weights, then free them afterwards. For the 7.5B example on 64 GPUs, the paper computes per-GPU model state falling from 120 GB to about 31.4 GB at stage 1, 16.6 GB at stage 2 and 1.88 GB at stage 3.

The idea was introduced by Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase and Yuxiong He of Microsoft in a paper first posted to arXiv in October 2019 and revised in May 2020. It reported training models of over 100 billion parameters on 400 GPUs at about 15 petaflops of throughput, an 8 times increase in trainable model size and up to 10 times better performance than the prior state of the art, and training models of up to 13 billion parameters without any model parallelism. A family of follow-ups came next. ZeRO-Offload (January 2021) moved optimizer state and the optimizer step onto CPU memory, reaching 40 teraflops on one V100 for a 10B model. ZeRO-Infinity (April 2021) added NVMe solid-state storage as a third memory tier and reported over 25 petaflops on 512 V100 GPUs. ZeRO++ (June 2023) attacked communication instead, cutting ZeRO's communication volume by 4 times with quantized collectives. PyTorch documented its own FSDP design in an April 2023 paper reporting performance comparable to standard Distributed Data Parallel on models that DDP cannot hold.

ZeRO is one of several ways to split a model. Tensor parallelism, as in NVIDIA's Megatron-LM, cuts individual matrix multiplications across GPUs, which needs very fast links and changes the model code. Pipeline parallelism assigns different layers to different GPUs and streams micro-batches through them, which introduces idle bubbles and scheduling complexity. ZeRO instead keeps every GPU running the whole model on its own data and only changes where state is stored, so model code usually needs no rewriting. In practice, the largest training runs combine all three, with ZeRO or FSDP handling the data-parallel dimension. The ZeRO paper itself paired its stages with Megatron-style model parallelism to reach its 100-billion-parameter results.

The costs are mostly communication. Stages 1 and 2 move the same amount of data per step as plain data parallelism, which the paper counts as twice the parameter count, because the all-reduce that DDP already performs can be split into a reduce-scatter and an all-gather. Stage 3 raises that to three times the parameter count, a 1.5 times increase, because weights must be gathered once in the forward pass and again in the backward pass. On fast interconnects that overhead can be hidden by overlapping communication with computation, but on slower networks stage 3 can become network-bound. Offloading to CPU or NVMe trades GPU memory for PCIe bandwidth, so it can enable a model that would otherwise not fit while training it noticeably more slowly. Sharding also does nothing for activation memory on its own, so activation checkpointing is usually needed alongside it.

Using ZeRO in practice is mainly configuration. In DeepSpeed, a JSON block named zero_optimization sets the stage and options such as reduce_bucket_size, overlap_comm, and offload_optimizer or offload_param with a device of cpu or nvme. In PyTorch, applying fully_shard to each transformer block and then to the root module gives stage 3 behavior, with reshard_after_forward controlling whether weights are freed after the forward pass. One operational detail catches many users: sharded checkpoints are saved as many per-rank files, so producing a single ordinary weight file requires a consolidation step, such as DeepSpeed's zero_to_fp32.py script, which the docs say needs about twice the final checkpoint size in CPU memory. Choosing a stage is a trade-off: use the lowest stage that fits the model, because each higher stage adds communication.