2 天训练出 15 亿参数大模型：国产开源项目力克英伟达 Megatron-LM

1 月 20 日消息，当今 AI 之势，影响纵深发展的矛盾是什么？

一方面，大模型风头正劲，效果惊艳，人人都想试试。但另一方面，硬件基础上动不动就是上万张 GPU 的大规模集群在日夜燃烧，钞能力劝退。

所以如果告诉你，现在只用一半数量的 GPU，也能完成同样的 GPT-3 训练呢？

你会觉得关键钥匙是什么？

不卖关子了。实现如此提升的，是一个名为 Colossal-AI 的 GitHub 开源项目。

而且该项目开源不久，就迅速登上了 Python 方向的热榜世界第一。

GitHub 地址：点击打开

不仅能加速 GPT-3，对于 GPT-2、ViT、BERT 等多种模型，Colossal-AI 的表现也都非常 nice：

比如半小时左右就能预训练一遍 ViT-Base / 32，2 天能训完 15 亿参数 GPT 模型、5 天可训完 83 亿参数 GPT 模型。

与业内主流的 AI 并行系统 -- 英伟达 Megatron-LM 相比，在同样使用 512 块 GPU 训练 GPT-2 模型时，Colossal-AI 的加速比是其 2 倍。而在训练 GPT-3 时，更是可以节省近千万元的训练费用。

此外在训练 GPT-2 时，显存消耗甚至能控制在 Megatron-LM 的十分之一以下。

Colossal-AI 究竟是如何做到的？

老规矩，我们从论文扒起。

高效 6 维并行方法

简单来说，Colossal-AI 就是一个整合了多种并行方法的系统，提供的功能包括多维并行、大规模优化器、自适应任务调度、消除冗余内存等。

首先来看多维并行。

所谓"多维"是指，目前主流的分布式并行方案往往使用多种并行方法。

比如英伟达的 Megatron-LM 使用了 3 种方法：数据并行、流水并行和张量并行。因此这种模式也被称为三维并行。微软的 DeepSpeed 调用 Megatron-LM 作为并行基础。

而 Colossal-AI 能将系统的并行维度，一下子拉升到 6 维 --

在兼容数据并行、流水并行的基础上，基于该项目团队自研的 2 维 / 2.5 维 / 3 维张量并行方法，以及序列并行实现。

其中，高维张量并行正是 Colossal-AI 提升大模型显存利用率和通信效率的关键所在。

其实张量并行并不新奇，只是过去我们常见的张量并行更多都是基于一维的。

它的原理是将模型层内的权重参数按行或列切分到不同的处理器上，利用分块矩阵乘法，将一个运算分布到多个处理器上同时进行。

比如英伟达的 Megatron-LM 就是一个典型的例子。

但这种并行方式存在一定弊端。

比如，每个处理器仍需要存储整个中间激活，使得在处理大模型时会浪费大量显存空间。

另一方面，这种单线方法还会导致每个处理器都需要与其他所有处理器进行通信。

这意味着假设有 100 个 GPU 的话，每个 GPU 都需要与其他 99 个 GPU 通信，每次计算需要通信的次数就高达 9900 次。

但如果将张量并行的维度扩展到 2 维，单次计算量能立刻下降一个量级。

因为每个 GPU 只需与自己同行或同列的 GPU 通信即可。

同样还是 100 个 GPU 的情况，每个 GPU 需要通信的 GPU 个数就能降到 9 个，单次计算仅需 900 次。

实际上在此基础上，Colossal-AI 还包含 2.5 维、3 维张量并行方法，可以进一步降低传输成本。

相较于 2 维并行方法，2.5 维并行方法可提升 1.45 倍效率，3 维方法可提升 1.57 倍。

针对大图片、视频、长文本、长时间医疗监控等数据，Colossal-AI 还使用了序列并行的方法，这种方法能突破原有机器能力限制，直接处理长序列数据。

值得一提的是，Colossal-AI 的 API 接口是可以定制的，这使得它可以便捷添加新的并行维度。

其次，大规模优化器也是 Colossal-AI 的亮点。

上面我们也提到了，在分布式并行系统中会使用多种并行方法，数据并行则是另一种常见方法。

这种方法的原理不难理解，就是把训练数据划分成若干份，让不同的机器运算不同的数据，然后通过一个参数服务器 （Paremeter Server）收集目标数据。

由此可以大幅提升 AI 模型训练过程中的批量大小，加速训练过程。

不过大批量训练有个"通病"，就是会产生泛化误差 （Generalization Gap），导致网络泛化能力下降，进而导致 AI 模型准确度下降。

所以，Colossal-AI 在系统中使用了自研的 LAMB、LARS 等大规模优化器。在保证训练精度的情况下，还将批大小从 512 扩展到 65536。

其中，LARS 优化器是通过逐层调整学习率，来减少因为学习率导致的无法收敛情况。

LAMB 优化器则是在 LARS 的基础上，将逐层调整学习率的思想应用到自适应梯度上。

由此，LAMB 能够很好解决此前 LARS 在 BERT 训练中存在差异的问题，最大批量达到了 64K。

此前，LAMB 优化器曾成功将预训练一遍 BERT 的时间，从原本的三天三夜缩短到一个多小时。

第三方面，Colossal-AI 使用自适应可扩展调度器来高效处理任务。

与现有常见的任务调度器不同，Colossal-AI 不是静态地通过 GPU 个数来判断任务规模，而是根据批大小来动态、自动管理每个任务.

通过演化算法，该任务调度器还能不断优化调度决策，更大程度提升 GPU 利用率。

评估结果表明，与当前最先进的方法相比，该方法在平均 JCT （job completion time）上能够缩短 45.6% 的时间，优于现有的深度学习任务调度算法。

此外，这种自适应可扩展调度器还能通过 NCCL 网络通信实现高效的任务迁移。

最后，消除冗余内存也是加速 AI 训练的一种解决思路。

在这方面，Colossal-AI 使用了 zero redundancy optimizer 技术（简称 ZeRO）。

这种方法主要通过切分优化器状态、梯度、模型参数，使 GPU 仅保存当前计算所需的部分，从而来消除数据并行、模型并行中存在的内存冗余。

尤其是在部署模型推理时，通过 zero offload 可以将模型卸载到 CPU 内存或硬盘，仅使用少量 GPU 资源，即可实现低成本部署前沿 AI 大模型。

综上不难看出，在技术层面 Colossal-AI 的加速效果非常明显。

而在应用层面，Colossal-AI 的设计也顾及了能耗问题和易用性两个维度。

考虑到数据移动会是能耗的主要来源，Colossal-AI 在不增加计算量的情况下尽可能减少数据移动量，以此来降低能耗。

另一方面，作为一个开源给所有人使用的系统，Colossal-AI 的使用门槛不高，即便是没有学习过分布式系统的人也能上手操作。

同时，只需要极少量的代码改动，Colossal-AI 就能将已有的单机代码快速扩展到并行计算集群上。

最新实验结果释出

Talk is cheap，效果如何，还是得把实验结果展开来看。

Colossal-AI 近日释出的最新实验结果表明，这一大规模 AI 训练系统具有通用性，在 GPT-3、GPT-2、ViT、BERT 等流行模型上均有亮眼的加速表现。

注：以下 GPU 均指英伟达 A100。

GPT-3 训练速度提高 10.7%

英伟达的 Megatron-LM 在加速训练 GPT-3 时，至少需要 128 块 GPU 才能启动；而从下表可以看出，使用相同的计算资源，Colossal-AI 可以将每次迭代花费的时间从 43.1 秒降至 38.5 秒。

这也就意味着，Colossal-AI 可以将 GPT-3 的训练速度进一步提高 10.7%。

站在工程的角度，考虑到训练这样的大模型往往需要投入数百万美元，这一提升比例带来的收益不言而喻。

另外，通过系统优化，Colossal-AI 还能在训练速度损失不大（43.1→48.5）的前提下，将 GPU 数量从 128 块减少到 96 块，大幅降低训练成本。

而进一步启用 ZeRO3（零冗余优化器）后，所需 GPU 数量甚至能减少一半 -- 至 64 块。

2 天内可完成 GPT-2 训练

在 GPT-2 的加速训练结果中，可以看到，无论是在 4、16 还是 64 块 GPU 的情况下，与 Megatron-LM 相比，Colossal-AI 占用的显存都显著减少。

也就是说，利用 Colossal-AI，工程师们可以在采用同等数量 GPU 的前提下，训练规模更大的模型，或设置更大的批量大小来加速训练。

从下表结果中还可以看出，随着批量大小的增加，Colossal-AI 的资源利用率会进一步提高，达到 Megatron-LM 速度的 2 倍。

研发团队在 256 块 GPU 上进行了实验，最终用时 82.8 个小时完成了 15 亿参数版 GPT-2 的训练。

据此预估，后续在 512 块 GPU 上进行 GPT-2 预训练，Colossal-AI 能将训练时间加速到 45 小时。

充分兼容多种并行模式

在 BERT 上进行的实验，则体现了 Colossal-AI 作为世界上并行维度最多的 AI 训练系统的优势。

与 Megatron-LM 相比，Colossal-AI 序列并行方法只需要更少的显存，就能够利用更大的批量大小来加速训练。同时，还允许开发者使用更长的序列数据。

Colossal-AI 的序列并行方法还与流水并行方法兼容。当开发者同时使用序列并行和流水并行时，可以进一步节省训练大模型的时间。

另外，在近期的学术热点 ViT 模型上，Colossal-AI 也展现了高维张量并行方法的优势。

在使用 64 张 GPU 的情况下，Colossal-AI 采用 2/2.5 维方式进行张量并行，充分利用更大的批量大小，达到了更快的处理速度。

背后团队：LAMB 优化器作者尤洋领衔

看到这里，是不是觉得 Colossal-AI 确实值得标星关注一发？

实际上，这一国产项目背后的研发团队来头不小。

领衔者，正是 LAMB 优化器的提出者尤洋。

在谷歌实习期间，正是凭借 LAMB，尤洋曾打破 BERT 预训练世界纪录。

据英伟达官方 GitHub 显示，LAMB 比 Adam 优化器快出整整 72 倍。微软的 DeepSpeed 也采用了 LAMB 方法。

说回到尤洋本人，他曾以第一名的成绩保送清华计算机系硕士研究生，后赴加州大学伯克利分校攻读 CS 博士学位。

2020 年博士毕业后，他加入新加坡国立大学计算机系，并于 2021 年 1 月成为校长青年教授（Presidential Young Professor）。

同样是在 2021 年，他还获得了 IEEE-CS 超算杰出新人奖。该奖项每年在全球范围内表彰不超过 3 人，仅授予在博士毕业 5 年之内，已在高性能计算领域做出有影响力的卓越贡献，并且可以为高性能计算的发展做出长期贡献的优秀青年学者。

与此同时，尤洋回国创办潞晨科技 -- 一家主营业务为分布式软件系统、大规模人工智能平台以及企业级云计算解决方案的 AI 初创公司。

其核心团队成员来自加州大学伯克利分校、斯坦福大学、清华大学、北京大学、新加坡国立大学、新加坡南洋理工大学等国内外知名高校，在高性能计算、人工智能、分布式系统方面有十余年的技术积累，并已在国际顶级学术刊物 / 会议上发表论文 30 余篇。

目前，潞晨科技已拿下创新工场和真格基金合投的超千万元种子轮融资。

传送门

有关 Colossal-AI，今天就先介绍到这里。

最后，附上传送门，感兴趣的小伙伴，自行取用~

GitHub 地址：点击打开

参考链接：点击打开