Het toevoegen van NVMe SSD's om het finetunen van 100B-modellen op een enkele GPU mogelijk te maken en te versnellen.

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

Samenvatting

Recente vooruitgang in grote taalmodellen heeft enorme waarde gebracht aan de wereld, met hun superieure capaciteiten die voortkomen uit het enorme aantal parameters dat ze gebruiken. Echter, zelfs de GPU's met de hoogste geheugencapaciteiten, die momenteel een maximum van 80GB bereiken, zijn lang niet voldoende om deze enorme parameters en hun bijbehorende optimizerstatussen te accommoderen bij het uitvoeren van stochastische gradiëntdaling-gebaseerde optimalisatie. Een benadering om dergelijke enorme modellen te hosten, is het samenvoegen van apparaatgeheugen van veel GPU's. Deze benadering brengt echter onoverkomelijke kosten met zich mee voor de meeste academische onderzoekers, die altijd een beperkt budget hebben voor veel high-end GPU-servers. In dit artikel richten we ons op het finetunen van enorme modellen op een enkele, zelfs low-end, GPU in een standaard server, die toegankelijk is voor de meeste AI-onderzoekers. In een dergelijk scenario lijdt het state-of-the-art werk ZeRO-Infinity aan twee ernstige problemen bij het draaien op een standaard server: 1) lage GPU-utilisatie door inefficiënt swappen, en 2) beperkte trainbare modelgrootte door CPU-geheugencapaciteit. De onderliggende reden is dat ZeRO-Infinity is geoptimaliseerd voor het draaien op high-end GPU-servers. Daarom presenteren we Fuyou, een low-cost trainingsframework dat efficiënt finetunen van 100B enorme modellen mogelijk maakt op een low-end server met een low-end GPU en beperkte CPU-geheugencapaciteit. Het kernidee is om de SSD-CPU-communicatie toe te voegen als een optimalisatiedimensie en zo zorgvuldig berekening en data-swapping te co-optimaliseren vanuit een systematische benadering om GPU-utilisatie te maximaliseren. De experimentele resultaten tonen aan dat 1) Fuyou in staat is om 175B GPT-3 te finetunen op een consumenten-GPU RTX 4090 met hoge GPU-utilisatie, terwijl ZeRO-Infinity hierin faalt; en 2) bij het trainen van een klein GPT-3 13B-model, Fuyou 156 TFLOPS bereikt op een RTX 4090 GPU, terwijl ZeRO-Infinity slechts 45 TFLOPS bereikt.

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.