単一GPUでの100Bモデルのファインチューニングを可能にし加速するためのNVMe SSDの追加

Adding NVMe SSDs to Enable and Accelerate 100B Model Fine-tuning on a Single GPU

March 11, 2024

著者: Changyue Liao, Mo Sun, Zihan Yang, Kaiqi Chen, Binhang Yuan, Fei Wu, Zeke Wang

cs.AI

要旨

大規模言語モデルの最近の進展は、その膨大なパラメータ数に由来する優れた能力により、世界に多大な価値をもたらしています。しかし、現在最高のメモリ容量を持つGPUでさえ、80GBに達する程度であり、確率的勾配降下法に基づく最適化を行う際に、これらの膨大なパラメータとそれに関連するオプティマイザの状態を収容するには全く不十分です。このような巨大なモデルをホストするための一つのアプローチは、多くのGPUのデバイスメモリを集約することです。しかし、このアプローチは、多くのハイエンドGPUサーバーに対して常に限られた予算を持つほとんどの学術研究者にとって、法外なコストを伴います。本論文では、ほとんどのAI研究者がアクセス可能な、一般的なサーバー内の単一の、さらにはローエンドのGPU上での巨大モデルのファインチューニングに焦点を当てます。このようなシナリオでは、最先端の研究であるZeRO-Infinityは、一般的なサーバーで実行する際に2つの深刻な問題に直面します：1)非効率的なスワッピングによる低いGPU利用率、2)CPUメモリ容量によるトレーニング可能なモデルサイズの制限です。その根本的な理由は、ZeRO-InfinityがハイエンドGPUサーバーでの実行に最適化されているためです。この目的のために、我々はFuyouを提案します。これは、ローエンドのサーバーとローエンドのGPU、そして限られたCPUメモリ容量で、効率的に100Bの巨大モデルのファインチューニングを可能にする低コストのトレーニングフレームワークです。その鍵となるアイデアは、SSD-CPU通信を最適化の次元として追加し、計算とデータスワッピングを体系的アプローチから慎重に共最適化することで、GPU利用率を最大化することです。実験結果は、1) FuyouがコンシューマーGPU RTX 4090上で175B GPT-3を高いGPU利用率でファインチューニングできるのに対し、ZeRO-Infinityはファインチューニングに失敗すること、2)小規模なGPT-3 13Bモデルをトレーニングする際、FuyouがRTX 4090 GPU上で156 TFLOPSを達成するのに対し、ZeRO-Infinityは45 TFLOPSしか達成できないことを示しています。

English

Recent advances in large language models have brought immense value to the world, with their superior capabilities stemming from the massive number of parameters they utilize. However, even the GPUs with the highest memory capacities, currently peaking at 80GB, are far from sufficient to accommodate these vast parameters and their associated optimizer states when conducting stochastic gradient descent-based optimization. One approach to hosting such huge models is to aggregate device memory from many GPUs. However, this approach introduces prohibitive costs for most academic researchers, who always have a limited budget for many high-end GPU servers. In this paper, we focus on huge model fine-tuning on a single, even low-end, GPU in a commodity server, which is accessible to most AI researchers. In such a scenario, the state-of-the-art work ZeRO-Infinity suffers from two severe issues when running in a commodity server: 1) low GPU utilization due to inefficient swapping, and 2) limited trainable model size due to CPU memory capacity. The underlying reason is that ZeRO-Infinity is optimized for running on high-end GPU servers. To this end, we present Fuyou, a low-cost training framework that enables efficient 100B huge model fine-tuning on a low-end server with a low-end GPU and limited CPU memory capacity. The key idea is to add the SSD-CPU communication as an optimization dimension and thus carefully co-optimize computation and data swapping from a systematic approach to maximize GPU utilization. The experimental results show that 1) Fuyou is able to fine-tune 175B GPT-3 on a consumer GPU RTX 4090 with high GPU utilization, while ZeRO-Infinity fails to fine-tune; and 2) when training a small GPT-3 13B model, Fuyou achieves 156 TFLOPS on an RTX 4090 GPU while ZeRO-Infinity only achieves 45 TFLOPS.