一个令人 "赞叹不已" 的开源项目！

大家好，我是小 G。

提到训练 AI 大模型，总能让人想起动辄几百上千块 GPU、天价训练费用、只有几家大厂才玩得起，普通 AI 玩家看着铺天盖地的大模型新闻只能默默流泪～

现在，仅有一块 GPU 的个人 PC 也可以训练高达 180 亿参数 GPT；普通的笔记本电脑，也能训练十几亿参数的模型，相比现有主流方案，可提升参数容量十余倍！

如此显著的提升来自 GitHub 上的一个国产项目：Colossal-AI，一个通用 AI 大模型高效训练系统。最重要的是，它完全开源，仅需极少量修改，即可让现有深度学习项目在单张消费级显卡上使用大得多的模型进行训练，每个人都可以在家训练 AI 大模型！尤其是大幅度降低了 AI 大模型微调和推理等下游任务和应用部署的门槛！

Colossal-AI 还可将现有项目便捷扩展到大规模计算集群，使用高效并行技术进一步加速。

GitHub 地址：https://github.com/hpcaitech/ColossalAI

巨头角力，争炼 AI 大模型

从 2018 年谷歌提出的 3 亿参数 BERT 起，大模型记录在短短几年时间内被不断刷新，OpenAI 1750 亿参数的 GPT-3，微软和英伟达联手发布的 5300 亿参数 MT-NLG ......

稠密单体模型已达千亿参数，而稀疏混合模型，如谷歌 2021 年发布的 Switch Transformer，更是将参数量推至万亿级别。

然而，从头训练如此之大的模型费用极高，通常需要同时使用数百甚至上千张 NVIDIA A100 等专业高性能计算 GPU，使用专用的 InfiniBand 高速网络互联组成超级计算机集群，单次训练成本可达千万美元。

使用单张消费级显卡训练 AI 大模型

对于在校学生和个人开发者等普通 AI 玩家来说，显然无法承担上述如此高昂的费用，最方便获得的计算资源还是个人电脑中的英伟达 RTX 系列消费级 GPU。

为了解放 AI 生产力，让大模型普惠到更多开发人员，真正实现 "多快好省" 的使用 AI 大模型，Colossal-AI 仅需几行代码，便可实现提升单卡训练模型容量十几倍的提升。

在各类型硬件上，Colossal-AI 的表现均远超原生 PyTorch 和主流的微软 DeepSpeed 方案。

对于大模型的代表 GPT，使用 Colossal-AI 在搭载 RTX 2060 6GB 的普通游戏笔记本上，也足以训练高达 15 亿参数模型；对于搭载 RTX3090 24GB 的个人电脑，更是可以直接训练 180 亿参数的模型；对于 Tesla V100 等专业计算卡，Colossal-AI 也能显示出显著改善。

Colossal-AI 还成功复现了谷歌的最新研究成果 PaLM (Pathways Language Model)，在各类硬件上也体现出卓越的性能提升，而微软 DeepSpeed 尚不支持 PaLM 模型。

关键技术：异构训练再升级

使用单张消费级显卡训练 AI 大模型的最大困难在于显存容量极其有限，严重限制了可容纳的模型参数量。微软 DeepSpeed 的提出 ZeRO-offload 方法，尝试将模型切分，利用更大容量、更低成本的内存。目前已经出现了多个基于 DeepSpeed 的魔改版本。但如下图左边所示，当 GPU 内存不足以满足其相应的模型数据要求时，即使当时 CPU 上仍有可用内存，系统也会崩溃。

不同于基于 DeepSpeed 的各种魔改方案，Colossal-AI 团队从头搭建了如 ZeRO 等核心关键技术，并针对 DeepSpeed 在 CPU 和 GPU 内存之间仅使用静态划分模型数据、对不同训练配置使用固定内存布局等问题做了诸多改进，进一步挖掘高效的 GPU 与 CPU 内存高效协同方案，毕竟内存条可比大显存的高端显卡便宜太多～

Colossal-AI 设计的 Gemini，就像双子星一样，高效管理和利用 GPU 与 CPU 的异构内存，让张量在训练过程中动态分布在 CPU-GPU 的存储空间内，从而让模型训练突破 GPU 的内存墙。

我们利用深度学习网络训练过程的迭代特性，按迭代次数将训练分为 warmup 和 non-warmup 两个阶段。在初期 warmup 阶段，监测内存信息；在 non-warmup 阶段利用已收集的信息来高效移动张量，以达到最小化 CPU-GPU 数据移动的目的。

听起来容易，实现起来却并不简单。非模型的内存使用量其实难以获取，因为非模型数据的生存周期并不归用户管理，现有的深度学习框架没有暴露非模型数据的追踪接口给用户。其次，CUDA context 等非框架开销也需要统计。

Colossal-AI 通过采样方式在 warmup 阶段获得 CPU 和 GPU 内存的使用情况。非模型数据的使用可以通过统计两个时刻之间系统最大内存使用 - 模型内存使用获得。模型的内存使用情况可以通过查询内存管理器得知，如下图黑色实线所示。

而所有模型数据张量则交给内存管理器管理，每个张量标记一个状态信息，包括 HOLD，COMPUTE，FREE 等。并根据动态查询到的内存使用情况，不断动态转换张量状态，调整张量位置，最终实现对 GPU 显存和 CPU 内存的高效利用，实现在硬件极其有限的情况下，最大化模型容量和平衡训练速度，对于 AI 民主化和低成本微调大模型下游任务等意义巨大。

更进一步：便捷高效并行扩展

并行分布式技术是进一步加速模型训练的重要手段，想要以最短时间训练当今世界最大最前沿的 AI 模型，仍离不开高效的分布式并行扩展。针对现有方案并行维度有限、效率不高、通用性差、部署困难、缺乏维护等痛点，Colossal-AI 通过高效多维并行和异构并行等技术，让用户仅需极少量修改，即可高效快速部署 AI 大模型训练。

例如对于 GPT-3 这样的超大 AI 模型，相比英伟达方案，Colossal-AI 仅需一半的计算资源，即可启动训练；若使用相同计算资源，则能提速 11%，可降低 GPT-3 训练成本超百万美元。

对于蛋白质结构预测应用 AlphaFold，基于 Colossal-AI 的加速方案的 FastFold，成功超越谷歌和哥伦比亚大学的方案，将 AlphaFold 训练时间从 11 天减少到 67 小时，且总成本更低，在长序列推理中也实现 9.3～11.6 倍的速度提升。

Colossal-AI 注重开源社区建设，提供中文教程，开放用户社群及论坛，对于用户反馈进行高效交流与迭代更新，不断添加 PaLM、AlphaFold 等前沿应用。

自开源以来，Colossal-AI 已经多次登顶 GitHub 热榜 Python 方向世界第一，与众多已有数万 star 的明星开源项目一起受到海内外关注！

项目团队

潞晨技术团队的核心成员均来自美国加州大学伯克利分校，斯坦福大学，清华大学，北京大学，新加坡国立大学，新加坡南洋理工大学等国内外知名高校；拥有 Google Brain、IBM、Intel、 Microsoft、NVIDIA 等知名厂商工作经历。公司成立即获得创新工场、真格基金等多家顶尖 VC 机构种子轮投资。

潞晨科技创始人尤洋教授：加州大学伯克利分校博士、IPDPS/ICPP 最佳论文、ACM/IEEE George Michael HPC Fellowship、福布斯 30 岁以下精英 (亚洲 2021)、IEEE-CS 超算杰出新人奖、UC 伯克利 EECS Lotfi A. Zadeh 优秀毕业生奖

潞晨 CSO Prof. James Demmel：加州大学伯克利分校杰出教授、ACM/IEEE Fellow，美国科学院、工程院、艺术与科学院三院院士

传送门

GitHub 地址：https://github.com/hpcaitech/ColossalAI

参考链接：

https://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc