DeepSpeed ZeRO-3 Offload: over 40 billion parameters on a single GPU
In one sentence Microsoft ships ZeRO-3 Offload in DeepSpeed, combining ZeRO Stage 3 (partitioning parameters too) with CPU offload: fine-tuning of 40B+ parameter models on a single 32 GB V100 and 2-trillion-parameter models on 512 GPUs.
Training a large AI model takes many GPUs, and each one usually keeps a full copy of the model's data in memory. With models of tens of billions of parameters, memory runs out quickly.
Microsoft's DeepSpeed library already split part of that data across GPUs. With ZeRO-3 it splits the model weights too: each GPU keeps only a slice and fetches the others when needed. On top of that, whatever doesn't fit on the GPU can be moved to the computer's main memory (RAM).
The stated result: fine-tuning a model with over 40 billion parameters on a single GPU, with almost no code rewrite, by turning on a couple of options in the configuration file.
Companies
Microsoft
Tools
DeepSpeed, ZeRO-3 Offload
Tags
Sources