ZeRO-Infinity

For a Transformer with nl layers and hidden dimension hd (params approx 12 nl hd^2, model-state memory 240 nl hd^2 bytes, Eq. 1-2), Model State Working Memory for the largest linear layer (hd -> 4 hd) is 16 hd^2 bytes (Eq. 4) and Activation Working Memory is bsz * seq * ci * (16 hd + 2 attn_heads * seq) bytes (Eq. 5); with per-iteration compute 8 * bsz * seq * params FLOPs (Eq. 7-8), the arithmetic intensities with respect to parameters and gradients, optimizer states, and activation checkpoints reduce exactly to ait_{param+grad} = seq * bsz (Eq. 9), ait_{opt} = seq * bsz / 4 (Eq. 10), and ait_{act_ckpt} = 24 * hd * ci (Eq. 11), governing bandwidth efficiency via efficiency = (ait * bw) / (ait * bw + peak_tp) (Eq. 6).

Memory-Centric Tiling decomposes a large linear operator into a mathematically equivalent sequence of T smaller parameter tiles that are fetched and released one at a time via ZeRO-3, reducing peak GPU Model State Working Memory by factor 1 / T without tensor-slicing model parallelism (Section 5.1.3, scaling hidden dimension from 8K to 64K in Figure 6b); simultaneously, Bandwidth-Centric Partitioning shards each parameter across all dp data-parallel processes and reconstructs it via allgather instead of single-owner broadcast, activating all dp PCIe links in parallel so effective slow-memory-to-GPU bandwidth scales linearly with dp (Section 6.1).

By decorating torch.nn.Module.__init__ to partition and offload each submodule's parameters immediately upon construction (Section 7.2) and injecting recursive pre/post forward-backward hooks with external-parameter auto-registration (Section 7.1), ZeRO-Infinity fine-tunes up to 1-trillion-parameter models on a single 16-GPU DGX-2 node without model parallelism and scales up to 32-trillion-parameter models across 32 DGX-2 nodes (512 V100 GPUs), sustaining over 25 petaflops (49 TFLOPS/GPU, ~40% of peak) with superlinear weak scaling.

Because optimizer states have 4x lower arithmetic intensity (seq * bsz / 4) than parameters and gradients and are updated at the end of the backward pass where they cannot be fully overlapped with forward-backward GPU compute, achieving high efficiency at small per-GPU batch sizes requires high aggregate NVMe/CPU bandwidth (~1.5 TB/s across the cluster), and at extreme 20T-parameter scale on 512 GPUs, CPU memory capacity for activation checkpoints limits per-GPU micro-batch size to 1.25, reducing throughput from 49 to 34 TFLOPS/GPU.