ImplementCore implementation

ZeRO: per-device model-state memory planner

Implement `run(psi, nd, stages)` returning the per-device model-state memory in bytes for mixed-precision Adam, given parameter count Ψ, data-parallel degree Nd, and a list of ZeRO stages (0 = full DP replication, 1 = P_os, 2 = P_os+P_g, 3 = all). Use K=12 optimizer multiplier. Round to the nearest integer byte.

python12 mindifficulty 2/4

From ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Sample tests

run([7500000000,64,0])→120000000000
run([7500000000,64,1])→31406250000
run([7500000000,64,2])→16640625000
run([7500000000,64,3])→1875000000

Hints

Full DP replicates everything: 2Ψ (fp16 params) + 2Ψ (fp16 grads) + 12Ψ (fp32 master, momentum, variance) = 16Ψ.
P_os shards optimizer states across Nd ranks; parameters and gradients stay replicated.
P_g additionally shards gradients; P_p shards parameters too. A stage applies cumulatively to each component it covers.
Stage math: stage 0 → 16Ψ; stage 1 → 2Ψ + 2Ψ + 12Ψ/Nd; stage 2 → 2Ψ + (2Ψ+12Ψ)/Nd; stage 3 → (2Ψ+2Ψ+12Ψ)/Nd.

Hints

Code runs in an isolated sandbox on the OpenSandbox judge. No model key required.