自从 Meta 入局开源了 LLaMA，大模型圈的创业浪潮似乎就没停下来过，数以万计的初创公司崭露头角，颇有当年个人计算机（PC）行业在 Linux 系统开源化之后的繁荣景象。
就在去年 12 月，曾经创立了 fast.ai 的 Jeremy Howard 和 Eric Ries 创建了一个名为 Answer.AI 的新人工智能研发实验室。该实验室得到了 Decibel VC 的 1000 万美元资金支持，旨在通过专注于利用现有技术创建实用的人工智能应用程序。
Jeremy Howard
3 月 6 日，Answer.AI 终于在自己的官网博客上正式公开了自己的第一个项目：一个完全开源的系统，首次能够高效地在配备两个或多个标准游戏显卡（如老黄家的 RTX 3090 或 4090）的普通台式电脑上训练一个 70b（700 亿参数）的大规模语言模型。目前，Hacker News 对该项目已经讨论得热火朝天。
在 GitHub Trending 上，这一项目今天也冲上了 Jupyter Notebook 的热度第一。
“端侧大模型”是今年相当热门的话题，之前开源了 7b 模型的法国公司 Mistral 就用一个小模型击败了 LLaMA，还得到了微软的青睐；而公众号曾分享过的 1 bit 大模型更是从底层掀翻了大模型牌局。Answer.AI 的这个项目基于 FSDP 和 QLoRA 技术，是 Answer.AI、Tim Dettmers、以及 Hugging Face 团队的 Titus von Koeller 和 Sourab Mangrulkar 共同合作的成果。
接下来，让我们一起来看看 Answer.AI 究竟用了哪些手段实现“在家训练大模型”，一同展望「人人都能创建个性化模型」的时刻还有多久到来。
原文链接：https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html#a-first-step
用游戏显卡训练大模型！
2024 年的今天，训练深度学习模型使用的硬件通常分为两大阵营。
一派是数据中心级别的高端硬件，比如 H100s 和 A100s，售价高达数十万美元；另一派则是配备有游戏显卡的普通台式电脑，例如一对 RTX 4090 显卡，总价不足 1 万美元（甚至可以采用二手部件拼装，俗称「捡垃圾」，成本不到预装整机的一半）。
关键在于，尽管游戏显卡的价格仅为数据中心级显卡的十分之一，但它们的性能却相当接近。设想若能有效利用这类成本低廉的游戏显卡训练大语言模型，无疑将是一个巨大的进步。
然而，现实情况是，顶级的数据中心级显卡拥有 80 GB 的 RAM，而游戏显卡的最大 RAM 则为 24 GB。鉴于只有大型乃至超大型模型才能取得最优结果，迄今为止，构建此类顶尖模型对于大多数个人和小型团队来说都是遥不可及的。
Answer.AI 的团队认识到，理论上存在使用低成本游戏显卡进行大规模模型训练的固有限制。关键是如何找到合适的方法，在满足显卡内存限制的前提下，合理地输入模型和数据。于是，自然产生了一个疑问：为何至今仍未有人实现这一目标呢？
学术界不乏擅长解决难题的杰出人才。但对于解决这个特定问题，学术界尚未找到答案。这是因为大学研究人员很难证明花费时间在这类工作上的合理性。将现有工具和技术组合在一起通常被认为不够新颖，不足以在高影响力期刊上发表论文。
大型科技公司也满载着解决难题的天才人士。但对于使用消费级GPU训练模型这个特定问题，他们并不需要解决——因为各大行业巨头实验室已拥有昂贵得多的硬件设备，因此缺乏动力去攻克这个问题。
创始人 Eric Ries 用一句话解释了这一困境：“今天的金融市场迫使企业优先考虑短期收益而非其他一切”。
简而言之，这活儿只能让初创公司甚至个人开发者去干。
Answer.AI 的核心创新观念很直接：探寻如何利用这些价格亲民、内存相对较小的游戏显卡训练出一流的开源大模型。具体目标设定为：仅用游戏显卡训练一个拥有 700 亿参数（70b）的模型，即训练用的单块 GPU 内存不超过 24 GB。这项任务颇具挑战性，因为每个参数通常占用 16 位（2 字节），因此即便是存储模型权重就需要至少 140 GB 的空间，还未计入诸如激活值、梯度以及优化状态等其它相关数据。
有两项技术发布为 Answer.AI 开辟了新的道路，即 QLoRA（由 Tim Dettmers 等研究人员提出）及 FSDP（来自 Meta 的 PyTorch 团队）。
QLoRA：在单个 GPU 上训练更大的模型
QLoRA 巧妙地结合了现代神经网络领域的两项关键技术突破：量化技术（quantization）和 LoRA。
量化技术是将原本需用 16 位乃至 32 位表示的神经网络权重缩减至仅用 4 位（甚至更少）表示。尽管 4 位数字仅有 16 种可能的取值，但 Tim Dettmers 等人证明对于当前广泛应用的大型语言模型而言，这一点已足够。
此外，Tim Dettmers 开发的 bits and bytes 库使得创建 4 位“量化”模型变得简便易行，同时 Hugging Face 团队也开始参与对该库的维护和文档编写工作。
然而，量化后的模型面临一个重大难题，即无法用常规方法继续训练——由于可供表达的数值有限，梯度下降法在几乎整个权重空间内都无法获得有效的梯度，从而无法对量化权重进行更新。这意味着量化技术目前只能应用于推理阶段，而非模型的预训练或微调阶段。尽管推理应用具有重要意义，但未来的建模不应该是单向的消费行为，而是人人都能参与共创的过程。
解决这一困境的关键在于引入 LoRA 技术，即“Low-Rank Adaptation of Large Language Models”。LoRA 并不直接对整个大型模型进行训练，而是添加体积非常小（通常不到整个模型体积的 1%）的“适配器”，在维持模型主体不变的前提下单独训练这些适配器。像 Stable Diffusion 这样的模型常常会使用此类适配器，这也是为何它们能以较小体积快速下载的原因。
Tim Dettmers 发现量化技术与 LoRA 可以有效结合：采用一个固定的量化基线模型，附加不可量化的、可训练的 LoRA 适配器，这就形成了 QLoRA 技术。借助这项技术，他的团队首次在一块 48 GB GPU 上成功训练了一个（未经量化时）超过 GPU 内存容量的 65b 模型（未量化时大小为 130 GB）。
Hugging Face 团队再度出手，创建了 PEFT 库，大大简化了 LoRA 训练流程，并将其与 bits and bytes 库直接集成，使得用户只需寥寥几行代码就能使用 QLoRA 技术。QLoRA 虽未能完全攻克 Answer.AI 在 24 GB 显卡上训练 70b 模型的目标，但它已经取得了前所未有的进展。若将 70b 模型量化至 4 位（即 0.5 字节），则占用 35 GB 空间，超过了 24 GB 游戏显卡的可用内存。
QLoRA 尚存在其他局限。首先，48 GB 的高端显卡价格不菲，而且仅能勉强放下 65b 模型的训练任务。除此之外，在训练过程中，还需要考虑存储大量额外数据，包括激活值、梯度以及模型的优化状态，这就要求有足够的剩余内存。若加载完模型权重后内存所剩无几，那么模型训练时的运行内存就会捉襟见肘。
举例来说，语言模型的一大优点是可以模拟对话或者分析长篇文本。为了让模型能够处理较长的序列，训练时就需要提供长序列样本。训练中使用的最长序列称为“序列长度”。在 48 GB 显卡上训练 65b QLoRA 模型时，如果选择过长的序列长度，将会因内存不足而导致训练失败，因为大部分内存都被用于存储模型权重，无法承载更多序列信息。
此外，如果模型每次只能处理一个序列，则完成整个训练数据集的迭代周期将极为漫长。为了加速训练，通常我们会采取“批处理”的方式，一次性处理多个序列，这里的“批次大小”指的就是一次处理的序列数目。当 GPU 内存紧巴巴地装入模型权重之后，仅能使用很小的批次大小，进而导致训练速度极度缓慢。
FSDP：在多个 GPU 上扩展训练规模
面对消费级 GPU 内存限制问题，一种直观的解决方案是采用多块 GPU 进行训练。
在开源社区中，常见做法是将模型的部分层分布在每张显卡上。例如，在训练过程中，首先在第一块 GPU 上运行前几层，然后在第二块 GPU 上处理接下来的层，以此类推。如此一来，一个 70b（140 GB）的大规模模型可以分布到 8 块 24 GB GPU上，每块 GPU 占用约 17.5 GB 内存。
因此，需要寻求一种既能分割模型至多个 GPU 上，又能充分利用所有 GPU 进行并行计算的方法。Meta 公司研发的 FSDP 库恰好填补了这一空白。FSDP 通过将大型模型参数跨多个 GPU 进行分片，使得所有 GPU 能够同时得到充分使用。在神经网络某一层于特定 GPU 上进行训练时，会临时复制必要的参数分片至该 GPU 进行计算，计算完成后删除这些临时数据。
举例来说，一个未量化的 70b（700 亿参数）模型需要 140 GB 内存，即使是最顶级的 NVIDIA H100 显卡（售价约为 4 万美元一块，配备 80 GB 内存）也无法满足需求。然而，借助 FSDP 技术，四块 H100 GPU 组合起来能提供总计 320 GB 的内存空间。
值得注意的是，构建这样一台搭载多块 H100 GPU 的机器可能需要花费大约 15 万美元左右。
找到方向后，Answer.AI 的首要任务之一便是探索如何结合 QLoRA 技术，使得普通台式机上的游戏级 GPU 也能高效地训练 70b 模型。
合二为一！
设想一下：通过 QLoRA 将模型大小压缩约 400%，使其适应 35 GB 内存，再利用 FSDP 将其分片分配到两个或更多 24 GB 消费者级别的显卡上，从而腾出足够的内存用于模型训练。
第一步是在 2023 年末，Answer.AI 创始人 Jeremy 和 Tim Dettmers 讨论了将 FSDP 和 QLoRA 相结合的可能性。Tim Dettmers 将 Jeremy 引荐给了 Hugging Face 的 Titus von Koeller，两人共同合作研究、理解和记录了整合这两个库时遇到的问题。
经过原型设计及调试，团队发现并解决了 FSDP 未能正确复制各分片所需的模型量化信息的问题。他们注意到，由于 FSDP 对同步至各 GPU 间的数据子集有严格的规范，必须确保量化状态（包含参数（去）量化所必需的信息）从参数转移至模型层内部，以避免在 FSDP 移动分片时丢失相关信息。
在这些问题得以解决后，研究人员成功运用 FSDP 和量化模型完成了第一批数据的训练，并将这些改进、相应的测试以及重构内容整理成针对 bits and bytes 项目的拉取请求。值得庆幸的是，bits and bytes 项目维护者们迅速响应，协助这些改动顺利通过审查流程。
随着项目接近尾声，研究人员意识到实际情况下仍然难以直接加载大于单个 GPU 容量的大型量化模型，因为模型加载和量化过程本身就需要先将整个模型置于单个 GPU 上。
Jeremy 又花费数周时间仔细研究了 Meta 公司的 Llama-Recipes 项目，该项目是他找到的最佳 FSDP 微调完整实现案例。研究团队调整后发现，可以逐层进行加载和离散化操作，从而避免在整个模型一次性位于单个 GPU 上的需求。与此同时，研究人员还编写了一种定制版 LoRA 算法实现，使之能与以上改进协同工作。
最终，Answer.AI 实现了在双 3090 游戏 GPU 上成功微调 70b 模型的目标。
如何上手？
Answer.AI 的官方博客有更详尽的办法，此处只进行简单介绍。
原文链接：https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html#a-first-step
首先，你需要拥有不止一块 GPU。当然，现在也有不少社区云服务都提供这类机器的租赁。
其次，需要安装最新版本的 Transformers、PEFT 和 bitsandbytes。接着，克隆 Answer.AI 的仓库并遵循其中的 README 文件操作。
仓库链接：https://github.com/AnswerDotAI/fsdp_qlora/tree/main
Answer.AI 的这一项目振奋人心，他们不仅证明了在低成本游戏级 GPU 上扩展资源高效的 QLoRA 训练规模是可行的，还引起了更多人关注降低模型训练成本这一问题。希望未来的 AI 能变得更易于访问，让更多人不仅能够消费模型，也能构建有价值的模型。
4 月 25 ~ 26 日，由 CSDN 和高端 IT 咨询和教育平台 Boolan 联合主办的「全球机器学习技术大会」将在上海环球港凯悦酒店举行，特邀近 50 位技术领袖和行业应用专家，与 1000+ 来自电商、金融、汽车、智能制造、通信、工业互联网、医疗、教育等众多行业的精英参会听众，共同探讨人工智能领域的前沿发展和行业最佳实践。欢迎所有开发者朋友访问官网 http://ml-summit.org、点击「阅读原文」或扫描下方海报中的二维码，进一步了解详情。
