Our team is struggling with "Out of Memory" errors when training our latest 70B parameter setup. We heard DeepSpeed offers offloading features. Is it possible to leverage system RAM or NVMe to keep costs down, or is that just going to slow everything down to the point where it’s not worth it for production?
3 answers
You are likely looking for ZeRO-Offload or ZeRO-Infinity. These features are game-changers for teams on a budget. ZeRO-Offload specifically moves the optimizer state and gradients to the CPU memory, while keeping the active computation on the GPU. Since the optimizer is often the biggest memory hog, this frees up a massive amount of VRAM. While it does add some overhead because of the PCI-E data transfer, the ability to train a model that literally wouldn't fit otherwise is invaluable. For a 70B model, you might see a performance dip, but the cost-to-capability ratio stays very high compared to renting 8x A100 nodes.
How much RAM does the host machine actually need to support ZeRO-Offload for a 70B model? I’ve heard the ratio can be quite demanding on the CPU side.
We used the NVMe offloading in DeepSpeed-Infinity for a massive research project. It was slower, but it allowed us to experiment with scales we never thought possible.
I agree with Kevin. The flexibility is the real value here. Even if it's slower, the fact that you can run these workloads on "consumer-grade" enterprise servers is a huge win for SEO and tech startups.
Jason, for a 70B model, you generally want at least 512GB of system RAM to be safe, as the optimizer states for Adam can take up about 12 bytes per parameter. It's a heavy requirement for the motherboard, but still significantly cheaper than buying or renting more high-end GPUs with larger HBM capacities.