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

近年来，大型语言模型的进展为世界带来了巨大的价值，其卓越的能力源于其所利用的大量参数。然而，即使是目前容量最高的GPU，其内存容量也仅达到80GB，远远不足以容纳这些巨大的参数及其相关的优化器状态，进行基于随机梯度下降的优化。一种处理这些庞大模型的方法是聚合来自多个GPU的设备内存。然而，这种方法对于大多数学术研究人员来说成本过高，他们总是有限的预算来购买许多高端GPU服务器。在本文中，我们专注于在单个、甚至是低端的GPU上进行巨型模型微调的工作，这在大多数AI研究人员中是可行的。在这种情况下，最先进的ZeRO-Infinity工作在普通服务器上运行时存在两个严重问题：1）由于交换效率低下，GPU利用率低；2）由于CPU内存容量有限，可训练模型大小受限。根本原因是ZeRO-Infinity针对高端GPU服务器进行了优化。为此，我们提出了Fuyou，这是一个低成本的训练框架，可以在低端服务器上进行高效的100B巨型模型微调，具有低端GPU和有限的CPU内存容量。其关键思想是将SSD-CPU通信作为一个优化维度，从系统化的方法中谨慎地协同计算和数据交换，以最大化GPU利用率。实验结果表明，1）Fuyou能够在消费级GPU RTX 4090上高效地微调175B GPT-3，而ZeRO-Infinity则无法微调；2）在训练小型GPT-3 13B模型时，Fuyou在RTX 4090 GPU上实现了156 TFLOPS，而ZeRO-Infinity仅实现了45 TFLOPS。

本论文旨在解决在单个低端服务器上使用低端GPU进行巨型模型微调时，ZeRO-Infinity存在的低GPU利用率和受限的可训练模型大小的问题。

Fuyou是一个低成本的训练框架，通过将SSD-CPU通信作为优化维度，从系统化的方法中优化计算和数据交换，以最大化GPU利用率。

实验结果表明，Fuyou能够在消费级GPU RTX 4090上高效微调175B GPT-3，而ZeRO-Infinity无法微调；当训练小型GPT-3 13B模型时，Fuyou在RTX 4090 GPU上实现了156 TFLOPS，而ZeRO-Infinity仅实现了45 TFLOPS。

在近期的相关研究中，还有一些关于巨型模型微调的研究，如Megatron、GShard等。