提到训练AI大模型，总能让人想起动辄几百上千块GPU、天价训练费用、只有几家大厂才玩得起，普通AI玩家看着铺天盖地的大模型新闻只能默默流泪~
现在，仅有一块GPU的个人PC也可以训练高达180亿参数GPT；普通的笔记本电脑，也能训练十几亿参数的模型，相比现有主流方案，可提升参数容量十余倍！
如此显著的提升来自Colossal-AI，一个通用AI大模型高效训练系统。最重要的是，它完全开源，仅需极少量修改，即可让现有深度学习项目在单张消费级显卡上使用大得多的模型进行训练，每个人都可以在家训练AI大模型！尤其是大幅度降低了AI大模型微调和推理等下游任务和应用部署的门槛！
Colossal-AI还可将现有项目便捷扩展到大规模计算集群，使用高效并行技术进一步加速。
开源地址：
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，美国科学院、工程院、艺术与科学院三院院士
传送门
项目地址：
https://github.com/hpcaitech/ColossalAI
参考链接：
-
https://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc
