L'ajout de SSD NVMe pour permettre et accélérer le fine-tuning de modèles de 100 milliards de paramètres sur une seule GPU.

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

March 11, 2024

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

cs.AI

Résumé

Les récentes avancées dans les modèles de langage de grande taille ont apporté une immense valeur au monde, leurs capacités supérieures découlant du nombre massif de paramètres qu'ils utilisent. Cependant, même les GPU dotés des plus grandes capacités de mémoire, actuellement limitées à 80 Go, sont loin d'être suffisants pour accueillir ces vastes paramètres et leurs états d'optimisation associés lors de l'optimisation basée sur la descente de gradient stochastique. Une approche pour héberger de tels modèles volumineux consiste à agréger la mémoire des appareils de plusieurs GPU. Cependant, cette approche entraîne des coûts prohibitifs pour la plupart des chercheurs académiques, qui disposent toujours d'un budget limité pour de nombreux serveurs GPU haut de gamme. Dans cet article, nous nous concentrons sur le fine-tuning de modèles volumineux sur un seul GPU, même bas de gamme, dans un serveur standard, accessible à la plupart des chercheurs en IA. Dans un tel scénario, l'état de l'art, représenté par ZeRO-Infinity, souffre de deux problèmes majeurs lors de son exécution sur un serveur standard : 1) une faible utilisation du GPU due à un échange inefficace, et 2) une taille de modèle entraînable limitée en raison de la capacité de mémoire CPU. La raison sous-jacente est que ZeRO-Infinity est optimisé pour fonctionner sur des serveurs GPU haut de gamme. À cette fin, nous présentons Fuyou, un cadre d'entraînement à faible coût qui permet un fine-tuning efficace de modèles de 100 milliards de paramètres sur un serveur bas de gamme avec un GPU bas de gamme et une capacité de mémoire CPU limitée. L'idée clé est d'ajouter la communication SSD-CPU comme dimension d'optimisation et ainsi de co-optimiser soigneusement le calcul et l'échange de données dans une approche systématique pour maximiser l'utilisation du GPU. Les résultats expérimentaux montrent que 1) Fuyou est capable de fine-tuner GPT-3 de 175 milliards de paramètres sur un GPU grand public RTX 4090 avec une utilisation élevée du GPU, alors que ZeRO-Infinity échoue ; et 2) lors de l'entraînement d'un petit modèle GPT-3 de 13 milliards de paramètres, Fuyou atteint 156 TFLOPS sur un GPU RTX 4090, tandis que ZeRO-Infinity n'atteint que 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.