Hinzufügen von NVMe-SSDs zur Aktivierung und Beschleunigung des Feinabstimmens eines 100B-Modells auf einer einzelnen GPU.

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

March 11, 2024

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

cs.AI

Zusammenfassung

Die jüngsten Fortschritte bei großen Sprachmodellen haben der Welt immense Werte gebracht, die aus ihrer überlegenen Leistungsfähigkeit resultieren, die auf der enormen Anzahl von Parametern beruht, die sie nutzen. Allerdings sind selbst die GPUs mit dem höchsten Speicherplatz, der derzeit bei 80 GB liegt, weit davon entfernt, diese umfangreichen Parameter und ihre zugehörigen Optimiererzustände bei der Durchführung einer Optimierung auf Basis stochastischer Gradientenabstiegsverfahren zu beherbergen. Ein Ansatz, um solche riesigen Modelle zu hosten, besteht darin, den Gerätespeicher von vielen GPUs zu aggregieren. Dieser Ansatz bringt jedoch hohe Kosten mit sich, die für die meisten akademischen Forscher unerschwinglich sind, da sie immer über ein begrenztes Budget für viele High-End-GPU-Server verfügen. In diesem Artikel konzentrieren wir uns auf das Feintuning riesiger Modelle auf einer einzelnen, selbst Low-End, GPU in einem handelsüblichen Server, der den meisten KI-Forschern zugänglich ist. In einem solchen Szenario leidet die hochmoderne Arbeit ZeRO-Infinity unter zwei schwerwiegenden Problemen, wenn sie auf einem handelsüblichen Server ausgeführt wird: 1) geringe GPU-Auslastung aufgrund ineffizienten Swappings und 2) begrenzte trainierbare Modellgröße aufgrund des Speicherplatzes der CPU. Der zugrunde liegende Grund dafür ist, dass ZeRO-Infinity für den Betrieb auf High-End-GPU-Servern optimiert ist. Zu diesem Zweck präsentieren wir Fuyou, ein kostengünstiges Schulungsframework, das ein effizientes Feintuning riesiger Modelle von 100B auf einem Low-End-Server mit einer Low-End-GPU und begrenztem CPU-Speicherplatz ermöglicht. Die Schlüsselidee besteht darin, die SSD-CPU-Kommunikation als Optimierungsdimension hinzuzufügen und somit Berechnungen und Datenaustausch aus einer systematischen Perspektive sorgfältig zu optimieren, um die GPU-Auslastung zu maximieren. Die experimentellen Ergebnisse zeigen, dass 1) Fuyou in der Lage ist, 175B GPT-3 auf einer Consumer-GPU RTX 4090 mit hoher GPU-Auslastung feinzutunen, während ZeRO-Infinity daran scheitert; und 2) bei der Schulung eines kleinen GPT-3 13B-Modells erreicht Fuyou 156 TFLOPS auf einer RTX 4090-GPU, während ZeRO-Infinity nur 45 TFLOPS erreicht.

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.