当涉及训练大型 AI 模型时，人们会想到使用数千个 GPU，昂贵的训练成本，只有少数几家科技巨头能够承受。而 AI 用户，如初创公司或大学的研究人员，只能眼睁睁地看着大型模型的训练，别无他法。
现在，一台只有一个 GPU 的个人电脑就可以训练高达 180 亿参数的 GPT，一台笔记本电脑也可以训练具有超过 10 亿个参数的模型。与现有的主流解决方案相比，参数容量可以增加十倍以上！
这样显著的改进来自于 Colossal-AI，它是一个通用的大型 AI 模型高效训练系统。最重要的是，它完全开源，只需要进行最小的修改即可让现有的深度学习项目在单个消费级图形卡上训练更大的模型，让每个人都可以在家中训练大型 AI 模型！特别是，它使下游任务和应用部署，如大型 AI 模型微调和推理变得更加容易！
通过提供各种受欢迎的高效并行处理，Colossal-AI 还可以帮助其用户轻松地将现有项目部署到大型计算集群中。
自 2018 年谷歌提出具有 3 亿参数的 BERT 模型以来，仅仅几年时间，大型模型的参数记录就多次被更新，如 OpenAI 提出的具有 1750 亿参数的 GPT-3，由微软和 NVIDIA 共同推出的具有 5300 亿参数的 MT-NLG 等等...
密集模型已经实现了超过千亿参数的规模，而稀疏的专家混合模型（MoE），例如谷歌于 2021 年发布的 Switch Transformer，已将参数数量提升到万亿级别。
然而，从头开始训练如此庞大的模型可能会非常昂贵。通常需要同时使用数百甚至数千个专业高性能 GPU，如 NVIDIA A100。如果我们使用专用的 InfiniBand 高速网络来构建超级计算机集群，训练成本甚至可能达到 1000 万美元。
一块消费级 GPU，即可训练大模型#
很明显，像大学生和个人开发者这样的 AI 用户无法承担这样高昂的成本来训练大模型，而在 AI 社区中，这类人最流行的计算资源就是 NVIDIA RTX GPU。
为了提高 AI 生产力，让更多的开发者受益于大模型，并真正实现我们的愿景，使大型 AI 模型的使用 “快速且廉价”，Colossal-AI 只需要几行代码就可以实现模型训练容量的增加十倍。这将使每个人都能够使用单个普通 GPU 来训练大型 AI 模型。
在所有类型的硬件上，Colossal-AI 的表现都比普通的 PyTorch 和主流的分布式解决方案，如微软的 DeepSpeed 更好。
对于大型模型的代表 ——GPT，Colossal-AI 能够在一台搭载 RTX 2060 6GB 的游戏笔记本上训练最多 1.5 亿个参数的 GPT。对于搭载 RTX3090 24GB 的电脑，Colossal-AI 可以帮助用户训练具有 180 亿个参数的 GPT。对于 Tesla V100 等高性能显卡，Colossal-AI 也可以带来显著的改进。
Colossal-AI 还成功地实现了谷歌最近发布的 PaLM（Pathways 语言模型）。在各种硬件上，它也表现出了出色的性能提升，而微软 DeepSpeed 尚未发布其官方的 PaLM 实现。
关键技术：增强异构训练#
使用单个消费级 GPU 来训练大型 AI 模型的最大问题是 GPU 内存容量极其有限，这严重限制了可以容纳的模型参数。微软 DeepSpeed 提出的 ZeRO-offload 方法尝试分割模型并利用 CPU 内存，其容量更大且成本更低。目前，已经有几个基于 DeepSpeed 进行的异构训练的修改版本。但是，如下图左侧所示，当 GPU 内存不足以满足其相应的模型需求时，即使 CPU 上仍有可用内存，系统也会崩溃。
Colossal-AI 与 DeepSpeed 派生的方法不同，它们从头开始构建了自己的核心技术，如 ZeRO，解决了 DeepSpeed 存在的问题。DeepSpeed 只在 CPU 和 GPU 内存之间静态地分割模型数据，并且对不同的训练配置只使用固定的内存布局。Colossal-AI 做了很多改进，提高了 GPU 和 CPU 内存的使用效率。毕竟，CPU 内存比带有大内存的高性能显卡便宜得多。
Colossal-AI 设计的 Gemini 机制有效地管理和利用 GPU 和 CPU 的异构内存，因此在训练过程中，张量会动态地分配到 CPU-GPU 存储空间中，从而突破了 GPU 的内存限制。
我们利用深度学习网络训练过程的迭代性质，将训练分为两个阶段：预热阶段和非预热阶段，根据迭代次数划分。在初始的预热阶段，监视内存信息；在非预热阶段，使用收集到的信息来高效地移动张量，以最小化 CPU-GPU 数据移动。
这听起来很简单，但实现起来却很麻烦。因为非模型数据的生命周期不由用户管理，而现有的深度学习框架也没有向用户公开非模型数据的跟踪接口，因此很难获得非模型数据的内存使用情况。其次，还需要考虑非框架开销，如 CUDA 上下文等。
Colossal-AI 通过在热身阶段进行采样来获取 CPU 和 GPU 内存的使用情况。而非模型数据的使用情况可以通过比较两个时刻之间的系统最大内存使用量和模型内存使用量来获得。通过查询内存管理器，可以知道模型的内存使用情况，如下图所示的黑色实线所示。
这听起来很简单，但实现起来却很麻烦。因为非模型数据的生命周期不由用户管理，而现有的深度学习框架也没有向用户公开非模型数据的跟踪接口，因此很难获得非模型数据的内存使用情况。其次，还需要考虑非框架开销，如 CUDA 上下文等。
Colossal-AI 通过在热身阶段进行采样来获取 CPU 和 GPU 内存的使用情况。而非模型数据的使用情况可以通过比较两个时刻之间的系统最大内存使用量和模型内存使用量来获得。通过查询内存管理器，可以知道模型的内存使用情况，如下图所示的黑色实线所示。
所有模型张量都由内存管理器管理，每个张量都标有状态信息，包括 HOLD、COMPUTE、FREE 等。根据动态查询的内存使用情况，Colossal-AI 改变张量状态，并不断调整张量位置。最终实现了 GPU 和 CPU 内存的高效使用，最大化了模型容量，并在极其有限的硬件情况下平衡了训练速度。这对于 AI 民主化和大型模型的低成本微调的下游任务具有重要意义。
此外，分布式并行技术是进一步加速模型训练的重要方法。针对现有解决方案的痛点，如并行维度有限、效率低、通用性差、部署困难和缺乏维护，Colossal-AI 使用了多维并行和异构并行等技术，使用户可以仅通过对其代码进行少量修改，即可高效快速地部署大型 AI 模型。
例如，对于像 GPT-3 这样的超大型 AI 模型，与 NVIDIA 解决方案相比，Colossal-AI 只需要一半的计算资源即可开始训练；如果使用相同的计算资源，速度可以进一步提高 11％，从而可以将 GPT-3 的训练成本降低超过 100 万美元。
对于用于蛋白质结构预测的 AlphaFold，我们团队引入了基于 Colossal-AI 加速方案的 FastFold。FastFold 成功超越了谷歌和哥伦比亚大学提出的其他方案，将 AlphaFold 的训练时间从 11 天缩短至 67 小时，同时也降低了总成本。此外，在长序列推理方面，我们实现了 9.3~11.6 倍的速度提升。
此外，Colossal-AI 非常重视开源社区建设，提供英文和中文教程，并支持最新的前沿应用，如 PaLM 和 AlphaFold。我们还将定期推出新的创新功能。我们始终欢迎社区的建议和讨论，如果您遇到任何问题，我们将非常乐意为您提供帮助。您可以在这里提出问题或在我们的论坛中创建讨论主题。我们非常感谢您的建议。最近，Colossal-AI 在 Github 的热门项目排行榜中排名第一，而这是在许多拥有 10K 星的项目的背景下实现的。
关于原作者
原作者都是 HPC-AI Tech 的核心成员，来自加州大学伯克利分校、斯坦福大学、清华大学、北京大学、新加坡国立大学、南洋理工大学等知名高校，同时他们也曾在 Google Brain、IBM、英特尔、微软、NVIDIA 等科技巨头工作过。该公司也获得了创新工场、真格基金等多家顶级风险投资机构的种子轮融资。
Prof. Yang You, Founder of HPC-AI Tech
Ph.D., University of California, Berkeley
IPDPS/ICPP Best Paper Author
ACM/IEEE CS George Michael Memorial HPC Fellowship
Forbes 30 Under 30 (Asia 2021)
IEEE-CS Outstanding Newcomer Award in Supercomputing
UC Berkeley EECS Lotfi A. Zadeh Prize
Prof. James Demmel, CSO of HPC-AI Tech
Distinguished Professor, University of California, Berkeley
ACM/IEEE Fellow
Member of the American Academy of Sciences, the Academy of Engineering, and the Academy of Arts and Sciences
Check out the project over here: https://github.com/hpcaitech/ColossalAI
原文地址（本文对原文进行了翻译）：https://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc