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.

Build 1.5 hr Python

Prerequisites: Phase 19 Track C lessons 42-49

#Learning Objectives

  • Shard optimiser state (first moment, second moment, fp32 master copy) across N ranks so each rank owns 1/N.
  • Use reduce_scatter 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.
Full Access lesson

Keep going. The rest of this lesson is one step away.

Unlock all 503 lessons, video, code, quizzes and every future update. Pay in your own currency with Paystack.

The text and code of this lesson are based on AI Engineering from Scratch by Rohit Ghumare, MIT licensed.