ZeRO Optimizer State Sharding
Adam stores two moment estimates per parameter, both in float32. A 7B-parameter model carries 56 GB of optimiser state. ZeRO stage 1 shards that across N ranks; each rank owns 1/N of the optimiser. After the local step the updated parameter shards broadcast back, every rank reconstructs the full model, and the next step begins. The win is a linear memory drop on the largest single allocation in the training stack. Shard optimiser state (first moment, second moment, fp32 master copy) across N ranks so each rank owns 1/N. Use reducescatter to deliver each rank only its shard's gradient sum, then allgather to broadcast the updated parameter shards back. Compute the memory savings table for stage 1, stage 2, stage 3 against vanilla DDP. Defend the choice of stage 1 vs stage 2 vs stage 3 on model size and bandwidth budget. Vanilla DDP replicates everything: parameters, gradients, and optimiser state are present in full on every rank. For a 7B-parameter model in fp16 that means 14 GB of parameters, 14 GB of gradients, and 28 GB of optimiser state per rank. The optimiser state is the largest term and the easiest to shard because it is only touched during the step, not during forward or backward. ZeRO stage 1 shards the optimiser state. Each rank holds 1/N of the…
ZeRO Optimizer State Sharding: Adam stores two moment estimates per parameter, both in float32. A 7B-parameter model carries 56 GB of optimiser state. ZeRO…
This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.
Browse the complete course catalog or open this lesson on GitHub.