新加坡国立大学尤洋：高性能AI如何突破？

来源：雷锋网

作者：黄楠

过去数年，AI 模型的参数发生了极大变化。尤洋指出，从 2016 年至 2021 年 1 月，AI 模型的参数量是每 18 个月增长 40 倍；从 2018 年 1 月到 2021 年 1 月，AI 大语言模型的参数量每 18 个月增长 340 倍。而相形之下，2016 年 1 月至 2021 年 1 月间，GPU 的内存增长每 18 个月仅有 1.7 倍。

由此可见，训练成本高、周期长，是当前大模型发展最需要克服的难题。

针对这一问题，尤洋提出了 Colossal-AI 系统，从高效内存系统、N 维并行系统和大规模优化三个层次出发，以实现同样的设备条件下将数据移动的最小化，将 GPU 的吞吐量扩大至最高点。

尤洋还指出，现阶段的模型参数量以 10 万倍扩大、但层数增加不多，这或意味着：如今的 AI 发展可能不再是深度学习、而是进入了宽度学习时代。在模型变得更宽的情况下，面对大规模、长时间的 GPU 训练任务，大模型训练系统的核心将是如何实现 GPU 并行计算，以实现大模型训练越快越省钱的目标。

以下为尤洋的现场演讲内容，雷峰网作了不改变原意的编辑及整理：

AI 大模型的机遇和挑战

首先展示一张图片。图上的横坐标是时间，纵坐标是 AI 模型的参数量。

从 2016 年至 2021 年 1 月，AI 大模型的参数量大概每 18 个月增长 40 倍；从 2018 年 1 月到 2021 年 1 月，AI 模型的参数量每 18 个月增长 340 倍。

2016 年，当时世界上最好的模型是 ResNet-50，而今天最好的模型是 GPT-4。从架构上来看，虽然 OpenAI 没有对外公布 GPT-4 的架构，但对比 ResNet-50 的 50 层神经网络和 GPT-3 未达 100 层的架构，可以说 AI 模型近年来的层数并没有产生的太大的变化。

从 ResNet-50 到 GPT-4，虽然参数量大了 10 万倍左右，其实是每一层都变得更加宽了。包括 LLaMA-65B 版本，也是几十层的网络。

所以我们可能不是深度学习，而是进入了一个宽度学习时代。

可以看到，自 2019 年开始， Transformer 架构基本统一了 AI 大模型赛道，当前最高效的 AI 大模型均是 Transformer 架构。上图中的两条虚线，既展示了大模型参数的变化趋势，实际上也展现了 GPU 的变化趋势。

虽然现在英伟达的股价涨了很多倍，但包括英伟达在内的厂商，其 GPU 内存的增长速度远跟不上大模型的发展速度。

相较于过去六年模型参数量的增长速度变化，2016 年 1 月至 2021 年 1 月间，英伟达 GPU 的计算增长速度每 18 个月仅增长了 1.7 倍。

以 A100 80G 为例计算 GPT-3 训练所需的内存数，GPT-3 有大概 1750 亿参数，为方便计算取整数 2000 亿，等于 200 乘以 10 的 9 次方，每个单精度占用 4 个字节，仅参数就要占 800G 内存，梯度也占了 800G 内存。按照当前的优化方法储存一阶矩（first moment）、二阶矩（second moment）等信息均为 800G。也就是说，如果训练一个什么事情都不干大模型，至少需要几 T 的内存，单个 A100 GPU 仅 80G 内存远远不够，加上中间结果的 batch size 越大，内存开销也越大。

这也是为什么从内存角度上看，训练大模型首先需要有成千上万个 GPU 的原因。

举一个可量化的案例。PaLM 是一个 5400 亿的大模型，根据当前云计算市场价，训练 PaLM 需要承包至少 1000 个 GPU，成本约 900 多万美金。而 Meta 此前曾明确提到，LLaMA 需要使用到 2000 个 A100 GPU，并且用三周时间才可完成一次训练，由此计算可得出 LLaMA 单次训练的成本在 500 万美金。

但由于炼大模型并不仅限于一次训练，可能一个好的大模型产品迭代至少需要五六次，前期都是在试错。因此，据公开渠道分析， GPT-4 单次训练成本在 6000 万美金左右，且训练一次需要至少几个月时间。这也是为什么目前 ChatGPT 即便将其升级至最新版本，其底层还是 2021 年 9 月版本的模型。也就是说，从 2021 年 9 月至今，OpenAI 实质上并没有升级其产品，根本原因就在于，模型的每次训练不仅成本很高，训练周期也很长，因此大模型训练的代价高就非常严重。

设想一个问题，今天有一个 1 万亿参数以及 1000 亿参数的大模型，我们能否用什么方法，去检测万亿参数大模型比千亿参数大模型二者哪个效果更好？也即是目前常说的，模型参数量增大、到底要增大到什么程度？

到今天而言，我觉得这个问题暂时无法用科学回答。原因有几个。

首先，训练神经网络存在非凸优化的问题，目前训练所收敛的点多为局部最优解、而非全局最优解。因此，我们要验证神经网络训练到什么程度，在现有计算资源情况下是无法验证的。

第二个难度在于，大模型训练往往只训练一两个 epoch，而此前的 CNN 模型中，ResNet 训练有 90 个 epoch，甚至自监督学习训练有 1000 个 epoch，因此大模型只训练一两个 epoch 的方式，相当于只将数据集过了一两遍，收敛就更加不充分了。因此在训练成本如此之高的情况下，我们很难验证，一个 1 万亿参数的模型和 2 万亿参数的模型二者谁更好，因为它们潜力都没有能通过实验得到充分发挥。因此我认为，今天 AI 大模型是一个实验性学科，如何能高效提升这个实验的效率，降低成本，对整个行业的普及具有根本性的作用。

回到现实之中，为什么今天人人都在追求大模型？从数学逻辑上看，模型参数越大、效果越好，这是绝对的。

与此同时，成本也再不断攀高。目前训练大模型需要成百上千、甚至上万个 GPU，如何将上万个 GPU 的成本进一步降低，挑战非常大的。

在 20 年前，由于当时依靠的是主频的，所有的程序都是串行的，假设将硬件速度提高 10 倍，在一行代码都不用的更改的情况下，其运行速度也可以提升 10 倍。但到了如今，如果想将代码速度提升 10 倍，假定硬件已经增速 10 倍，但如果不优化代码，很可能速度反而会变慢。原因就在于，机器规模更大的情况下，比如 GPU 内部，GPU 内存和 CPU 之间的数据移动，或是 GPU 间的数据移动，再加上服务器实现的数据移动，会占据整个系统的大部分时间，把大部分时间都花在了数据移动上，模型的扩展性也会变得不好。

我认为，未来一个好的分布式软件和一个差的分布式软件，在上千个 GPU 上，甚至 500 个 GPU 上，其速度可能相差 10 倍。

Colossal-AI 如何运行？

基于上述的挑战，我们提出了大模型训练系统 Colossal-AI，提供优化方法，降低数据移动的代价，将模型扩展性效率提到最高。

一个具体的数据是，使用最简单的 PyTorch 训练 GPT-3，成本为 1000 万美金，英伟达经过优化后，用 Megatron 可将其成本减少至 300 万美金，而使用 Colossal-AI 后，成本可以降低到 130 万美金。可以看到，相同的设备条件下，数据移动的最小化将数据移动占比降低最低，能够把 GPU 吞吐量扩大至最高点。

针对上述问题，Colossal-AI 提出了三个层次。其他类似的软件也包括了这三个层次。

第一层是优化内存，先确保单个 GPU、单个服务器的内存效率最高，这是基础。

第二层是 N 维的并行。当前我们使用上千、上万个 GPU 时，其核心技术就是 ParallelComputing（并行计算）。从 1 个 GPU 扩到 10 个 GPU，因为其规模比较小，我们可以轻易获得 7 倍加速；从 10 个 到 100 个 GPU 时，往往可能只获得 4 倍加速，因为并行规模变大，它的通信代价变高了。而从 100 个 GPU 到 1000 个 GPU，因为通信代价进一步加高，很可能只获得 2 倍的加速。而从 1000 个 GPU 到 1 万个 GPU，如果软件运行情况不佳时，不仅可能无法加速，甚至还会更慢，因为设备将所有时间花耗在了更高密度的通信上。

其次是优化问题，未来 AI 大模型的发展方向我认为有两层，第一层是模型变得更加智能，设计出更好的结构，比如说从 BERT 到 GPT，或者从 ResNet 到 BERT等，都是在不断地尝试改变模型结构。

此外还有优化方法的进步，从 SGD 过渡到 MOMENTUM、ADAGRAD，到现在有 ADAM，未来又会有哪些更好的优化方法能够将效率提升 10 倍，这一点也非常重要。

具体到实际操作训练大模型的并行问题。

首先是数据并行，这是最简单、也是最高效的并行方法。数据并行指的是，假设现有 1 万张图片，每次循环处理 1000 张图片，如果有 10 个机器，每个机器分配 100 张，10 个循环即可完成所有图片的处理。

在数据并行的过程中需要进行汇总，每个机器用不同的数据获得不同梯度，机器在不同数据上学习不同的更改，并更新参数梯度，最后算出全局梯度，目前采用的是加和求平均的方式，效果已经非常好了。此前 Colossal-AI 在数据并行中的 LARS 方法，就为谷歌、 Meta、腾讯、索尼等公司，将 ImageNet 的训练时间从一小时缩短至一分钟。

数据并行是最基本的，同时也是最稳定的。将数据划分之后，假设未来有 1 万个 GPU，很容易发生的情况是，隔几个小时就有四五个 GPU 崩溃了，运维 1 万个 GPU 的集群很难，但数据并行的稳定之处在于，即便有 1 万个GPU 崩溃了十几个，但大体结果是不会变的，因为它是梯度加和求平均。

基于这个考虑，我认为数据并行是一个根本性的基础设施。

当然，仅用数据并行并不够，原因在于：数据并行有一个假设，必须将模型拷贝到每个 GPU 或服务器内，由服务器或 GPU 去交换梯度。但如果 GPU 仅 80G 内存时，万亿参数的模型则需要几十 T 的内存，这在 GPU 中是无法存放的，需要将模型切割至不同的 GPU 上再汇总结果。这种方法叫做模型并行。模型并行包括两种，第一种是张量并行（ tensor paralism），即层内的模型并行。例如 GPT-3 的层数大概为八九十层，每层切割一次模型，将其层内计算分隔成多份，算完一层再算下一层，依次类推，这就是张量并行。

另一种方式则是 Pipeline Parallelism（流水线并行），介于数据并行和张量并行外的一种模型并行方式。通过构建几个数据 pipe（管道），每个数据 pipe 的数据点不同，相当于将一个大尺寸分割为多个小尺寸，通过这种方式进行 pipe 计算。假如有 10 个 pipe，10 个 pipe 代表十组不同的数据，第一个 pipe 计算第一层的数据，第二个 pipe 计算第二层......以此方式并行，类似我们盖楼一样，10 个工程队盖 1000 层楼，当第一个工程队在盖第一栋楼的第一层，第二个工程队盖第二栋楼的第二层，依此类推。

当楼数越多，楼和工程队之间的比值越高，效率也越高，相当于 10 个工程队在同时运转。其中每个工程队就相当于一个 GPU，每个楼就相当于一个 pipe，楼的层数相当于这个神经网络的层数，这就流水线并行的核心逻辑。

目前工业界已经做了相关的工作，除了 Colossal-AI 之外，还有英伟达的 TensorRT 和微软的 DeepSpeed，他们也是技术壁垒最高的两家公司。

但 Colossal-AI 与其不同之处是，Colossal-AI 专注于未来大模型的发展方向。可以看到，当前的模型还在变得更宽，而不是变得更深，张量并行将会更加重要，但它最大的弊端就在于，因为它是切割的是整个层，通信开销太大。这也是为什么英伟达 CEO 在 GTC 峰会上首次介绍 3D 并行时特别说明其通信开销太大的问题，只能放到一个服务器内去做。因此，Colossal-AI 主打 2D 张量并行和 2.5D 张量并行，将计算成本降低了一个数量级。

这就意味着用一维张量并行，1 万个机器里，每个机器都需要跟 9999 个机器打交道，而 2D 张量并行则是将其分成了各个子单元，每个机器只需要跟 96 个机器打交道。它的核心逻辑是，用一些 local synchronization（局部同步）去取代global synchronization（全局同步），以更多的局部通信去取代全局通信，这个过程中，设计调度是最困难的。

3D 张量并行也是同样，每升高一个维度，它的设计复杂度会高一个数量级，最终通信复杂度下降了。

在内存优化方面，目前 AI 大模型训练的内存开销很大，即便什么事情都不做，也需要几 T 的内存，如果不进行人工干预，一旦使用起来，可能需要几十 T 甚至是几百 T 的内存。

为了让模型预测效果更好，我们往往需要长序列数据，当前大模型的本质是通过一个单词的输出来预测下一个单词的概率，长序列成为刚需。对此， Colossal-AI 也推出了 Sequence Parallelism（序列并行）。

具体而言，在将序列进行切割后，会面临一个严重的问题是：在进行 attention score 时，每个 token 都需要跟全局序列中的其他 token 去评估，而切分后的服务器上只有部分 token，其他服务器上也会分布不同的 token，以至于每个服务器运行是需要同其他服务器打交道。

也就是说，假设今天屋子里 200 个人每人分别拿了一包零食，我希望每个人能品尝下其他所有人的零食，至少需要 200 个平方次交换，才能让每个人都尝到其他人的所有零食。那么最简单的方式是：所有人围一个圈，每个人将自己吃过的零食递给右手边的人，从自己的左手边获得零食，仅需 n-1 次，即 199 次的传递即可完成。从而降低了整个通信成本。

总结一下，目前 AI 大模型训练的核心技术栈，其实就是并行计算，因为我们要处理成百上千上万个 GPU 核心，把所有 GPU 并行利用起来。数据并行、张量并行、流水线并行以及数据序列并行是并行中较为核心的模块。

目前在内存优化方面，我们处于一个没有太多选择的环境，英伟达 GPU 是最好的，我们好像也没有其他更好的方案能够去取代它。但美中不足的是，英伟达 GPU 的内存有限，在这种情况下，我们能否思考如何利用 CPU 内存、NVMe内存，核心思想就是，GPU 内存放不下就移到 CPU 上，CPU 放不下就放到 NVMe 上，相当于在盖楼时，所需要的原材料自家楼下工地放不下，那我们就将其放到隔壁工厂。其技术的核心也在于最小化数据移动，即最小化 CPU、 GPU 之间的数据移动，最强化 CPU 和 NVMe 之间的数据移动，从而将吞吐力速度提升到最高。

走向开源

Colossal-AI 是一个开源软件，同时我们也做了一个商业化的平台，对没有 GPU 的用户，可以直接在平台上去训练部署自己的大模型。我们也提供了 LLaMA、PaLM、 GPT 等各种模型，最快仅需两三天就可完成一个模型微调。相比之前用户可能需要几周、乃至几个月来处理硬件、软件等基础设施，效率得到了极大提升。同时，Colossal-AI 也保护用户的隐私，平台不会保留、访问用户数据，这是 Colossal-AI 与 OpenAI ChatGPT 的本质区别。我们将数据上传至 Google Cloud 时，很多时候 Google 并没有碰我们的数据，但是 OpenAI GPT 会进行分析，AI 模型的不可解释性、训练不彻底等风险普遍存在。因此，未来也会有很多企业训练自己的大模型，Colossal-AI 做的，是最大化保护用户的隐私，同时提供大模型训练的工具。

在性能上，Colossal-AI 在同样的硬件上可以训练24 倍大的模型，相对于 DeepSpeed 的 3 倍加速，即便是一个低端的服务器，也可以借助 Colossal-AI 完成相应的模型训练。例如 LLaMA-65B 的训练，在 Colossal-AI 上使用同样的代码放，可以直接获得约 50% 的加速效率。

一个简单的比喻，比如说现在大模型是挖金子，英伟达是卖铲子的，那我们就是卖手套、卖衣服的，把挖金子的效率提到最高。