LAB / BUILD PLAN

Simulating ZeRO on 32 virtual GPUs

Build 32 virtual GPUs, run a small model across them, implement data parallel and all three ZeRO stages, and measure how memory and communication actually change. The deliverables are a notebook, a repository, and a README that explains the concepts in your own words.

What makes it convincing: not that the code runs, but that the measured numbers match what ZeRO predicts and that you can say why. The plan below is arranged so the measurements are checkable against the formulas from the earlier pages.

The decision that shapes everything

What a virtual GPU is determines what can be measured. Two workable routes.

Two ways to build 32 virtual GPUs
	A — explicit simulator	B — real torch.distributed
what a rank is	an object holding its own tensors	a process, gloo backend
collectives	written by hand, every byte counted	real all-reduce and friends
byte accounting	exact, by construction	inferred, or read from counters
32 ranks in Colab	trivial, one process	heavy — 32 processes
verdict	primary: it is what makes the numbers auditable	a 4-rank corroboration section

Route A as the main build, because the prototype has to show how memory and computation change, and that needs exact per-rank accounting rather than an operating system memory reading. Then a short Route B section with four real ranks, checking the simulator's formulas against genuine collectives. Doing only B leaves you unable to prove the memory claims; doing only A invites the objection that it is all bookkeeping.

The eight steps

Virtual GPU and communication layer
A Rank holding its own shard of each band, and a Comm class wrapping all-reduce, reduce-scatter and all-gather. Every method increments a byte counter tagged by collective and by rank. This counter is the whole experiment — build it first and trust nothing that does not go through it.
The demo model
A three-layer MLP or one transformer block, sized so the parameter count divides evenly by 32. Adam with real moments, an fp32 master copy, and fp16 working copies, so all four bands physically exist rather than being asserted.
Data parallel as the baseline
Every rank holds everything; split the batch, all-reduce the gradients, apply the same update everywhere. Assert that all 32 ranks hold bitwise-identical weights after the step. If that assertion fails nothing later is meaningful.
The three ZeRO stages behind one interface
Each stage exposes the same step(batch) -> loss. Stage 1 shards optimizer state and all-gathers parameters after the update. Stage 2 also shards gradients and drops the non-owned ones after the reduce-scatter. Stage 3 shards parameters and gathers each layer on demand, releasing it immediately.
Instrumentation
Three numbers per step: resident bytes per rank broken down by band and sampled at peak, bytes communicated by collective, and FLOPs plus wall time.
The two checks
Correctness, then theory. Both are described below, and they are the part anyone reading this will look at first.
Sweep the world size
N = 1, 2, 4, 8, 16, 32 for all four schemes. Plot per-rank memory against N, which reproduces the memory-wall chart with measured data instead of formulas.
Write it up
Notebook, repository with the simulator as an importable module and tests, and the README.

Check one: they must all learn the same thing

Same seed, same data order, same initial weights. All four schemes must produce the same loss curve, because all four compute the same gradient and apply the same update. A divergence is a bug, not a property.

The most convincing plot in the notebook

Four loss curves lying exactly on top of each other, with a second panel showing the maximum absolute difference between schemes at each step sitting at or near zero. That single figure demonstrates you implemented ZeRO rather than something that merely uses less memory. Show the numbers, not just the lines: report the largest weight difference across all ranks and all steps.

Check two: measured against predicted

The formulas from the earlier pages, evaluated at N = 32. Your measured column should reproduce these to within rounding.

What to expect at 32 virtual GPUs
Scheme	Formula	Predicted bytes/param	vs DP	Traffic
data parallel	`16`	16.0000	1.00×	2P
ZeRO-1	`4 + 12/N`	4.3750	3.66×	2P
ZeRO-2	`2 + 14/N`	2.4375	6.56×	2P
ZeRO-3	`16/N`	0.5000	32.00×	3P

Two traps worth knowing before you hit them.

The 2P figure is an approximation. The exact ring cost is 2(N−1)/N × P, which is 1.9375P at N = 32, not 2P. If your measured traffic comes out slightly under 2P, that is correct and worth saying so in the README rather than quietly rounding.

Peak memory is not end-of-step memory. Under ZeRO-3 a rank holds its shard plus whichever layer is currently gathered. Sampling at the end of the step will show 16/N and miss the peak, which is the number that decides whether a real run fits. Measure both and report both.

A suggested repository shape

layout

zero-sim/
  README.md            the write-up, in your own words
  notebook.ipynb       the narrative run with plots
  zerosim/
    comm.py            Comm, with the byte counters
    ranks.py           Rank, band-level storage
    model.py           the demo model
    schemes.py         dp, zero1, zero2, zero3
    measure.py         memory and FLOP accounting
  tests/
    test_equivalence.py   all schemes match DP bitwise
    test_memory.py        measured bytes match the formulas
    test_traffic.py       measured traffic matches 2(N-1)/N * P

The tests are worth more than they cost. test_equivalence.py is the prototype's central claim asked as an assertion, and having it pass in CI is a stronger claim than a paragraph saying it works.

What the README has to argue

The instruction is explicit that this part is yours rather than an agent's, so the honest division is that the numbers and structure come from the experiment and the reasoning comes from you. Six things it should answer.

README checklist
Question	What a good answer contains
What does each stage shard?	optimizer state, then gradients, then parameters — and which of the sixteen bytes that is
Why are stages 1 and 2 free?	an all-reduce already contains a reduce-scatter; they keep the intermediate instead of discarding it
Why does stage 3 cost 3P?	parameters must be gathered before the forward pass and again in the backward
Why can ZeRO-1 never fit some models?	the 4-byte replicated floor does not divide by N
When would you pick each?	your own measured step times, and which constraint was binding
Where does ZeRO stop being enough?	when one layer alone exceeds a card — then the model, not its state, has to be split

On the "not your agent" instruction

The defensible way to use help here is to let it write the plumbing and the plots, and to write the interpretation yourself — including the parts where your measurements disagree with the tidy formulas. A README that says "traffic came out at 1.94P, not 2P, and here is why" is evidence of understanding in a way that a correct restatement of the paper is not.

What to build first

The byte counter, then data parallel, then the equivalence test. Once those three pass, each ZeRO stage is a small change against a harness that already tells you when you have broken something. Building the stages first and the instrumentation last is the way this kind of prototype usually goes wrong.