The four schemes
One base class holding the bands and the Adam update, and four subclasses that differ only in which bands are sharded and which collectives run. Written this way the diff between stages is the explanation.
Three flags describe all four schemes
| Scheme | shard_opt | shard_grad | shard_param | Collectives per step |
|---|---|---|---|---|
| data parallel | — | — | — | all-reduce |
| ZeRO-1 | yes | — | — | reduce-scatter, all-gather |
| ZeRO-2 | yes | yes | — | reduce-scatter, all-gather |
| ZeRO-3 | yes | yes | yes | all-gather ×2, reduce-scatter |
Note that ZeRO-2 subclasses ZeRO-1 and changes one flag, with no new code at all. That is the honest shape of the difference: stage 2 does not do anything extra, it declines to keep gradients it does not own. If your implementation needs new logic for stage 2, something has been misunderstood.
Where the third P comes from
ZeRO-3 is the only stage that overrides step, and the reason is worth getting
right because it is easy to implement a version that looks correct and undercounts.
A card holding only its own slice of the parameters cannot run the forward pass, so it
gathers. It then releases, because holding the gathered copy is what it was trying to avoid. So
when the backward pass needs those same weights, it has to gather again. Two
gathers plus one reduce-scatter is 3(N−1)/N, which is the 3P from the
session.
The mistake to avoid
Gathering once and keeping the result through both passes gives the right loss and the wrong numbers: traffic comes out at 2P and peak memory never shows the gathered copy. The prototype would then contradict the thing it was built to demonstrate. Release between the passes even though it costs a second gather.
"""Data parallel and ZeRO stages 1-3 over a flat parameter vector.
Every scheme exposes the same `step(batches) -> loss` and the same
`bytes_per_rank()` accounting, so they can be compared directly.
"""
from __future__ import annotations
import numpy as np
from .comm import Comm
# bytes per parameter, per band
B_P16, B_GRAD, B_MASTER, B_M, B_V = 2, 2, 4, 4, 4
class Base:
name = "base"
shard_grad = False
shard_param = False
shard_opt = False
def __init__(self, model, world, lr=1e-2, betas=(0.9, 0.999), eps=1e-8):
self.model, self.world, self.lr = model, world, lr
self.b1, self.b2, self.eps = betas[0], betas[1], eps
self.comm = Comm(world)
self.n = model.n_params
assert self.n % world == 0, "parameter count must divide by world size"
self.shard = self.n // world
self.t = 0
# fp32 master, sharded or not depending on the scheme
if self.shard_opt:
self.master = [model.theta0[i * self.shard:(i + 1) * self.shard].copy()
for i in range(world)]
else:
self.master = [model.theta0.copy() for _ in range(world)]
self.m = [np.zeros_like(x) for x in self.master]
self.v = [np.zeros_like(x) for x in self.master]
# fp16 working copy: every rank holds all of it unless params are sharded
self.p16 = [model.theta0.astype(np.float16).copy() for _ in range(world)]
# ---- the update, run on whatever slice this rank owns -------------
def _adam(self, r, g):
self.m[r] = self.b1 * self.m[r] + (1 - self.b1) * g
self.v[r] = self.b2 * self.v[r] + (1 - self.b2) * g * g
mh = self.m[r] / (1 - self.b1 ** self.t)
vh = self.v[r] / (1 - self.b2 ** self.t)
self.master[r] -= self.lr * mh / (np.sqrt(vh) + self.eps)
def full_theta(self):
"""The fp32 model as a single vector, for loss curves and comparisons."""
if self.shard_opt:
return np.concatenate(self.master)
return self.master[0].copy()
# ---- memory accounting -------------------------------------------
def bytes_per_param(self):
n, w = self.n, self.world
p = B_P16 / (w if self.shard_param else 1)
g = B_GRAD / (w if self.shard_grad else 1)
o = (B_MASTER + B_M + B_V) / (w if self.shard_opt else 1)
return p + g + o
def peak_bytes_per_param(self):
"""ZeRO-3 also holds one gathered layer at a time; others do not."""
return self.bytes_per_param()
def bytes_comm_per_rank(self):
return self.comm.bytes_per_rank()
class DP(Base):
name = "data parallel"
def step(self, batches):
self.t += 1
losses, grads = [], []
for r in range(self.world):
x, y = batches[r]
l, g = self.model.loss_and_grad(self.master[r], x, y)
losses.append(l); grads.append(g)
avg = self.comm.all_reduce_mean(grads)
for r in range(self.world):
self._adam(r, avg[r])
self.p16[r] = self.master[r].astype(np.float16)
return float(np.mean(losses))
class ZeRO1(Base):
name = "ZeRO-1"
shard_opt = True
def step(self, batches):
self.t += 1
theta = np.concatenate(self.master)
losses, grads = [], []
for r in range(self.world):
x, y = batches[r]
l, g = self.model.loss_and_grad(theta, x, y)
losses.append(l); grads.append(g)
parts = self.comm.reduce_scatter_mean(grads) # rank r gets slice r
for r in range(self.world):
self._adam(r, parts[r])
gathered = self.comm.all_gather([m.astype(np.float32) for m in self.master])
for r in range(self.world):
self.p16[r] = gathered[r].astype(np.float16)
return float(np.mean(losses))
class ZeRO2(ZeRO1):
name = "ZeRO-2"
shard_grad = True # non-owned gradients are dropped after the scatter
class ZeRO3(ZeRO2):
name = "ZeRO-3"
shard_param = True
def step(self, batches):
self.t += 1
shards = [m.astype(np.float32) for m in self.master]
# gather for the forward pass, then release
theta_list = self.comm.all_gather(shards)
losses, acts = [], []
for r in range(self.world):
x, y = batches[r]
l, _ = self.model.loss_and_grad(theta_list[r], x, y)
losses.append(l)
theta_list = None # released
# gather again for the backward pass: this is the third P
theta_list = self.comm.all_gather(shards)
grads = []
for r in range(self.world):
x, y = batches[r]
_, g = self.model.loss_and_grad(theta_list[r], x, y)
grads.append(g)
theta_list = None
parts = self.comm.reduce_scatter_mean(grads)
for r in range(self.world):
self._adam(r, parts[r])
return float(np.mean(losses))
def peak_bytes_per_param(self):
# the persistent shard, plus one full fp32 copy while it is gathered
return self.bytes_per_param() + B_MASTER
SCHEMES = [DP, ZeRO1, ZeRO2, ZeRO3]
Two accounting methods, not one
bytes_per_param is what a rank holds between steps. peak_bytes_per_param
is what it holds at the worst moment. For three of the four schemes these are the same number,
and for ZeRO-3 they are not, because of the gathered copy.
| Scheme | Resident | Peak | Why they differ |
|---|---|---|---|
| data parallel | 16.000 | 16.000 | nothing is ever gathered |
| ZeRO-1 | 5.500 | 5.500 | parameters were never sharded |
| ZeRO-2 | 3.750 | 3.750 | same |
| ZeRO-3 | 2.000 | 6.000 | plus one gathered fp32 copy |
Reporting only the resident figure makes ZeRO-3 look three times better than it is at this size. Reporting only the peak hides what the sharding achieved. The prototype should print both, and the README should explain why the gap exists.