一块Nvidia3090单挑了一个180亿参数的大模型。国内的开源项目这次疯了。什么？单个GPU可以训练大型模型吗？还是20系列能拿下的那个？？？不开玩笑，事实已经摆在我们面前：RTX20606GB通用游戏本能地训练了一个15亿参数的模型；RTX309024GB主机直接单挑一个180亿参数的大模型；TeslaV10032GB甚至可以拿下240亿参数。与PyTorch和业界主流的DeepSpeed方法相比，参数容量可提升10倍以上。而且这个方法是完全开源的，几行代码就可以搞定，修改量很小。这一波操作真的把大模型的训练门槛砍了一半。老黄不是要赔钱吗？那么，打出如此大名的大佬是谁呢？就是国内的开源项目Colossal-AI。自开源以来，多次霸占GitHub最受欢迎榜单。△开源地址：https://github.com/hpcaitech/ColossalAI主要是加速各种大模型的训练，GPT-2、GPT-3、ViT、BERT等模型都可以搞定。比如ViT-Base/32预训练半小时左右，15亿参数GPT模型2天训练完成，83亿参数GPT模型5天训练完成。同时可以节省GPU。例如训练GPT-3时使用的GPU资源可以只有NVIDIAMegatron-LM的一半。那么这一次，它是如何让单个GPU训练出百亿参数的大型模型呢？深入了解一下原理~GPU+CPU异构显存的高效利用为什么单一消费级显卡很难训练AI大模型？有限的显存是最大的困难。如今的大尺度模特风头正劲，效果不错。谁不想感受一下呢？但是动不动就“CUDAoutofmemory”，实在让人受不了。目前业界主流的方法是微软DeepSpeed提出的ZeRO（ZeroReduencyOptimizer）。它的主要原理是对模型进行拆分，将模型内存平均分配给单个GPU。数据并行度越高，GPU上的内存消耗就越低。这种方法仅使用CPU和GPU内存之间模型数据的静态分区，内存布局对于不同的训练配置也是不变的。这导致两个问题。首先，当GPU或CPU内存不足以满足相应模型数据需求时，即使其他设备上有内存可用，系统也会崩溃。第二，细粒度张量在不同内存空间传输时，通信效率会很低；当模型数据可以提前放在目标计算设备上时，CPU-GPU通信流量就不需要了。目前DeepSpeed已经有很多修改版本，并提出使用电脑硬盘来动态存储模型，但是硬盘的读写速度明显低于内存和显存，而且训练速度还是会变慢。针对这些问题，Colossal-AI采用了高效利用GPU+CPU异构显存的方案。具体来说，它利用深度学习网络训练过程中不断迭代的特点，根据迭代次数将整个训练过程分为预热和形式化两个阶段。在warm-up阶段，监控收集非模型数据内存信息；正式阶段，根据收集到的信息，预留计算设备上下一个算子所需的内存峰值，将部分GPU模型张量移至CPU内存。大体逻辑如下：这里稍微展开一下，模型数据由参数、梯度和优化器状态组成，它们的足迹与模型结构的定义有关。非模型数据由算子生成的中间张量组成，这些中间张量根据训练任务的配置（如批量大小）动态变化。他们经常做的事情就是抢GPU显存。因此，有必要在GPU显存不够用的时候让CPU帮忙，同时避免其他情况下的内存浪费。Colossal-AI高效利用了GPU+CPU的异构内存，就是这样的逻辑。在上面的过程中，其实很难获取到非模型数据的内存使用情况。由于非模型数据的生命周期不受用户管理，现有的深度学习框架没有向用户暴露非模型数据跟踪接口。其次，还需要计算CUDA上下文等非框架开销。Colossal-AI这里的解决方案是在warm-up阶段使用采样来获取非模型数据对CPU和GPU的内存占用。简单来说，这是一个加减运算：非数据模型使用量=两个统计时刻之间的最大系统内存使用量-模型数据内存使用量已知，通过查询管理器可以知道模型数据内存使用量。具体是这样的：将所有的模型数据张量交给内存管理器进行管理，每个张量都标记一个状态信息，包括HOLD、COMPUTE、FREE等。然后根据动态查询到的内存使用情况，不断动态转换张量的状态，调整张量的位置，更高效地利用GPU内存和CPU内存。使用非常有限的硬件最大限度地提高模型容量和平衡训练速度。这对于人工智能的普及、大型模型下游任务的低成本微调等都有着深远的意义。而且最重要的是——加个记忆棒比买个高端显卡划算多了卡片。前不久，Colossal-AI也成功复现了谷歌最新的研究成果PaLM（PathwaysLanguageModel），表现也很不错，而微软DeepSpeed目前还不支持PaLM模型。Colossal-AI还能做什么？如前所述，Colossal-AI可以挑战许多任务，例如加速训练和节省GPU资源。那么它是如何做到的呢？简单来说，Colossal-AI是一个集成多种并行方式的系统，提供包括多维并行、大规模优化器、自适应任务调度、消除冗余内存等功能。目前，基于Colossal-AI的加速程序FastFold，可以将蛋白质结构预测模型AlphaFold的训练时间从原来的11天缩短至仅67小时。而且总成本更低，在长序列推理任务中也能实现9到11.6倍的速度提升。这种方法成功地超越了谷歌和哥伦比亚大学的方法。此外，Colossal-AI只需一半数量的GPU即可训练GPT-3。与NVIDIA的解决方案相比，Colossal-AI只需要一半的计算资源就可以开始训练；如果使用相同的计算资源，速度可以提升11%，可以减少GPT-3训练成本超过100万美元。同时，Colossal-AI也非常重视开源社区的建设，提供中文教程，开放用户社区论坛，并根据大家的反馈不断更新迭代。比如有读者留言说，如果Colossal-AI能跑在普通的消费级显卡上就好了。这不，几个月后，已经安排好了~背后的团队：LAMB优化器作者游阳主导此帖。你觉得Colossal-AI真的值得明星关注吗？事实上，这个国产项目背后的研发团队不小。带头人是LAMB优化器的创始人游阳。以第一名保送清华大学计算机系研究生，后赴加州大学伯克利分校攻读博士学位。在CS中。曾获得IPDPS/ICPP最佳论文、ACM/IEEEGeorgeMichaelHPCFellowship、福布斯30岁以下精英（亚洲2021）、IEEE-CS超级计算杰出新人奖、加州大学伯克利分校EECSLotfiA.Zadeh杰出毕业生奖。在谷歌实习期间，游洋用LAMB方法打破了BERT预训练世界纪录。根据Nvidia官方GitHub，LAMB比Adam优化器快72倍。微软的DeepSpeed也采用了LAMB的方法。2021年，游阳回国创办了绿辰科技，一家主营分布式软件系统、大型人工智能平台、企业级云计算解决方案的人工智能初创公司。团队核心成员来自美国加州大学伯克利分校、哈佛大学、斯坦福大学、芝加哥大学、清华大学、北京大学、新加坡国立大学、南洋理工大学等国内外知名高校新加坡;拥有GoogleBrain、IBM、Intel、Microsoft、NVIDIA等知名厂商工作经验。公司成立后即获得创新工场、真格基金等多家顶级VC机构的种子轮投资。LuchenCSOJamesDemmel教授是加州大学伯克利分校的杰出教授、ACM/IEEE院士，以及美国科学、工程和艺术与科学学院的成员。传送门：https://github.com/hpcaitech/ColossalAI
