明敏整理自 MEET2023
量子位公众号 QbitAI
AI 模型急速增大与硬件算力缓慢增长的矛盾，已愈加尖锐。
尤其在 AI 落地进入深水区的背景下，如何降本增效成为了行业热点议题。
但具体到实际动作，该怎么做？
在 MEET 2023 智能未来大会现场，潞晨科技创始人、新加坡国立大学校长青年教授尤洋指明了一个方向：未来，我们迫切需要一个可扩展性的高效计算基础设施。
而且，潞晨科技已经用实际行动给出论证，它们推出的 Colossal-AI 系统，在过去一年里迅速成长为开源加速方案中的明星项目，GitHub 上揽星超过 7k。
为了完整体现尤洋的分享及思考，在不改变原意的基础上，量子位对他的演讲内容进行了编辑整理。
关于 MEET 智能未来大会：MEET 大会是由量子位主办的智能科技领域顶级商业峰会，致力于探讨前沿科技技术的落地与行业应用。今年共有数十家主流媒体及直播平台报道直播了 MEET2023 大会，吸引了超过 300 万行业用户线上参会，全网总曝光量累积超过 2000 万。
演讲要点
- AI 落地正面临一个大问题：大模型/最先进 AI 技术训练成本太高了。
- 我们迫切需要一个可扩展性的高效计算基础设施。
- Colossal-AI 由三部分组成：高效内存管理系统、自动N维并行技术、大规模优化技术
- Colossal-AI 积极融入全球生态系统，为 PyTorch、Hugging Face、PyTorch Lighting 提供支持
（以下为尤洋演讲全文）
AI 大模型落地成本太高了
我今天分享的主题为《Colossal-AI：一种全新深度学习系统，面向未来各种大模型应用场景的低成本落地》。
在此先简单介绍一下 Colossal-AI 的团队成员。
我自己是新加坡国立大学校长青年教授，在加州大学伯克利分校获得博士学位，并很荣幸获得了 IEEE Early Career Excellence Award。
另一位主要成员是James Demmel 教授，他是 UC 伯克利前计算机系主任和 EECS 院长，也是 UC 伯克利杰出教授，美国科学院、工程院院士。
还有几位是负责 Colossal-AI 开发推广的朋友（如图所示）。
下面进入正题——
Colossal-AI 到底想解决什么样的问题？
可以先来看看这张图。
横坐标为时间，纵坐标为 AI 模型的参数量。
图中有 3 条线，第一条从 2016-2021 年，模型大小从 200 多万增长到了 1.6 万亿，翻了成千上万倍，量化起来是每 18 个月翻 40 倍。
如果从 2018 年算起就更夸张了，每 18 个月翻了340 倍。
与之形成鲜明对比的是，GPU 显存增长相对有限，每 18 个月只翻了1. 7 倍。
在未来，训练 AI 大模型是一种必然。
但目前如 OpenAI 等顶尖 AI 公司，他们训练 AI 的时候用了成百上千个 GPU 才能完成这件事。
核心在于：
- 第一模型计算量太大
- 第二参数量很大，导致显存需求也很大
即使用模型微调或推理，也需要几千 GB 的内存，这就是今天要面对的问题。
未来，我们迫切需要一个可扩展性的高效计算基础设施。
所以我们就专注这一点，打造了 Colossal-AI 这个系统。
另外，中小型企业的需求也十分重要。比如 GPT-2、Transformer，都是谷歌、微软等科技巨头的产物，中小型企业到底是什么情况？
右边这张图可以说明情况。
横坐标是时间，纵坐标是 AI 模型的参数量。
从这里可以看出，中小型企业使用的模型量也呈指数级上升。
这有一个很有意思的现象，不论企业大小，大家普遍都在把模型做得更大，核心的原因在于模型更大性能可以更好。
如上左图表明，横坐标体现了模型参数不断增大，纵坐标则是性能表现。
总结起来，AI 在落地方面正面临着一个很大的问题，大模型或者是最先进的 AI 技术训练成本太高了。
Stability AI 每年光计算就需要花费约 2000 万美元，这显然是小公司、科研单位难以承受的。
即使只拿大模型去做微调和推理，本地应用的挑战也很大，需要好几千 GB 的内存。
最终，想把大模型部署到生产线上，如果一个企业或者单位从零开始自己做，需要的人力为 70 人左右，而在欧美地区光是养活 70 个人，成本就需要 2000 万美金。
所以说大模型的训练成本提高、和自身复杂性使得整个生态变得更加困难。
何为 Colossal-AI？
由此，我们打造了 Colossal-AI 整个系统。
它主要包括 3 个部分：
- 高效内存管理系统（因为大模型本质上还是太吃内存，无法直接容纳）
- 自动N维并行技术，提升计算效率
- 大规模优化技术提升扩展性
从三方面把 AI 模型训练部署的性能提到最高。
目前来看，目标是希望用户只需要在自己的单机笔记本上写好代码，通过 Colossal-AI 能够无缝地部署到云端或者是超级计算机上。
我来简单谈一下背景。
目前训练大模型大概有 3 种并行方式：数据并行、张量并行以及流水线并行。
- 数据并行是指，比如有 1 万个数据表把它分到 10 台机器上，每台机器获得 1000 个数据。
- 张量并行，是在层内划分数据。
- 流水线并行，是在层与层之间去划分数据。
对于我们而言，之前在数据并行上已经做了很多工作。
比如说，我们设计的LARS、LAMB 方法，先后帮助腾讯、索尼、谷歌等公司把模型训练时间从 1 小时缩短到 1 分钟，我们的 LAMB 方法，也帮助谷歌把 BERT 训练时间从 3 天缩短到 76 分钟。
Colossal-AI 的解决方案，首先是支持了上述主流并行方案，然后我们创新地打造了 2D 张量并行、2.5D 张量并行以及 3D 张量并行，以及提出了数据序列并行，还提供了降低显存消耗的异构内存管理和大规模并行优化，把它们整合起来提供一套自动并行的解决方案。
其实 AI 工程师、研究员不需要理解背后的技术细节，只需要提供模型信息和计算资源，就可以自动地把计算资源能力发挥到最大化，同时完成虚拟模型训练和自动部署，轻松低成本应用 AI 大模型。
在此简单讲一下张量并行。
一维张量并行是最简单的层内划分，二维张量并行也是层内划分，但逻辑上把这些数据移动按二维的方式进行分布。
这样的好处是，通过我们精心的设计算法，可以使得传输数据通信的时候，每一个机器只用和同维度的机器打交道就行了。
原本一维情况下，1 台机器需要和 1 万台机器打交道，现在只需要跟数百个机器打交道。
三维的效果更明显，在 100 万台机器的情况下，每台机器都需要跟其他每一台机器打交道，即 1000000-1 台。但现在只需要 100 万的根号 3 次方，即 100 个机器打交道就行了，进一步降低了它的通信复杂度。
这张图可以展示二维张量并行和三维张量并行以及跟其他并行方式的区别。
与此同时，我们也提出了数据序列并行。
原因在于，未来真实 AI 运行中，往往是长序列数据。这些长数据非常棘手，为什么？
因为本身模型应用已经很大了，如果在蛋白质折叠或者 5G、元宇宙场景，长数据再来增加内存开销的话，会直接导致系统崩溃。
为了减缓系统的崩溃，我们引入了数据序列并行，主要有两个点：
对 MLP 层比较友好。因为它本身就是相互独立的。更麻烦的是要在 Transformer 中算 Attention score 的时候，需要所有 token。
这里有一个问题是，如果用比较简单的方法，每个机器需要跟所有机器打交道，显然得不偿失。
比如有p个 GPU，通信代价就变成了p的平方，100 个 GPU 通信需要每步循环 1 万个消息，不划算。
由此，我们设计了环状 self-attention 分析算法，核心思路是每次从左手邻居获得一个k，通过p-1 次就可以把这件事解决，将通信次数从 1 万降低到 99。
同时也提供自动并行技术。
自动并行的作用，使得用户只需把模型写好，不需要知道并行计算分布式系统下层的各种软件基础设施。
因为这些东西的学习负担太高了，用户只需告诉我们有多少个机器、多少资源、训练多少模型，就能自动把模型划分到它的计算资源上，从而帮助完成训练、部署。
另一个方面是内存优化技术。
众所周知 GPU 内存是非常有限的，假如只用 GPU 内存，比如说用 PyTorch baseline 训练 GPT-3，至少需要几百个 GPU 才能完成。
毕竟每个 GPU 都只有 16/40/80G 内存。
所以我们需要充分利用各个层级的内存。
GPU 内存放不下，就把部分数据迁移到 CPU 上，CPU 放不下再放到 NVMe 硬盘上。
但要知道，数据移动开销是远大于计算的，也就是 GPU 和 CPU 之间的数据传输，远慢于 CPU 和 NVMe 之间的数据移动。所以要如何尽可能多地把数据移出 CPU、GPU，同时能够减少移动量。
这就需要动态的智能系统，所以我们打造了 Colossal-AI 系统有 4 个模块去统计 tensor 信息、确定 tensor 状态，可以实时移动合适的 tensor。
总之，通过这种方式可以大幅提升效率。
这有一个比较，右上角图中红色代表普通 PyTorch。可以看出在 6G 内存的情况下，用 PyTorch 只能训练 1 亿参数，但是用 Colossal-AI 可以训练 15 亿参数，扩大了 15 倍。
相同硬件下，在 GPT-2 上的效果也更好，使用 PyTorch 本来只能训练 1 亿参数，通过 Colossal-AI 可以扩大到 120 亿。
加速效果还可以比 Megatron-LM 更快。
像 Stable Diffusion 这样的新模型上，Colossal-AI 也能取得6. 5 倍加速，并可以把成本降到原本的1/7。
Colossal-AI 也支持推理，比如在一台服务器上可以很好把 1750 亿参数的 OPT 部署起来。
目前，Colossal-AI 在开源后的增长情况还是非常好的，和之前比较成功的开源软件 Grafana、Spark 等相比，至少在前 13 个月 Colossal-AI 的增长速度是远超它们的，目前达到 7k 颗 GitHub Star。
Colossal-AI 的用户也是遍布于全球的，包括中国、美国、欧洲、印度、东南亚等。
与此同时，Colossal-AI 还在积极融入全球生态系统。
比如目前世界上第一大 AI 生态系统是 Facebook 的 PyTorch，它最好的模型是 OPT，在 OPT 的官网上直接指向了 Colossal-AI。
第二大生态系统 Hugging Face 也找我们做了优化。
第三大 AI 生态系统 PyTorch Lighting 依赖于 Colossal-AI，他们的用户能直接应用 Colossal-AI 提升效率。
如下目前开源社区上 Colossal-AI 的用户，都是一些国内外大型企业。
Colossal-AI 还有很多具体行业的落地应用，比如蛋白质折叠，可以大幅提升效率。
Colossal-AI 现已对外开源，代码公开免费，可以自由使用。
如果有问题，可以访问 Colossal-AI 的 GitHub 主页进行提问，我们免费提供技术咨询，也有微信群可以交流讨论。
