Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

ZeRO Optimization › Memory Savings Across Stages

Michael BrenndoerferJanuary 21, 20264 min read

Language AI HandbookAll chapters

Part 31 of 67Training Infrastructure

Chapter 225ZeRO Optimization › Memory Savings Across Stages

Reading section 6 of 14 · Memory Savings Across Stages · 4 minSection 6 of 14 · 4 min

Section 6 of 14 · 4 min read

ZeRO Optimization › Memory Savings Across StagesOne section · Sections in this chapter6 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's bring all three stages together and visualize their memory savings as a function of GPU count.

Out[10]:

Visualization

Line plot showing per-GPU memory in GB vs GPU count for baseline, ZeRO Stage 1, 2, and 3. — Per-GPU memory usage across ZeRO stages for a 7B parameter model as the number of GPUs increases. The A100 80 GB memory limit is shown as a reference line. All three ZeRO stages fit below the boundary at 2 GPUs, but only Stage 3 continues toward zero because it partitions every component.

The A100 80 GB line marks the feasibility boundary for an important class of hardware. The baseline never dips below 112 GB, making a 7B model impossible to train on A100s under standard data parallelism regardless of how many GPUs you add. All three ZeRO stages cross below the feasibility line at 2 GPUs in this model, but they diverge sharply as the GPU count grows. The curve for Stage 3 is a true hyperbola (linear reduction), while Stages 1 and 2 flatten out as the unpartitioned components dominate at high GPU counts.

Out[11]:

Visualization

Bar chart comparing per-GPU memory across ZeRO stages for models from 1B to 70B parameters. — Per-GPU memory across ZeRO stages for model sizes from 1B to 70B parameters at 64 GPUs. The red line shows the 80 GB A100 capacity. Baseline is infeasible above 1B, while Stage 3 keeps every model shown below 20 GB per GPU.

At 64 GPUs, Stage 3 fits a 70B model into 17.5 GB per GPU, well within even older GPU hardware with 24 or 32 GB. The baseline requires 1.12 TB per GPU for a 70B model, which is impossible on any current GPU. Stage 1 reduces that footprint to about 293 GB per GPU, so parameter and gradient replication still make it infeasible. This visualization illustrates why the choice of ZeRO stage has practical consequences: Stage 1 makes moderate models practical, while Stage 3 gives you access to a completely different tier of model sizes.