ZeRO
In plain wordsA memory-saving technique for training big models on many chips: instead of every chip keeping its own complete copy of all the training bookkeeping, the chips divide it among themselves and share pieces only when needed.
A family of memory optimizations for data parallelismdata parallelism, n.: Splitting the training data among many chips that each hold a full copy of the model; every chip learns from its own share, and they pool what they learned after each step so all copies stay identical. ('Zero Redundancy Optimizer') that removes the redundancy of every worker storing a full copy of the model state. In mixed-precision Adam training each parameter costs about 16 bytes — 16-bit weights and gradients plus 32-bit master weights and two optimizer moments — most of it optimizer state. ZeRO partitions this state across the data-parallel workers in three stages: stage 1 shards optimizer states (about 4× memory reduction), stage 2 also shards gradients (about 8×), and stage 3 also shards the parameters (a reduction that grows linearly with the number of workers), gathering each layer's parameters just in time for computation and discarding them afterwards. Stages 1 and 2 keep data parallelism's communication volume; stage 3 adds about 50%. Extensions offload state to CPU memory or NVMe storage (ZeRO-Offload, ZeRO-Infinity) or compress communication (ZeRO++). PyTorch's FSDPfully sharded data parallel, n.: A feature of the popular PyTorch software that lets many chips train a model too big for any one of them, by spreading the model and its training bookkeeping evenly across all the chips. is a closely related implementation.
ZeRO stage 2 let us fit the 13B model on 32 GPUs without tensor parallelism.
Origin
Introduced by Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase and Yuxiong He at Microsoft (arXiv, October 2019; SC20) and released in February 2020 in the DeepSpeed library, where it helped train Microsoft's 17-billion-parameter Turing-NLG.
Usage note
Despite the name, not an optimizer algorithm like AdamAdam, prop. n.: The most widely used method for training neural networks. It smooths the direction of updates over time and automatically sets a separate step size for each setting based on how large its recent changes have been.: ZeRO changes where optimizer and model state are stored, not how updates are computed. Written with capital Z, R and O and a lowercase e.
Related
- fully sharded data parallelfully sharded data parallel, n.: A feature of the popular PyTorch software that lets many chips train a model too big for any one of them, by spreading the model and its training bookkeeping evenly across all the chips.
- data parallelismdata parallelism, n.: Splitting the training data among many chips that each hold a full copy of the model; every chip learns from its own share, and they pool what they learned after each step so all copies stay identical.
- DeepSpeedDeepSpeed, prop. n.: Free software from Microsoft for training and running AI models too big for ordinary tools, spread over many chips; it is best known for a memory-saving technique called ZeRO.
- mixed-precision trainingmixed-precision training, n.: Training a model with shorter, less exact numbers for most of the arithmetic and full-length numbers only where accuracy really matters, which saves a great deal of memory and lets modern chips train far faster.
- AdamAdam, prop. n.: The most widely used method for training neural networks. It smooths the direction of updates over time and automatically sets a separate step size for each setting based on how large its recent changes have been.
- offloadingoffloading, n.: Moving part of an AI model's data or work out of a chip's scarce fast memory onto ordinary memory, a disk or another computer, so that models too big for the chip can still run.
- distributed trainingdistributed training, n.: Training one AI model on many computer chips at once, splitting up the data or the model itself among them and keeping them in sync, so that jobs that would take a single chip thousands of years finish in weeks.
- model parallelismmodel parallelism, n.: Splitting a single AI model that is too big for one chip into pieces spread over several chips, which must constantly pass intermediate results to each other as they work.
Sources
- Rajbhandari, S., Rasley, J., Ruwase, O. & He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20. arXiv:1910.02054. arxiv.org/abs/1910.02054
- Microsoft Research (2020). ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters. Microsoft Research Blog. microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters
Mentioned in
data parallelismdata parallelism, n.: Splitting the training data among many chips that each hold a full copy of the model; every chip learns from its own share, and they pool what they learned after each step so all copies stay identical., distributed trainingdistributed training, n.: Training one AI model on many computer chips at once, splitting up the data or the model itself among them and keeping them in sync, so that jobs that would take a single chip thousands of years finish in weeks., full fine-tuningfull fine-tuning, n.: Fine-tuning in which every one of a model's internal numbers is allowed to change — the most flexible option, but also the most expensive in memory and storage., gradient checkpointinggradient checkpointing, n.: A memory-saving trick: instead of storing every intermediate result needed for learning, the computer keeps only a few and recalculates the rest when needed, trading extra computing time for memory., optimizeroptimizer, n.: The rule a training program uses to turn measured errors into actual changes to a model's internal numbers, deciding the direction and size of each adjustment., 3D parallelism3D parallelism, n.: Combining three ways of dividing up training work — by data, by layer stage and by slices within layers — so that thousands of chips can train one giant model efficiently., collective communicationcollective communication, n.: The standard patterns by which a whole group of chips exchanges data at once — for example, everyone adding up their numbers and all getting the total — which large training jobs use constantly., DeepSpeedDeepSpeed, prop. n.: Free software from Microsoft for training and running AI models too big for ordinary tools, spread over many chips; it is best known for a memory-saving technique called ZeRO., fully sharded data parallelfully sharded data parallel, n.: A feature of the popular PyTorch software that lets many chips train a model too big for any one of them, by spreading the model and its training bookkeeping evenly across all the chips., offloadingoffloading, n.: Moving part of an AI model's data or work out of a chip's scarce fast memory onto ordinary memory, a disk or another computer, so that models too big for the chip can still run.