Python starts when you choose Run.
ZeRO’s Memory Budget, Byte by Byte
Parameters, gradients, and optimizer state account for different parts of training memory. ZeRO changes who stores each part.
Original dtype-explicit arithmetic; persistent state only, with no framework tensor inspection or peak-memory measurement.
The code uses Python 3 and the standard library. Saved output is from the checked CPU run.
stage=0: params=2.0 GB, grads=2.0 GB, optimizer=12.0 GB, total=16.0 GB stage=1: params=2.0 GB, grads=2.0 GB, optimizer=3.0 GB, total=7.0 GB stage=2: params=2.0 GB, grads=0.5 GB, optimizer=3.0 GB, total=5.5 GB stage=3: params=0.5 GB, grads=0.5 GB, optimizer=3.0 GB, total=4.0 GB
Sources
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models: https://arxiv.org/html/1910.02054v3
Python runs locally in this page. Saved results and figures are visible before execution.