Results and write-up

The sweep across world size, the figures it produces, and what the write-up has to argue. These are measured numbers from the reference implementation, not predictions.

Resident bytes per parameter, measured

Three things to point at in the write-up. The DP column never moves — adding hardware does nothing for it. ZeRO-1 and ZeRO-2 flatten toward 4 and 2, their replicated floors, and will never go below them. ZeRO-3 keeps halving, because it replicates nothing.

At N = 1 all four are identical, which is worth including rather than skipping: with one rank there is nothing to shard, and a correct implementation should say so.

Peak, which is the number that decides fit

The ZeRO-3 peak falls much more slowly than its resident figure, and at N = 32 it is worse than ZeRO-2's resident 2.44. That is an artefact of the toy model gathering all 33,088 parameters at once, and it is worth saying so plainly: a real ZeRO-3 gathers one layer, so the added term is the largest layer rather than the whole model. The prototype overstates this cost, and pretending otherwise would be the easiest way to draw a wrong conclusion from correct code.

The figures to produce

import numpy as np
from zerosim.model import MLP
from zerosim.schemes import SCHEMES

m, STEPS = MLP(), 3
rows = {}
for N in (1, 2, 4, 8, 16, 32):
    rng = np.random.default_rng(3)
    batches = [[(rng.standard_normal((4, 64)).astype(np.float32),
                 rng.standard_normal((4, 64)).astype(np.float32))
                for _ in range(N)] for _ in range(STEPS)]
    out = []
    for S in SCHEMES:
        s = S(m, N)
        for t in range(STEPS):
            s.step(batches[t])
        out.append((s.bytes_per_param(),
                    s.peak_bytes_per_param(),
                    s.bytes_comm_per_rank() / STEPS / (m.n_params * 4)))
    rows[N] = out

The write-up

The numbers and the structure come from the experiment. The reasoning is the part that has to be yours, and the most useful thing to write about is the places where the measurements do not match the tidy story.