浙江大学研究团队如何在消费级GPU上实现对100B模型微调

近年来，大型语言模型的快速发展为世界带来了巨大的价值，其优越性能源自它们所利用的庞大参数数量。然而，即使是目前内存容量最高的GPU，也只有80GB，远远不足以容纳这些庞大的参数及其相关的优化器状态，尤其在进行基于随机梯度下降的优化时。

上周，我们介绍了国外一个团队，他们提出的方法可以在仅具有24GB内存的单一消费级GPU（RTX 4090）上预训练一个70亿参数的LLM。女科学家提出GaLore：为消费级GPU上高效训练LLM铺平道路

这样的技术突破相当惊人，也相当诱人，尤其对缺乏GPU资源的研究人员来说，GaLore为人工智能研究带来巨大的推动力，也打破了资本密集的科研天花板。

中国受美国GPU卡禁制令影响，无法获得NVIDIA和AMD高端大显存GPU卡，对于中国发展自己的大模型是有影响的。中国的AI研究团队只能在LLM的训练方法上寻求突破。今天我们要介绍来自浙江大学的团队发的一篇名为《Simple and Scalable Strategies to Continually Pre-train Large Language Models》的论文。

论文地址：https://arxiv.org/pdf/2403.06504.pdf

痛点

在论文里研究团队提到：一种处理这样庞大模型的方法是从许多GPU中聚合设备内存。特别地，为了微调一个具有1000亿参数的模型，需要32个NVIDIA A100 GPU。然而，这种方法对于大多数学术研究人员来说造成了不可承受的成本，因为他们总是对许多高端GPU服务器有着有限的预算。

可见，这样的痛苦对于大多数AI研究人员来说，如鲠在喉。

如何使用一块GPU进行大模型的微调

于是，他们决定着眼于在一台单个甚至低端的商品服务器上使用一块GPU进行巨大模型的微调，这对大多数AI研究人员来说是可行的。而目前现有的方式是利用异质存储体来训练LLM，例如现行的ZeRO-Infinity技术，将LLM的参数、梯度和优化器状态从GPU内存卸载到CPU内存甚至SSD存储，并在必要时将激活（activation）卸载到主机内存，因而能够以有限的GPU内存对参数庞大的LLM进行微调。

不过，类似ZeRO-Infinity的技术有个问题，就是它们都是针对高阶GPU服务器设计的。高阶服务器上的微调实际上并不需要将激活和优化器状态卸载到SSD（因为GPU内存容量充足），一旦移转到低阶GPU或CPU内存容量有限的环境（例如你家的台式电脑），ZeRO-Infinity就使不上力。再加上由于低效的交换，GPU利用率低，继而导致可训练模型大小受限。

为此，该团队提出了Fuyou——一个低成本的训练框架，可以在低端服务器上的低端GPU和有限的CPU内存容量上实现高效的1000亿巨型模型微调。其关键思想是将SSD-CPU通信作为一个优化维度，并从系统化的角度精心协调计算和数据交换，以最大化GPU利用率。

三项创新

Fuyou克服了上述的瓶颈。他们添加SSD作为优化维度，在低端机器上实现高效的100B巨大模型微调，主要有三大创新：

首先，团队提出了一种与反向传播重叠的同步外核CPU优化器，以最大化GPU利用率。传统的优化器运算阶段，CPU负责更新优化器状态，GPU则完全闲置，为了提高GPU的使用率，Fuyou移除了传统的优化器阶段，让模型训练的优化器与反向传播重叠，而且不影响训练的收敛速度。

其次，团队提出了一种GPU-CPU-SSD完全流水线化的激活交换机制，允许显著更大的模型微调。该技术实现GPU内存、CPU内存和SSD之间的高效数据交换，使低阶机器能够微调LLM，该模型大小受SSD容量限制，而非CPU/GPU内存大小。

第三，团队提出了一种自动激活交换管理，可以自动确定最佳的交换激活数量，以最小化epoch时间。Fuyou使用一个成本模型来预测给定数量的激活交换的epoch时间（数据交换涉及到SSD带宽和GPU PCIe带宽），然后从中选择成本最小的激活交换量。

核心组件

Fuyou包括四个核心组件：

Fuyou Overview

1. 分析阶段：收集Fuyou自动交换管理所需的关键数据。
2. 同步外核CPU优化器：与反向传播重叠，防止GPU在优化期间空闲，同时不影响训练收敛速度。
3. 完全流水线化的激活交换机制：实现了完全流水线化的GPU-CPU-SSD两级激活交换，实现了对更大模型尺寸的微调。
4. 自动激活调度策略：自动确定交换激活的数量，进一步减少时代时间。

通过整合这些组件，Fuyou优化了训练过程，实现了在低端硬件上实现高效率，不仅最大化了GPU利用率，还增强了Fuyou可以微调的模型尺寸，使其成为AI研究人员克服资源限制的宝贵工具。

总结

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

Fuyou与具有不同batch size的基线模型的端到端GPU吞吐量比较

根据团队的说法：Fuyou是一个划时代的训练框架，旨在解决现有方法（如ZeRO-Infinity）在高效微调大型模型方面的局限性。有了Fuyou，甚至是配备普通GPU的低端服务器也能有效地完成这项任务。其关键创新在于将SSD-CPU通信作为一个优化维度，实现计算和数据交换的无缝协调，以最大化GPU利用率。