△开源地址：https://github.com/hpcaitech/ColossalAI
主要做的事情就是
加速各种大模型训练
，GPT-2、GPT-3、ViT、BERT等模型都能搞定。
比如能半小时左右预训练一遍ViT-Base/32，2天训完15亿参数GPT模型、5天训完83亿参数GPT模型。
同时还能省GPU。
比如训练GPT-3时使用的GPU资源，可以说只是英伟达Megatron-LM的一半。
那么这一回，它又是如何让单块GPU训练百亿参数大模型的呢？
我们深扒了一下原理~
高效利用GPU+CPU异构内存
为什么单张消费级显卡很难训练AI大模型？
显存有限
，是最大的困难。
当今大模型风头正盛、效果又好，谁不想上手感受一把？
但动不动就“CUDA out of memory”，着实让人遭不住。
目前，业界主流方法是微软DeepSpeed提出的
ZeRO
（Zero Reduency Optimizer）。
它的主要原理是将模型切分，把模型内存平均分配到单个GPU上。
数据并行度越高，GPU上的内存消耗越低。
这种方法在CPU和GPU内存之间仅使用
静态划分模型数据
，而且内存布局针对不同的训练配置也是恒定的。
由此会导致两方面问题。
第一
，当GPU或CPU内存不足以满足相应模型数据要求时，即使还有其他设备上有内存可用，系统还是会崩溃。
第二
，细粒度的张量在不同内存空间传输时，通信效率会很低；当可以将模型数据提前放置到目标计算设备上时，CPU-GPU的通信量又是不必要的。
目前已经出现了不少DeepSpeed的魔改版本，提出使用
电脑硬盘
来动态存储模型，但是硬盘的读写速度明显
低于
内存和显存，训练速度依旧会被拖慢。
所以，就需要在GPU显存不够时CPU能来帮忙，与此同时还要避免其他情况下内存浪费。
Colossal-AI高效利用GPU+CPU的异构内存，就是这样的逻辑。
而以上过程中，获取非模型数据的内存使用量其实
非常难
。
因为非模型数据的生存周期并不归用户管理，现有的深度学习框架没有暴露非模型数据的追踪接口给用户。其次，CUDA context等非框架开销也需要统计。
在这里Colossal-AI的解决思路是，在预热阶段用
采样
的方式，获得非模型数据对CPU和GPU的内存的使用情况。
简单来说，这是道加减法运算：
非数据模型使用 ＝ 两个统计时刻之间系统最大内存使用 — 模型数据内存使用
已知，模型数据内存使用可以通过查询管理器得知。
具体来看就是下面酱婶的：
他曾以第一名的成绩保送清华计算机系硕士研究生，后赴加州大学伯克利分校攻读CS博士学位。
拿过IPDPS/ICPP最佳论文、ACM/IEEE George Michael HPC Fellowship、福布斯30岁以下精英（亚洲 2021）、IEEE-CS超算杰出新人奖、UC伯克利EECS Lotfi A. Zadeh优秀毕业生奖。
在谷歌实习期间，凭借LAMB方法，尤洋曾打破BERT预训练世界纪录。
据英伟达官方GitHub显示，LAMB比Adam优化器快出整整72倍。微软的DeepSpeed也采用了LAMB方法。
2021年，尤洋回国创办
潞晨科技
——一家主营业务为分布式软件系统、大规模人工智能平台以及企业级云计算解决方案的AI初创公司。
团队的核心成员均来自美国加州大学伯克利分校、哈佛大学、斯坦福大学、芝加哥大学、清华大学、北京大学、新加坡国立大学、新加坡南洋理工大学等国内外知名高校；拥有Google Brain、IBM、Intel、 Microsoft、NVIDIA等知名厂商工作经历。
公司成立即获得创新工场、真格基金等多家顶尖VC机构种子轮投资。
潞晨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
游戏里的芯片战争
十七年过去，MacBook Pro 从小众走向大众
平板电脑销售数据出炉：完全卖不动，苹果小米也无力回天！
华硕灵耀 13 2023 体验：轻薄之上的一块好屏
教育硬件市场起起伏伏，只有平板电脑留了下来？
从游戏到AI，英伟达今天成了ChatGPT的动力之源
ChatGPT低成本复现流程开源！单张消费级显卡、1.62GB显存可体验
一个时代有一个时代的计算架构
暗黑4：买英伟达显卡送标准版游戏 坏消息：不包括中国大陆
假如开车也需要充会员
美国的芯片制造野心，首次详细披露
自主造芯：256TOPS算力，功耗至35W，首个存算一体智驾芯片交卷