Incorporación de unidades SSD NVMe para habilitar y acelerar el ajuste fino de modelos de 100B en una ú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

Resumen

Los recientes avances en los modelos de lenguaje de gran escala han aportado un valor inmenso al mundo, con capacidades superiores que derivan del enorme número de parámetros que utilizan. Sin embargo, incluso las GPU con las mayores capacidades de memoria, que actualmente alcanzan los 80 GB, están lejos de ser suficientes para alojar estos vastos parámetros y sus estados asociados del optimizador al realizar optimizaciones basadas en descenso de gradiente estocástico. Un enfoque para alojar modelos tan grandes es agregar la memoria de dispositivos de múltiples GPU. No obstante, este enfoque introduce costos prohibitivos para la mayoría de los investigadores académicos, quienes siempre tienen un presupuesto limitado para muchos servidores con GPU de gama alta. En este artículo, nos enfocamos en el ajuste fino de modelos enormes en una sola GPU, incluso de gama baja, en un servidor común, que es accesible para la mayoría de los investigadores en IA. En este escenario, el trabajo de vanguardia ZeRO-Infinity enfrenta dos problemas graves al ejecutarse en un servidor común: 1) baja utilización de la GPU debido a un intercambio ineficiente, y 2) tamaño limitado del modelo entrenable debido a la capacidad de memoria de la CPU. La razón subyacente es que ZeRO-Infinity está optimizado para ejecutarse en servidores con GPU de gama alta. Para abordar esto, presentamos Fuyou, un marco de entrenamiento de bajo costo que permite el ajuste fino eficiente de modelos enormes de 100B en un servidor de gama baja con una GPU de gama baja y capacidad limitada de memoria de CPU. La idea clave es añadir la comunicación SSD-CPU como una dimensión de optimización y, por lo tanto, co-optimizar cuidadosamente el cálculo y el intercambio de datos desde un enfoque sistemático para maximizar la utilización de la GPU. Los resultados experimentales muestran que 1) Fuyou es capaz de ajustar GPT-3 de 175B en una GPU de consumo RTX 4090 con alta utilización de la GPU, mientras que ZeRO-Infinity no logra ajustarlo; y 2) al entrenar un modelo pequeño GPT-3 de 13B, Fuyou alcanza 156 TFLOPS en una GPU RTX 4090, mientras que ZeRO-Infinity solo alcanza 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.