本文探讨了使用消费级 GPU 训练大型语言模型的挑战与解决方案。核心思想是利用 QLoRA 和 FSDP 技术，降低模型训练的硬件成本，使更多人能够参与到模型创建中。QLoRA 通过量化和 LoRA 适配器减少显存占用，FSDP 则将模型分片到多个 GPU 上实现并行训练。
译自：You can now train a 70b language model at home – Answer.AI
作者：Answer.AI
核心思想
训练深度学习模型需要两种截然不同的硬件。一种是数据中心级别的硬件，例如 H100 和 A100，售价高达 数十万美元。另一种是包含游戏 GPU 的台式电脑，例如双 4090，售价 低于 1 万美元（如果从二手零件组装，价格不到预装系统的一半）。
但关键在于：游戏 GPU 的性能与价格高出 10 倍以上的数据中心 GPU 相似！如果我们能使用这些便宜 10 倍（但速度几乎一样快）的显卡来训练大型语言模型，那就太好了，但我们不能，因为它们的内存要小得多。目前最好的数据中心显卡有 80GB 内存，而游戏显卡的内存最多只有 24GB。由于只有最大的模型才能产生最佳结果，因此创建最佳模型在很大程度上是大多数人无法实现的。
我们意识到，实际上并没有内在的原因导致这种情况。超快的硬件就在那里，等待被使用——我们只需要一种方法来以满足其内存约束的方式为其提供模型和数据。显而易见的问题是：为什么还没有人这样做呢？所有大型行业实验室都已经拥有了价格高出 10 倍的硬件，因此他们没有真正的动力来解决这个问题。
这里的核心思想很简单：弄清楚如何使用这些更便宜、内存更低的游戏 GPU 来训练现有的最佳开源模型。因此，目标是：仅使用游戏 GPU 训练一个 700 亿参数 (70b) 的模型，这意味着我们的每个 GPU 内存最多为 24GB。这将是一个挑战，因为每个参数通常占用 16 位（2 字节），因此即使仅存储权重也需要 70*2=140GB——这还不包括所有其他数据，例如激活、梯度和优化状态！
为什么要做这个项目？
Answer.AI 是一种非常不同寻常的组织——一个以盈利为目的的研发实验室，其精神更接近 19 世纪的电力实验室，而不是今天的 AI 研究小组。弄清楚如何使大型模型训练变得廉价且易于访问正是 Eric Ries 和 Jeremy Howard 在组织 于 NeurIPS 上启动 时希望我们能够做到的事情。
解决这个问题很难。它需要理解许多独立的库（例如 bitsandbytes、PEFT、Transformers、Accelerate 和 PyTorch），以及计算机科学和数学概念（例如离散化、分布式计算、GPU 编程、线性代数、梯度检查点等 SGD 概念），以及它们如何交互。
学术界有很多聪明人，他们解决了难题。但是学术界并没有解决这个特殊的问题。那是因为大学研究人员很难证明花时间从事这类工作是合理的。将现有工具和技术结合在一起通常不被认为是“新颖的”，不足以在高影响力的期刊上发表，但这是学者们需要的。此外，通常期望学者在其领域内变得高度专业化，这使得将如此多的部分整合到单个解决方案中具有挑战性。
而且，当然，大型科技公司也到处都是解决难题的聪明人。但是这个特殊的问题，即使用消费级 GPU 训练模型，并不是他们需要解决的问题——他们已经购买了昂贵的 GPU！许多初创公司也到处都是解决难题的聪明人！但是，正如 Eric Ries 解释的那样，“今天的金融市场迫使企业优先考虑短期收益，而不是其他一切”。初创公司很难向投资者证明他们为什么要将资金花费在开源软件和公共研究上。
虽然学术界、大型科技公司和初创公司都有充分的理由不解决这个问题，但这些 正是 这个问题非常适合 Answer.AI 的原因。在该公司工作的每个人都构建了我们必须在此问题上使用的那种系统，因此我们能够理解所有部分如何组合在一起。喜欢深入理解软件和 AI 基础知识，并且喜欢破解有趣且引人入胜的端到端系统的人是 Answer.AI 所吸引的人，反之亦然。
我们选择一起解决的问题是由将进行解决的同一批人选择的。因此，我们倾向于选择涉及将多个想法结合在一起以创建实用解决方案的项目。而且，由于我们是一家公益公司，其章程规定从 AI 中产生 长期 利益，因此开源软件和公共研究直接符合我们的使命。
QLoRA：在单个 GPU 上训练更大的模型
最近发布了两个项目，它们朝着实现这一目标迈出了关键的第一步：QLoRA（由 Tim Dettmers 等人）和 FSDP（由 Meta 的 PyTorch 团队）。
QLoRA 是现代神经网络中两个至关重要的进步的简单而辉煌的结合：量化 和 LoRA。量化是一种技术，其中使用 4 位（甚至更少）来存储神经网络的权重，而不是使用 16 位甚至 32 位。4 位数字只有 16 个可能的值，但是 Dettmers 和 Zettlemoyer 表明，这对于当今流行的大型语言模型来说可能足够了。得益于他的 bitsandbytes 库，Tim Dettmers 使创建这些 4 位“量化”模型变得容易，并且最近 Hugging Face 介入以帮助 维护和记录 该库，特别是由于 Titus von Koeller 的倡议。
不幸的是，一旦模型被量化，就不能再使用常规方法进行训练——只有 16 个可能的值，用于模型训练的梯度下降方法几乎会在所有地方观察到零梯度，因此它无法对量化的权重进行任何更新。这是一个主要问题，因为它意味着量化只能用于推理，而不能用于持续的预训练或微调。虽然推理有用且重要，但它实际上只是 消耗 模型。但是我们希望每个人都能够 贡献 于 创建 模型！
避免此限制的诀窍是使用 LoRA——“大型语言模型的低秩自适应”。LoRA 根本不训练整个大型语言模型，而是添加“适配器”，这些适配器是非常小的矩阵（通常小于完整模型的 1%），它们经过训练，同时保持模型的其余部分不变。如果您玩过像 Stable Diffusion 这样的模型，您可能已经多次看到这些适配器；这就是这些模型通常共享的方式，以及为什么它们如此小且下载速度快的原因。
Tim 意识到 LoRA 可以与量化结合使用：使用量化的基本模型，该模型在训练中根本不会改变，并添加未量化的可训练 LoRA 适配器。这种组合就是 QLoRA。Tim 的团队能够使用它来首次训练一个（未量化的）大于 GPU 的模型：他们在 48GB 的显卡上训练了一个 65b 的模型（未量化时为 130GB）。
Hugging Face 再次介入，创建了 PEFT 库，这使得 LoRA 训练更加简单，并且还将其直接与 bitsandbytes 集成，从而使任何人都可以仅使用几行代码来使用 QLoRA。Hugging Face 团队一直在幕后不知疲倦地工作，以确保开源社区可以使用这些技术来训练他们的模型。如果您曾经使用 Transformers 通过单个函数参数加载 4 位模型，那么您应该感谢他们（即使您没有这样做，您也几乎肯定使用了使用此生态系统构建其模型的那些人的工作）。
QLoRA 并没有完全解决我们提出的问题，即在 24GB 显卡上训练一个 70b 的模型，但它比以前任何时候都更接近。当量化为 4 位（即 0.5 字节）时，70b 模型占用 70/2 = 35 GB，这大于我们要使用的 24GB 游戏 GPU。
QLoRA 还有其他限制。48GB 显卡非常昂贵，并且在这样的显卡上训练 65b 模型只能勉强容纳。这可能是一个问题，因为我们还需要存储许多其他东西，包括模型在训练期间的激活、梯度和优化状态。如果在加载模型权重后剩余的内存不多，则没有足够的可用内存来支持训练。
例如，语言模型的好处之一是我们可以使用它们来“聊天”、理解或分析长文档或对话。为了使模型能够处理像这样的长序列，我们需要在训练期间向它们展示长序列的示例。训练中使用的最长序列称为“序列长度”。在 48GB 显卡上训练 65b QLoRA 模型时，尝试使用短序列长度以外的任何其他长度都会导致错误，因为没有足够的内存来存储有关该序列的所有信息；几乎所有内存都仅用于存储模型本身。
此外，如果模型一次只能查看一个序列，那么遍历训练集中的所有数据将需要很长时间。因此，我们希望能够一次“批处理”几个序列。包含的序列数是“批处理大小”。如果在加载模型权重后 GPU 上剩余的空间很少，我们只能使用非常小的批处理大小，从而导致训练速度极慢。
FSDP：将训练扩展到多个 GPU
解决单个消费级 GPU 的 RAM 限制的一个显而易见的解决方案是使用多个 GPU！开源社区中一种非常常见的方法是将模型的几层简单地放置在每个显卡上。然后，为了进行训练，您首先在第一个 GPU 上运行前几层，然后在第二个 GPU 上运行接下来的几层，依此类推。例如，可以使用 8 个 24GB GPU 来分配一个 70b（140GB）的模型，每个 GPU 上使用 17.5GB。Hugging Face Transformers 甚至还有一个方便的设置 device_map='auto'，您可能已经使用过；这实际上是幕后正在做的事情。这可以完成任务，但存在一个巨大的缺点：一次只有一个 GPU 处于活动状态，因为所有其他 GPU 都在等待它们的“轮到”。这意味着 ⅞ 的计算被浪费了。
分布式数据并行 (DDP) 以前是在多个 GPU 上高效训练模型的黄金标准方法。这需要在每个 GPU 上保留完整的模型——如果您有一个小型模型（例如，一个占用 4GB RAM 的 2b 模型），您可以简单地将整个模型分别加载到每个 GPU 上，并让每个 GPU 并行地处理训练示例。因此，例如，如果您有 4 个 GPU，那么训练速度将提高 4 倍。但是，如果模型无法容纳到 GPU 上，并且没有足够的剩余空间来存储训练过程所需的数据，则 DDP 将不起作用。
因此，我们需要一些可以将模型拆分到多个 GPU 上（如 device_map='auto'），并且还可以并行使用它们（如 DPP）的东西。这就是 Meta 的 完全分片数据并行 (FSDP) 库的用武之地。它通过将大型模型的参数拆分到多个 GPU 上来“分片”该模型，从而允许同时使用所有 GPU。当神经网络的某一层在训练期间在特定的 GPU 上计算时，所有必需的分片都会被复制到那里。然后，进行计算，最后从该 GPU 中删除复制的部分。虽然这听起来效率极低，但实际上，通过在当前层忙于计算的同时智能地复制下一层的数据，这种方法可能会导致与 DDP 相比没有减速。
FSDP 将 DDP 的性能带到大于任何一个 GPU 的模型的能力是一个启示。例如，一个 70b（700 亿参数）的未量化模型占用 140GB 的 RAM（因为每个参数都存储为 16 位，即 2 字节），但即使是 NVIDIA 的 H100 显卡（单张卡的价格约为 40,000 美元！）也无法满足需求，其 RAM 为 80GB。但是，使用 FSDP，可以将四个 H100 GPU 组合起来，总共获得 320GB 的 RAM。
（请注意，这样一台机器的价格约为 150,000 美元……）
