今天，我们发布了 Answer.AI 的第一个项目：一个完全开源的系统，首次可以在普通台式电脑上使用两个或更多标准游戏 GPU（RTX 3090 或 4090）高效训练一个 70 亿参数的大型语言模型。这个系统结合了 FSDP 和 QLoRA，是 Answer.AI、Tim Dettmers（华盛顿大学）以及 Hugging Face 的 Titus von Koeller 和 Sourab Mangrulkar 合作的成果。
这个系统将帮助开源社区发布更好的模型。Teknium，极其受欢迎的OpenHermes模型和数据集的创建者，拥有超过50万次下载，表示：
“有了这种能力，我们可以在本地将庞大的模型提升到新的高度，小型实验室现在可以访问庞大的、数百亿参数的模型。”
在 Answer.AI，我们将这视为我们的首要项目，因为这是我们的北极星的关键基础：帮助让有用的人工智能对每个人都可用。仅仅使用 其他 人的模型是不够的。我们希望每个人都能够创建他们 自己 的个性化模型，这样他们就能控制自己的人工智能系统。
训练深度学习模型使用了两种非常不同级别的硬件。一种是数据中心级别的硬件，比如 H100 和 A100，成本高达数十万美元。另一种是包含游戏 GPU 的台式电脑，比如双 4090，成本不到 1 万美元（并且可以从二手零件中组装，成本不到预装系统价格的一半）。
但这里的关键是：游戏显卡的性能与成本超过10倍的数据中心显卡相似！如果我们能够使用这些成本低10倍（但速度几乎一样快）的显卡来训练大型语言模型将会很棒，但我们不能，因为它们的内存要少得多。目前最好的数据中心显卡具有80GB RAM，而游戏显卡的最大内存为24GB RAM。由于只有最大的模型才能产生最佳结果，创建最佳模型对大多数人来说基本上是不可及的。
我们意识到实际上并没有固有的原因。超快速硬件已经准备就绪，只等待被使用 - 我们只需要一种方法以满足其内存限制的方式提供模型和数据。显而易见的问题是：为什么还没有这样做呢？所有大型行业实验室已经拥有了昂贵10倍的硬件，因此他们并没有真正的动力去解决这个问题。
这里的重要想法很简单：找出如何使用这些更便宜、内存更低的游戏 GPU 来训练最好的开源模型。因此，目标是：仅使用游戏 GPU 训练一个 700 亿参数（70b）的模型，这意味着我们每个 GPU 的内存最多为 24GB。这将是一个挑战，因为每个参数通常需要 16 位（2 字节），所以存储权重需要 70*2=140GB，甚至不包括所有其他数据，如激活、梯度和优化状态！
Answer.AI 是一种非常不寻常的组织 - 一家盈利的研发实验室，更接近于19世纪的电力实验室，而不是当今的人工智能研究团体。找出如何使大型模型训练变得廉价且易于获取，正是埃里克·里斯（Eric Ries）和杰里米·霍华德（Jeremy Howard）希望我们在去年在NeurIPS上推出时能够做到的事情。
解决这个问题很困难。需要理解许多独立的库（例如 bitsandbytes、PEFT、Transformers、Accelerate 和 PyTorch），以及计算机科学和数学概念（例如离散化、分布式计算、GPU 编程、线性代数、SGD 概念如梯度检查点），以及它们之间的相互作用。
学术界充满了解决难题的聪明人。但学术界尚未解决这个特定问题。这是因为大学研究人员很难证明在这类工作上花费时间是合理的。将现有工具和技术结合在一起通常不被认为足够“新颖”，无法发表在高影响力期刊上，但这是学者们所需要的货币。此外，学者们通常被期望在自己的领域内变得高度专业化，这使得将这么多要素融合成一个解决方案变得具有挑战性。
当然，大型科技公司也有很多聪明人解决困难问题。但使用消费级GPU训练模型这个特定问题，并不是他们需要解决的问题——他们已经购买了昂贵的大型GPU！许多初创公司也有很多聪明人解决困难问题！但正如Eric Ries解释的，“当今的金融市场迫使企业将短期收益置于一切之上”。初创公司很难向投资者证明他们为什么要将资金用于开源软件和公共研究。
尽管学术界、大科技公司和初创企业没有解决这个问题的充分理由，但这些正是 确切的原因 使得这个问题非常适合 Answer.AI。公司的每个人都曾构建过我们在这个问题上必须处理的那种系统，因此我们能够理解所有部分是如何相互配合的。热爱深入理解软件和人工智能基础的人，同时也喜欢研究有趣且有趣的端到端系统的人，正是被吸引到 Answer.AI 的人，反之亦然。
我们选择共同解决的问题是由将要解决问题的同一群人选定的。因此，我们倾向于选择涉及汇集多种想法以创造实用解决方案的项目。而且，由于我们是一家公益公司，其宗旨是从人工智能、开源软件和公共研究中获得长期利益，因此开源软件和公共研究与我们的使命完全一致。
最近已经发布了两个项目，它们迈出了实现这一目标的第一步：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%），在保持其余模型恒定的同时进行训练。如果你玩过像 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 模型只能勉强放得下这样的显卡。这可能会成为一个问题，因为我们还需要存储许多其他东西，包括训练期间模型的激活2、梯度和优化状态。如果在加载模型权重后剩余的内存不多，那么就没有足够的工作内存来支持训练。
例如，语言模型的一个好处是我们可以使用它们来“聊天”，或理解，或分析长文档或对话。为了制作能够处理这样长序列的模型，我们需要在训练过程中向它们展示长序列的示例。在训练中使用的最长序列称为“序列长度”。尝试使用除短序列长度之外的任何内容在训练 65b QLoRA 模型时会导致错误，因为在 48GB 卡上没有足够的内存来存储有关序列的所有信息；几乎所有内存都用于存储模型本身。
此外，如果模型一次只能查看一个序列，那么要处理训练集中的所有数据将需要很长时间。因此，我们希望能够一次将几个序列“批量”处理。包含的序列数量即为“批量大小”。当在加载模型权重后 GPU 上剩余空间很少时，我们只能使用非常小的批量大小，导致训练速度极慢。
解决单个消费级GPU的RAM限制问题的一个明显方法是使用多个GPU！在开源社区中，一种非常常见的方法是简单地将模型的几层放在每张卡上。因此，为了训练，您可以在第一个GPU上运行前几层，然后在第二个GPU上运行接下来的几层，以此类推。例如，一个70b（140GB）的模型可以分布在8个24GB的GPU上，每个GPU使用17.5GB。在Hugging Face Transformers中甚至有一个方便的设置，device_map='auto'，您可能已经使用过；这实际上是在幕后执行的操作。这样做虽然能完成任务，但有一个巨大的缺点：一次只有一个GPU处于活动状态，其他所有GPU都在等待它们的“轮次”。这意味着计算的7/8都被浪费了。
分布式数据并行（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（70亿参数）的非量化模型需要140GB的RAM（因为每个参数以16位存储，即2字节），但即使是NVIDIA的H100显卡（单卡售价约为40,000美元！）也无法满足需求，因为其只有80GB的RAM。但是使用FSDP，四张H100 GPU可以组合使用，总共达到320GB的RAM。
（要知道，这样一台机器将花费你大约15万美元…）
在 Answer.AI，我们的北极星是让有用的人工智能更易获取。花费 15 万美元来创建自己的高质量个性化模型绝对不算易获取！因此，我们着手进行的第一个项目是使得使用带有消费级游戏 GPU 的台式机能够高效训练一个 70 亿模型成为可能。我们认为，如果我们能够使用 QLoRA 将模型大小减少约 400%（因此 70 亿模型将适应 35GB RAM），然后我们使用 FSDP 将其分片到两个或更多的 24GB 消费级显卡上，那么剩下的 RAM 就足够训练一个模型了。
2023年底，Jeremy和Tim讨论了将FSDP和QLoRA合并的想法。Tim将Jeremy与Titus von Koeller联系在一起，Jeremy和Titus共同努力，探索，理解并记录了将这两个库合并时出现的问题。
Answer.AI的Johno Whitaker提出了一个重要的第一步：一个简单的独立测试脚本，让我们更深入地了解问题，并测试解决方案。一个关键突破发生在2024年初，当Answer.AI的Benjamin Warner和Titus独立提出了一个关键想法：将量化参数存储在可选择的数据类型中，其中存储数据类型与模型的“计算类型”相同3。
Benjamin 在开发这个想法的24小时内就完成了原型设计，但后来我们发现另一个问题：FSDP 没有复制每个碎片所需的量化信息，以便使用模型！这是因为 FSDP 对于它将在 GPU 之间同步的数据子集有着明确的看法。我们意识到，如果我们在每个 GPU 上对模型进行量化，那么缺失的元数据将保持在所有 GPU 上不变。此外，我们不得不将“量化状态”（(de)quantize 参数所需的信息）从参数移动到层中，以确保在 FSDP 移动碎片时不会被移除。
一旦我们解决了这些问题，我们就能够成功地使用FSDP训练我们的第一批数据，使用量化模型！Benjamin和Answer.AI的Kerem Turgutlu能够将所有必要的测试和重构打包到一个pull request中，以供bitsandbytes使用。我们非常感谢bitsandbytes项目的维护者，在引导我们的PR通过他们的流程中非常负责任。
在这一点上，我们再次意识到我们会很快解决问题，但再次低估了任务的复杂性！我们意识到的第一件事是，实际上仍然不可能加载大于单个GPU的量化模型，因为加载和量化过程本身需要整个模型放在一个GPU上。
Jeremy花了几周时间仔细研究Meta的出色Llama-Recipes项目，这是他找到的最好的完整FSDP微调实现，通过密切跟踪它是如何与bitsandbytes、Hugging Face的PEFT、Transformers和Accelerate项目一起工作的，他设法构建了一个最小的独立脚本，手动完成了微调模型所需的所有步骤。
Benjamin 意识到通过一些调整，可以逐层进行加载和离散化，从而避免将整个模型放在单个 GPU 上的必要性。他还找到了阻止 PEFT 库将量化状态移动到 CPU 的方法。Kerem 编写了 LoRA 算法的自定义实现，以便它可以与 Benjamin 的更改配合使用。
有了这个，我们第一次成功地在双 3090 游戏 GPU 上对 70b 模型进行了微调！
为了使这项工作成功，我们不仅受益于FSDP和QLoRA，还受益于学术界和开源社区在过去几年中开发的大量聪明技术。我们使用了：
- 梯度检查点（也称为激活检查点）以避免存储完整梯度，而是在模型中的若干“检查点”处保存激活，然后根据需要重新运行前向计算步骤来重新计算梯度
- CPU卸载将权重存储在CPU RAM中，而不是在GPU上当它们不在使用时，大大减少了所需的GPU内存。这种技术对于使用 H100 GPU 的“GPU富裕”用户并不是非常有用，因为它们具有高度优化的方法来相互传递权重。但对于我们的用例来说，这是绝对必要的，因为游戏GPU和主板没有这些系统
- Flash Attention 2 以使用内存优化的Cuda内核高效计算注意力。
将这些与FSDP和QLoRA一起配合使用并不总是一帆风顺。例如，在使用CPU卸载时，本杰明发现并修复了bitsandbytes中的一个问题。每次将“卸载”的权重复制回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良好协同工作所需的步骤几乎与bitsandbytes所需的步骤完全相同，结果他在几天内完成了一个完整的工作示例。mobius.ml folks在确保我们的PR成功合并方面反应迅速且乐于助人 - 因此，我们很高兴地宣布FSDP也与HQQ兼容！
要使用FSDP，当然需要不止一个GPU。如果您无法访问这样的系统，可以从Runpod社区云租用一台双3090的服务器，每小时约为0.60美元。还有许多其他供应商可供选择；cloud-gpus是一个很好的地方，可以查看提供的服务。
您需要安装最新版本的Transformers、PEFT和bitsandbytes（如果您使用HQQ也需要安装）。然后，克隆我们的存储库，并按照那里的README进行操作。运行python train.py --help将显示可用选项。要在两张24GB卡上对包含的alpaca数据集进行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/预览版本。虽然我们已经成功地使用它在各种硬件上训练了多个实用模型，但现在仍处于早期阶段。如果您不熟悉测试和调试模型，我们建议您等待几个月，让社区更全面地测试这种方法。
最初我们计划在本文中展示一些基准测试数据，并根据基准测试结果提供指导，以便更好地利用FSDP/QLoRA。然而，由于我们每隔几天就会进行重大改进，所以我们不得不一再推迟发布这篇文章！因此，我们将在未来几周内发布一篇关于基准测试和建议的文章。
我们展示的只是第一步。我们有许多改进的想法，相信开源社区会有许多我们甚至没有想到的其他想法！
我们希望通过看到这个概念验证，证明可以在廉价的游戏GPU上扩展资源高效的QLoRA训练，将有助于更多人关注降低模型训练成本的问题。让人工智能更易接触符合所有人的利益，使更多人不仅能够消费，还能够构建有价值的模型。