Adicionando SSDs NVMe para Habilitar e Acelerar o Ajuste Fino de Modelos de 100B em uma Única GPU

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

March 11, 2024

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

cs.AI

Resumo

Os recentes avanços em modelos de linguagem de grande escala trouxeram um valor imenso ao mundo, com suas capacidades superiores decorrentes do enorme número de parâmetros que utilizam. No entanto, mesmo as GPUs com as maiores capacidades de memória, atualmente chegando a 80GB, estão longe de serem suficientes para acomodar esses vastos parâmetros e seus estados associados de otimizador ao realizar a otimização baseada em gradiente descendente estocástico. Uma abordagem para hospedar modelos tão grandes é agregar a memória de dispositivos de várias GPUs. No entanto, essa abordagem introduz custos proibitivos para a maioria dos pesquisadores acadêmicos, que sempre têm um orçamento limitado para muitos servidores de GPU de alta qualidade. Neste artigo, focamos no ajuste fino de modelos enormes em uma única GPU, mesmo de baixo custo, em um servidor comum, que é acessível para a maioria dos pesquisadores de IA. Nesse cenário, o trabalho state-of-the-art ZeRO-Infinity sofre de dois problemas graves ao ser executado em um servidor comum: 1) baixa utilização da GPU devido à troca ineficiente de dados, e 2) tamanho limitado do modelo treinável devido à capacidade de memória da CPU. A razão subjacente é que o ZeRO-Infinity é otimizado para ser executado em servidores de GPU de alta qualidade. Para isso, apresentamos Fuyou, um framework de treinamento de baixo custo que permite o ajuste fino eficiente de modelos enormes de 100B em um servidor de baixo custo com uma GPU de baixo custo e capacidade limitada de memória da CPU. A ideia principal é adicionar a comunicação SSD-CPU como uma dimensão de otimização e, assim, co-otimizar cuidadosamente a computação e a troca de dados a partir de uma abordagem sistemática para maximizar a utilização da GPU. Os resultados experimentais mostram que 1) Fuyou é capaz de ajustar finamente o GPT-3 de 175B em uma GPU de consumo RTX 4090 com alta utilização da GPU, enquanto o ZeRO-Infinity falha em ajustar finamente; e 2) ao treinar um pequeno modelo GPT-3 de 13B, Fuyou alcança 156 TFLOPS em uma GPU RTX 4090, enquanto o ZeRO-Infinity alcança apenas 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.