Aggiunta di SSD NVMe per abilitare e accelerare il fine-tuning di modelli da 100B su una singola GPU

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

March 11, 2024

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

cs.AI

Abstract

I recenti progressi nei modelli linguistici di grandi dimensioni hanno portato un immenso valore al mondo, con le loro capacità superiori derivanti dall'enorme numero di parametri che utilizzano. Tuttavia, anche le GPU con le maggiori capacità di memoria, attualmente fino a 80GB, sono ben lontane dall'essere sufficienti per ospitare questi vasti parametri e i loro stati associati dell'ottimizzatore durante l'ottimizzazione basata sulla discesa del gradiente stocastica. Un approccio per ospitare modelli così grandi è quello di aggregare la memoria dei dispositivi da molte GPU. Tuttavia, questo approccio introduce costi proibitivi per la maggior parte dei ricercatori accademici, che hanno sempre un budget limitato per molti server GPU di fascia alta. In questo articolo, ci concentriamo sul fine-tuning di modelli enormi su una singola GPU, anche di fascia bassa, in un server di tipo consumer, accessibile alla maggior parte dei ricercatori di IA. In tale scenario, il lavoro all'avanguardia ZeRO-Infinity soffre di due gravi problemi quando eseguito in un server di tipo consumer: 1) bassa utilizzazione della GPU a causa di uno scambio inefficiente, e 2) dimensioni limitate del modello addestrabile a causa della capacità della memoria della CPU. La ragione sottostante è che ZeRO-Infinity è ottimizzato per l'esecuzione su server GPU di fascia alta. A tal fine, presentiamo Fuyou, un framework di addestramento a basso costo che consente un efficiente fine-tuning di modelli enormi da 100B su un server di fascia bassa con una GPU di fascia bassa e una capacità limitata della memoria della CPU. L'idea chiave è quella di aggiungere la comunicazione SSD-CPU come dimensione di ottimizzazione e quindi co-ottimizzare attentamente il calcolo e lo scambio di dati da un approccio sistematico per massimizzare l'utilizzazione della GPU. I risultati sperimentali mostrano che 1) Fuyou è in grado di eseguire il fine-tuning di GPT-3 da 175B su una GPU consumer RTX 4090 con un'elevata utilizzazione della GPU, mentre ZeRO-Infinity non riesce a eseguire il fine-tuning; e 2) durante l'addestramento di un piccolo modello GPT-3 da 13B, Fuyou raggiunge 156 TFLOPS su una GPU RTX 4090, mentre ZeRO-Infinity raggiunge solo 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.