The two checks
Two assertions decide whether the prototype is worth anything. The first is that all four schemes learn identically. The second is that the measured bytes match the formulas. Everything else is presentation.
Check one: identical learning
Same seed, same batches in the same order, same initial weights. All four schemes compute the same averaged gradient and apply the same update, so the parameters should agree exactly — not approximately.
import numpy as np
from zerosim.model import MLP
from zerosim.schemes import SCHEMES
def test_all_schemes_match_data_parallel():
m, N, STEPS = MLP(), 8, 10
rng = np.random.default_rng(7)
batches = [[(rng.standard_normal((8, 64)).astype(np.float32),
rng.standard_normal((8, 64)).astype(np.float32))
for _ in range(N)] for _ in range(STEPS)]
final = {}
for S in SCHEMES:
s = S(m, N)
for t in range(STEPS):
s.step(batches[t])
final[s.name] = s.full_theta()
base = final["data parallel"]
for name, theta in final.items():
assert np.array_equal(theta, base), name
This passes on the reference implementation with a maximum difference of exactly zero at every step, for every scheme.
| Scheme | Final loss | max |θ − θDP| |
|---|---|---|
| data parallel | 1.021900 | 0 |
| ZeRO-1 | 1.021900 | 0 |
| ZeRO-2 | 1.021900 | 0 |
| ZeRO-3 | 1.021900 | 0 |
Why exactly zero, and when to expect otherwise
The reference implementation reduces with the same summation order in every scheme, so the floating-point result is bitwise identical. A real framework will not give you this: NCCL's ring sums in a different order than a naive reduction, so you would see differences around 1e-7 growing slowly over steps. Either result is correct — but if you see exact zeros, say why, and if you see 1e-7, say that too. Claiming bitwise equality without having checked the summation order is the kind of thing worth being careful about.
Check two: measured against predicted
The formulas are not fitted to the measurements; they come from counting bands. So this is a real test rather than a restatement.
import pytest
from zerosim.model import MLP
from zerosim.schemes import DP, ZeRO1, ZeRO2, ZeRO3
PREDICTED = {
DP: lambda N: 16.0,
ZeRO1: lambda N: 4 + 12 / N,
ZeRO2: lambda N: 2 + 14 / N,
ZeRO3: lambda N: 16 / N,
}
@pytest.mark.parametrize("N", [1, 2, 4, 8, 16, 32])
@pytest.mark.parametrize("S", list(PREDICTED))
def test_resident_bytes_match_formula(S, N):
s = S(MLP(), N)
assert s.bytes_per_param() == pytest.approx(PREDICTED[S](N))
import numpy as np, pytest
from zerosim.model import MLP
from zerosim.schemes import DP, ZeRO1, ZeRO2, ZeRO3
@pytest.mark.parametrize("N", [2, 4, 8, 16, 32])
@pytest.mark.parametrize("S,mult", [(DP, 2), (ZeRO1, 2), (ZeRO2, 2), (ZeRO3, 3)])
def test_traffic_matches_ring_cost(S, mult, N):
m = MLP()
s = S(m, N)
rng = np.random.default_rng(0)
batch = [(rng.standard_normal((4, 64)).astype(np.float32),
rng.standard_normal((4, 64)).astype(np.float32)) for _ in range(N)]
s.step(batch)
unit = m.n_params * 4 # the communicated dtype
got = s.bytes_comm_per_rank() / unit
assert got == pytest.approx(mult * (N - 1) / N)
What the traffic test actually asserts
Not 2P and 3P, but 2(N−1)/N and 3(N−1)/N. Writing the
test against the exact ring cost rather than the quoted approximation is what makes it a test:
an implementation that accidentally does a naive all-to-all would pass a loose 2P check and
fail this one.
| N | DP | ZeRO-1 | ZeRO-2 | ZeRO-3 | 2(N−1)/N |
|---|---|---|---|---|---|
| 2 | 1.0000 | 1.0000 | 1.0000 | 1.5000 | 1.0000 |
| 4 | 1.5000 | 1.5000 | 1.5000 | 2.2500 | 1.5000 |
| 8 | 1.7500 | 1.7500 | 1.7500 | 2.6250 | 1.7500 |
| 16 | 1.8750 | 1.8750 | 1.8750 | 2.8125 | 1.8750 |
| 32 | 1.9375 | 1.9375 | 1.9375 | 2.9062 | 1.9375 |
The three left columns being identical is the point of the session made numerically: stages 1 and 2 cost exactly what data parallelism already cost. And nothing reaches 2P or 3P at any size — it approaches from below, which is worth stating in the README rather than rounding away.