DeepSpeed ZeRO: Three-Stage Training-Memory Optimization, Principles and Practice

Note: This is a review article (no first-hand experimental data), based on public documentation and mechanism analysis.

Experiment Objective

Large-model training cannot avoid the memory wall: a 70B model's BF16 weights are 140GB; adding gradients and optimizer state (Adam's two momentums), a single 80GB card cannot fit them at all. Ch09's FSDP article covered one path; this article covers the other mainstream route, DeepSpeed ZeRO: is it a "replacement" for FSDP? What do the three stages shard, and how much do they save? Where is the boundary of offloading to CPU/NVMe? Goal: a complete ZeRO picture — three-stage principles, the correspondence with FSDP, script integration, and the "which tier should I use" decision.

Code Walkthrough

Part 1: Integrating ZeRO into a training script

import deepspeed

# Retrofitting an existing PyTorch training script: one init + engine swap
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=params,
    config_params=ds_config,
)
for step, batch in enumerate(train_loader):
    loss = model_engine(batch)          # forward
    model_engine.backward(loss)         # backward + ZeRO internal communication
    model_engine.step()                 # optimizer update

Key points:

1. Low integration cost: deepspeed.initialize returns a model_engine that replaces the native model; forward/backward/step calls stay the same — ZeRO's communication orchestration (All-Gather/Reduce-Scatter) is encapsulated inside the engine.
2. config_params is where all the magic lives: besides zero code changes, all tuning is in the config file — optimizer type, ZeRO stage, offload switches, communication backend.

Part 2: The three-stage config

# deepspeed_config.json: three ZeRO stage examples
{
  "train_batch_size": 256,
  "optimizer": { "type": "AdamW", "params": { "lr": 3e-5 } },
  "zero_optimization": {
    "stage": 2,                        # 1 / 2 / 3
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "allgather_partitions": true,
    "reduce_scatter": true,
    "overlap_comm": true               # overlap communication with compute
  },
  "communication_data_type": "bf16"
}

• stage 1: shard only optimizer state (Adam's dual momentums dominate, ~2× weights).
• stage 2: additionally shard gradients (gradients = weight size); each card keeps only its shard's gradients.
• stage 3: shard weights too (like FSDP full-sharding); memory drops to "one layer fits" level.
• offload_optimizer to CPU: optimizer state moves to CPU RAM, saving more GPU memory at the cost of CPU↔GPU transfer overhead.
• overlap_comm: overlaps gradient communication with the next compute segment, offsetting some sharding communication latency.

Principles

What the three stages shard, and how much they save

With 70B, BF16, AdamW (estimated):

Item	Full (per card)	ZeRO-1	ZeRO-2	ZeRO-3
Optimizer state (Adam)	2× weights	1/DP×	1/DP×	1/DP×
Gradients	1× weights	1×	1/DP×	1/DP×
Weights	1× weights	1×	1×	1/DP×
Total (DP=8)	4× weights	~3×	~1.5×	~0.5×+one layer

(DP = data-parallel degree; each row is per-card footprint in weight multiples.)

Understanding points:

• Optimizer state is the big one: Adam holds fp32 master weights plus two momentums ≈ 2× weights, often in fp32 (16 bytes/parameter). Sharding it at stage 1 halves memory immediately.
• Stage 3 is "load one layer at a time during training": after weight sharding, a card only needs the current compute layer, theoretically training arbitrarily large models — equivalent to FSDP's Full Shard; the mechanisms are nearly identical (FSDP later became the native PyTorch implementation; ZeRO-3 is its DeepSpeed implementation).
• Saving memory costs communication: the more sharding, the more All-Gather/Reduce-Scatter traffic. ZeRO-2 communicates less than ZeRO-3; prefer the lower stage when it fits.

ZeRO vs FSDP: not a one-or-the-other

Dimension	ZeRO-3 (DeepSpeed)	FSDP (native PyTorch)
Mechanism	shard weights/gradients/optimizer state	same (Full Shard)
Integration	deepspeed.initialize	torch.distributed.fsdp.FullyShardedDataParallel
Ecosystem	DeepSpeed family (Offload, gradient compression)	native PyTorch, works with torch.compile
Fit	existing DeepSpeed deps / need Offload	pure-PyTorch projects first

Selection judgment: new pure-PyTorch projects prefer FSDP (native, stable); use ZeRO when you need CPU/NVMe Offload or DeepSpeed ecosystem integration (MoE training, etc.). The mechanisms are homologous; migration cost is low.

The Offload boundary

offload_optimizer: device: cpu moves optimizer state to CPU RAM — at the cost of a CPU↔GPU transfer every step, potentially dropping training throughput 10%-30%. Offload is "a remedy when memory is insufficient," not a default. NVMe Offload (ZeRO-Infinity) goes further but has even larger IO latency — only for ultra-large models in slow training. Judgment: if memory fits (lower stage or non-offload higher stage), never offload.

Results

(Estimated, not measured; DP=8, 70B, BF16 per-card memory)

Scheme	Per-card memory estimate	Comm overhead	Fit
DDP full	~4× weights (560GB+)	none	small models
ZeRO-2	~1.5× weights (~200GB)	medium	common production tier
ZeRO-3	~0.5× weights + one layer	high	very large models
ZeRO-2 + CPU offload	weights + a little	medium + transfers	tight memory

Engineering Takeaways

1. Compute "how short you are" before choosing the stage: use the multiplier table to calculate the current memory gap; if ZeRO-2 solves it, do not jump to ZeRO-3 — higher sharding tiers bring heavier communication and diminishing throughput.
2. Offload is a remedy, not a default: do not offload when memory fits; when it does not, prefer CPU offload of optimizer state (least impact); NVMe offload is the last resort.
3. Turn on overlap_comm: compute-communication overlap hugely affects sharded-training throughput; it is the highest value-per-switch config in ZeRO/FSDP.
4. Link with the Checkpoint/stability articles: under sharded training, what is saved is "shard snapshots" (each card saves its shard), and recovery must gather them — ch10's restart strategies must adapt to the ZeRO shard format.

Next article preview: with the training-side memory problem solved (and its inference counterpart covered in ch12's KV Cache), next we return to applications, dissecting "Scientific Agents: From Tool Calling to Workflow Automation."