您现在可以在家训练 70b 语言模型

原文：Answer.AI - You can now train a 70b language model at home

我们正在发布一个基于 FSDP 和 QLoRA 的开源系统，可以在两个 24GB GPU 上训练 70b 模型。

已发表

2024 年 3 月 6 日

概括

今天，我们发布了 Answer.AI 的第一个项目：一个完全开源的系统，首次可以在具有两个或更多标准游戏 GPU（RTX 3090 或 4090）的常规台式计算机上高效地训练 70b 大型语言模型。该系统结合了 FSDP 和 QLoRA，是 Answer.AI、Tim Dettmers（华盛顿大学）以及 Hugging Face 的 Titus von Koeller 和 Sourab Mangrulkar 之间合作的结果。

该系统将帮助开源社区发布更好的模型。 Teknium 是极受欢迎的 OpenHermes 模型和数据集的创建者，下载量超过 50 万次，他表示：

“有了这种能力，我们可以在本地将巨大的模型提升到新的高度，现在小型实验室也可以访问巨大的、数千亿个参数的模型。 ”

在 Answer.AI，我们将此作为我们的第一个项目，因为它是我们北极星的关键基础：帮助每个人都能使用有用的人工智能。仅仅能够使用其他人的模型是不够的。我们希望每个人都能够创建自己的个性化模型，以便他们控制自己的人工智能系统。

背景

大创意

用于训练深度学习模型的硬件有两种截然不同的级别。还有数据中心级的硬件，比如H100s、A100s，价格就几十万美元。然后是包含游戏 GPU 的台式计算机，例如双 4090，成本低于 10,000 美元（并且可以用二手零件组装而成，价格不到预构建系统价格的一半）。

但关键是：游戏 GPU 的性能与数据中心 GPU 相似，但价格高出 10 倍以上！如果我们可以使用这些便宜 10 倍（但速度几乎一样快）的卡来训练大型语言模型，那就太好了，但我们不能，因为它们的内存要少得多。目前最好的数据中心卡具有 80GB RAM，而游戏卡最大可达 24GB RAM。由于只有最大的模型才能产生最好的结果，因此大多数人基本上无法创建最好的模型。

我们意识到这实际上没有内在原因。超快的硬件就在那里，等待使用——我们只需要一种方法来为其提供模型和数据，以满足其内存限制。显而易见的问题是：为什么当时没有这样做？所有大型工业实验室都已经拥有昂贵 10 倍的硬件，因此他们并没有动力去解决这个问题。

这里的主要想法很简单：弄清楚如何使用这些更便宜、内存更低的游戏 GPU 来训练最好的可用开源模型。所以目标是这样的：仅使用游戏 GPU 训练1 个700 亿个参数 (70b) 的模型，这意味着我们的每个 GPU 内存最多为 24GB。这将是一个挑战，因为每个参数通常需要 16 位（2 个字节），因此甚至需要 70*2=140GB 来存储权重 - 而且这还不包括所有其他数据，例如激活、梯度和优化状态！

为什么是这个项目？

答：AI 是一种非常不寻常的组织类型——一个营利性研发实验室，在精神上更接近19 世纪的电力实验室，而不是今天的人工智能研究小组。 Eric Ries 和 Jeremy Howard去年在 NeurIPS 上成立该组织时，希望我们能够找到如何让大型模型训练变得廉价且易于使用的方法。

解决这个问题很难。它需要了解许多独立的库（例如，bitsandbytes、PEFT、Transformers、Accelerate 和 PyTorch）以及计算机科学和数学概念（例如，离散化、分布式计算、GPU 编程、线性代数、SGD 概念，例如梯度检查点），以及它们如何全部相互作用。

学术界充满了解决难题的才华横溢的人。但学术界尚未解决这个特殊问题。这是因为大学研究人员很难证明花时间在此类工作上是合理的。将现有工具和技术结合在一起通常被认为不够“新颖”，不足以在高影响力期刊上发表，但这正是学术界所需要的货币。此外，学者们通常被期望在其领域内变得高度专业化，这使得将如此多的部分整合到一个解决方案中具有挑战性。

当然，大型科技公司也充满了解决难题的优秀人才。但使用消费级 GPU 训练模型这一特殊问题并不是他们需要解决的问题——他们已经购买了昂贵的 GPU！许多初创公司也充满了解决难题的才华横溢的人！但是，正如埃里克·里斯 (Eric Ries) 所解释的那样，“当今的金融市场迫使企业将短期收益置于其他一切之上”。对于初创公司来说，向投资者解释为什么他们将资金用于开源软件和公共研究是极其困难的。

虽然学术界、大型科技公司和初创公司有充分的理由不解决这个问题，但这些正是这个问题非常适合 Answer.AI 的确切原因。在公司工作的每个人都构建了我们必须使用的系统来解决这个问题，因此我们能够了解所有部分如何组合在一起。那些既喜欢深入了解软件和人工智能的基础，又喜欢破解有趣的端到端系统的人，就是被 Answer.AI 所吸引的人，反之亦然。

我们选择共同解决的问题是由负责解决的同一个人选择的。因此，我们倾向于选择将多种想法结合在一起以创建实用的解决方案的项目。因为我们是一家公益公司，我们的使命是从人工智能中产生长期利益，所以开源软件和公共研究与我们的使命直接一致。

QLoRA：在单个 GPU 上训练更大的模型

最近发布的两个项目为实现这一目标迈出了关键的第一步：QLoRA（由Tim Dettmers 等人开发）和 FSDP（由 Meta 的PyTorch 团队开发）。

QLoRA 是现代神经网络中两个至关重要的进步的简单而出色的组合：量化和LoRA。量化是一种不使用 16 位甚至 32 位来存储神经网络权重的技术，而是使用 4 位（甚至更少）。 4 位数字只有 16 个可能的值，但Dettmers 和 Zettlemoyer 表明，这对于当今流行的大型语言模型来说已经足够了。 Tim Dettmers 借助他的 bitsandbytes 库使这些 4 位“量化”模型易于创建，最近 Hugging Face 介入帮助维护和记录该库，特别感谢 Titus von Koeller 的倡议。

不幸的是，一旦模型被量化，就无法使用常规方法对其进行进一步训练——只有 16 个可能的值，用于模型训练的梯度下降方法几乎到处都会观察到零梯度，因此它无法对模型进行任何更新。量化权重。这是一个主要问题，因为这意味着量化只能用于推理，而不能用于持续的预训练或微调。虽然推理很有用且重要，但它实际上只是消耗模型。但我们希望每个人都能为创建模型做出贡献！

避免这种限制的技巧是使用LoRA—— “大型语言模型的低阶适应”。 LoRA 根本不训练整个大型语言模型，而是添加“适配器”，它们是经过训练的非常小的矩阵（通常小于整个模型的 1%），同时保持模型的其余部分不变。如果您玩过“稳定扩散”等模型，您可能会多次看到这些适配器；这就是这些模型通常如何共享的原因，也是它们体积如此之小、下载速度如此之快的原因。

Tim 意识到 LoRA 可以与量化相结合：使用训练时根本不会改变的量化基础模型，并添加未量化的可训练 LoRA 适配器。这个组合就是QLoRA。 Tim 的团队首次能够使用它来训练（未量化）比 GPU 更大的模型：他们在 48GB 卡上训练了 65b 模型（未量化为 130GB）。

Hugging Face 再次介入，创建了PEFT库，使 LoRA 训练变得更加简单，并将其直接与 BitsandBytes 集成，让任何人只需几行代码即可使用 QLoRA。 Hugging Face 团队一直在幕后不知疲倦地工作，以确保开源社区可以使用这些技术来训练他们的模型。如果您曾经使用 Transformers 使用单个函数参数加载 4 位模型，那么您应该感谢他们（即使您没有，您几乎肯定使用过那些构建了 Transformer 的人的工作成果）他们的生态系统模型）。

QLoRA 并没有完全解决我们想要解决的问题，即在 24GB 卡上训练 70b 模型，但它比以前任何事情都更接近。当量化为 4 位（即 0.5 字节）时，70b 模型需要 70/2 = 35 GB，这比我们想要使用的 24GB 游戏 GPU 更大。

QLoRA 还有其他限制。 48GB 卡非常昂贵，训练 65b 模型只适合这样的卡。这可能是一个问题，因为我们还需要存储许多其他内容，包括训练期间模型的激活2、梯度和优化状态。如果加载模型权重后剩余的内存不多，则说明没有足够的工作内存来支持训练。

例如，语言模型的好处之一是我们可以使用它们来“聊天”、理解或分析长文档或对话。为了创建可以处理这样的长序列的模型，我们需要在训练期间向它们展示长序列的示例。训练中使用的最长序列称为“序列长度”。在 48GB 卡上训练 65b QLoRA 模型时，尝试使用除短序列长度之外的任何内容都会导致错误，因为没有足够的内存来存储有关序列的所有信息；几乎所有内存都用于存储模型本身。

此外，如果模型一次只能查看一个序列，那么将需要很长时间才能遍历训练集中的所有数据。因此，我们希望能够一次将几个序列“批处理”在一起。包含的序列数就是“批量大小”。当加载模型权重后 GPU 上剩余的空间非常少时，我们只能使用非常小的批量大小，从而导致训练速度极其缓慢。

FSDP：将训练扩展到多个 GPU

解决单个消费类 GPU 的 RAM 限制问题的一个明显的解决方案是使用多个 GPU！开源社区中一种非常常见的方法是简单地将模型的几层放置在每张卡上。因此，要进行训练，您需要在第一个 GPU 上运行前几层，然后在第二个 GPU 上运行接下来的几层，依此类推。例如，70b (140GB) 模型可以分布在 8 个 24GB GPU 上，每个 GPU 使用 17.5GB。拥抱脸部变形金刚中甚至还有一个方便的设置，device_map=’auto’您可能已经使用过；这就是这实际上在幕后所做的事情。这可以完成工作，但有一个巨大的缺点：一次只有一个 GPU 处于活动状态，因为所有其他 GPU 都在等待“轮到”。这意味着 ⅞ 的计算被浪费了。

分布式数据并行(DDP) 以前是跨多个 GPU 高效训练模型的黄金标准方法。这需要在每个 GPU 上保留完整的模型 - 如果您有一个小模型（例如 2b 模型，需要 4GB RAM），您可以简单地将整个模型分别加载到每个 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美元！）也达不到要求配备 80GB 内存，满足您的需求。但通过 FSDP，四个 H100 GPU 可以组合起来形成总共 320GB RAM。

（请注意，这样一台机器将花费您大约 150,000 美元......）

将 FSDP 和 QLoRA 结合在一起

在 Answer.AI，我们的北极星正在让有用的人工智能变得更容易使用。花费 150,000 美元来创建自己的高品质个性化模型绝对算不上容易获得的！因此，我们开展的第一个项目是让使用带有消费级游戏 GPU 的桌面来高效训练 70b 模型成为可能。我们认为，如果我们可以使用 QLoRA 将模型的大小减少大约 400%（因此 70b 模型将适合 35GB RAM），然后我们使用 FSDP 将其分片到两个或更多 24GB 消费卡上，那么就剩下剩余足够的 RAM 来训练模型。

第一步

Jeremy 和 Tim 在 2023 年底讨论了将 FSDP 和 QLoRA 结合在一起的想法。 Tim 将 Jeremy 与 Titus von Koeller 联系起来，Jeremy 和 Titus 一起尝试、探索、理解和记录两个库合并时出现的问题。

Answer.AI 的 Johno Whitaker 完成了重要的第一步：一个简单的独立测试脚本，它使我们能够更深入地理解问题并测试解决方案。 2024 年初，Answer.AI 的 Benjamin Warner 和 Titus 独立提出了一个关键想法：将量化参数存储在可选择的数据类型中，其中存储数据类型与“计算类型”相同。模型3 .

本杰明在提出这个想法后 24 小时内就完成了这个原型，但随后我们发现了另一个问题：FSDP 没有复制每个分片使用该模型所需的量化信息！这是因为 FSDP 对于将在 GPU 之间同步的数据子集非常固执4。我们意识到，如果我们在每个 GPU 上量化模型，丢失的元数据将在所有 GPU 上保持不变。此外，我们必须将“量化状态”（对参数进行（反）量化所需的信息）从参数移动到各层中，以确保它们在 FSDP 移动分片时不会被删除。

一旦我们解决了这些问题，我们就能够使用 FSDP 使用量化模型成功训练我们的第一批数据！ Benjamin 和 Answer.AI 的 Kerem Turgutlu 能够将其与所需的所有测试和重构打包到bitandbytes 的拉取请求中。我们非常感谢bitsandbytes 项目的维护者，他们在他们的流程中非常积极地指导我们的公关。

任务完成，差不多了

此时，我们再次认为我们很快就会把事情搞定，但我们再次低估了任务的复杂性！我们意识到的第一件事是，实际上仍然不可能加载比单个 GPU 更大的量化模型，因为加载和量化过程本身需要将整个模型放在一个 GPU 上。

Jeremy 花了几周时间仔细研究 Meta 的Llama-Recipes项目，这是他发现的最完整的 FSDP 微调实现，并密切跟踪它如何与 BitsandBytes 以及 Hugging Face 的 PEFT、Transformers 和 Accelerate 项目一起工作，他设法构建了一个最小的独立脚本，手动完成微调模型所需的所有步骤。

本杰明意识到，通过一些调整，可以一次加载和离散化一层，从而避免需要在单个 GPU 上运行整个模型。他还弄清楚了如何防止 PEFT 库将量化状态转移到 CPU。 Kerem 编写了 LoRA 算法的自定义实现，以便它可以适应 Benjamin 的更改。

至此，我们首次能够在双 3090 游戏 GPU 上微调 70b 模型！

为了实现这项工作，我们不仅受益于 FSDP 和 QLoRA，还受益于过去几年学术界和开源社区开发的大量巧妙技术。我们用了：

• 梯度检查点（也称为激活检查点）以避免存储完整梯度，而是在整个模型的多个“检查点”处保存激活，然后根据需要重新运行前向计算步骤来重新计算梯度
• CPU 卸载，在不使用权重时将权重存储在 CPU RAM 中，而不是存储在 GPU 上，从而大大减少了所需的 GPU 内存。这项技术对于使用 H100 GPU 的“GPU 丰富”来说并不是很有用，因为 H100 GPU 具有高度优化的方式来相互传递权重。但对于我们的用例来说，这是绝对必要的，因为游戏 GPU 和主板没有这些系统
• Flash Attention 2使用内存优化的 Cuda 内核有效计算注意力。

让这些与 FSDP 和 QLoRA 一起工作并不总是那么简单。例如，在使用 CPU 卸载时，Benjamin 发现并修复了位和字节的问题。每次“卸载”权重被复制回 GPU 时，它都会自动重新量化，从而有效地将预训练模型转变为随机权重！我们向bitsandbytes发出了一个拉取请求，跟踪哪些参数已经被量化，这样我们就可以避免冗余计算。

完成所有这些工作后，我们非常高兴地看到我们可以使用消费级 GPU 训练大型模型。 Jeremy 在一系列 GPU 硬件上对原始 llama-recipes 进行了详细的基准测试，Kerem 为新项目开发了一个全面的基准测试系统。在比较两者时，我们意识到我们仍然无法使用我们希望的序列长度或批量大小 - 由于某种原因，使用的内存比我们预期的要多。

当我们仔细观察时，发现这根本不是由于我们的 FSDP/QLoRA 集成所致 - 但实际上，即使没有 FSDP，当我们以位和字节为单位增加 seqlen 时，内存使用量也会呈超线性上升，最终导致内存更高比没有量化的用法！事实证明我们并不是第一个发现这个问题的人。我们还没有比特和字节解决方案（但正在研究中），但它确实让我们有了一个令人兴奋的发现……

发现HQQ

我们喜欢与志同道合的人合作，因此当我们看到 Daniel Han 在Unsloth上所做的出色工作时，我们想了解更多信息，看看我们是否能够互相帮助。我们问 Daniel 在这个领域是否还有其他有趣的项目值得我们关注，他向我们推荐了HQQ。

为了解释 HQQ，我们需要先给出一些背景知识……bitsandbytes 完成的 4 位量化使用了一种简单、快速且巧妙的方法，其中每组参数都被归一化到一致的范围，然后每个参数是放置在一个桶中，其中桶断点基于参数呈正态分布的假设。这样做的好处是量化几乎是即时的，但由于真实的模型参数不会完全符合假设的分布，因此准确性可能会受到影响。

其他方法（例如 GPTQ 和更新的 AWQ）则走向不同的方向，即当代表性数据传递给模型时，根据模型的实际行为来优化量化参数。5这些方法往往会产生更准确的模型，每个参数甚至可能少于 4 位；但它们的缺点是每个模型的优化过程可能需要数小时甚至数天的时间。

HQQ 结合了两全其美。与 GPTQ 相比，处理 70b 模型的速度快 50 倍，但精度更高。 Kerem 决定调查 HQQ 是否能与 FSDP 良好配合。

他发现，使 HQQ 和 FSDP 很好地协同工作所需的步骤几乎与位和字节所需的步骤完全相同，因此他在几天内完成了一个完整的工作示例。mobius.ml 人员在确保我们的 PR 成功合并方面反应迅速且乐于助人 - 因此我们现在很高兴地宣布 FSDP 也可以与 HQQ 配合使用！

如何使用 FSDP/QLoRA

要使用 FSDP，您当然需要多个 GPU。如果您无法访问此类系统，您可以从Runpod Community Cloud租用双 3090 盒子，价格约为 0.60 美元/小时。还有许多其他提供商可供选择；cloud-gpus是查看所提供内容的好地方。

您需要安装最新版本的 Transformers、PEFT 和 bitsandbytes（如果您使用的话，还需要安装 HQQ）。然后，克隆我们的存储库并按照其中的自述文件进行操作。运行python train.py --help将显示可用选项。要在两张 24GB 卡上包含的羊驼数据集上训练 llama2-7b，您可以运行：

python train.py --train_type qlora --dataset alpaca --batch_size 8 --gradient_accumulation_steps 2 --output_dir qlora_output --log_to wandb

我们将所需的所有内容都塞进了这个文件中，以便更轻松地查看正在发生的情况并在必要时进行修改。

您应该将此脚本视为 alpha/预览版本。虽然我们已经使用它在一系列硬件上成功训练了各种实用的模型，但现在还处于早期阶段。如果您对测试和调试模型不满意，我们建议推迟几个月，以便社区更全面地测试该方法。

更新：我们很高兴看到Hugging Face 生态系统已经通过对Accelerate、Transformers、TRL和PEFT的更改提供支持。我们的代码也已被纳入 Axolotl 微调库中，并用于训练 Mixtral和其他模型。

第一步

最初，我们计划在本文中展示一些基准测试，以及基于基准测试结果的有关如何最好地利用 FSDP/QLoRA 的指导。然而，我们不得不继续推迟发布这个，因为我们每隔几天就会做出重大改进！因此，我们将在接下来的几周内跟进一篇基准测试和建议文章。

我们在这里展示的只是第一步。我们有很多改进的想法，并且我们确信开源社区会有许多我们还没有想到的其他想法！

我们希望看到这个概念证明，即可以在廉价的游戏 GPU 上扩展资源高效的 QLoRA 训练，将有助于引起更多人关注降低模型训练成本的问题。让人工智能更容易使用，让更多人不仅能够消费，而且能够构建有价值的模型，符合每个人的利益。

脚注

1. 在本文中，“训练”可以指预训练或微调。↩︎

2. 严格来说，我们通常不存储实际的激活，而只是存储使用梯度检查点按需重新计算它们所需的中间部分。↩︎

3. FSDP 仅支持浮点类型，但大多数量化库以整数类型存储量化权重。可选择的存储数据类型解决了这一差异。↩︎

4. FSDP 仅支持同步 PyTorch 参数和缓冲区，而大多数量化库将“量化状态”元数据存储在字典中。↩︎

5. 这非常重要，因为校准数据偏差是使用这些数据相关方法可能面临的另一个主要问题。↩︎