01 大模型训练的挑战

1. 大模型发展现状及问题

上图中展示了大模型的发展历程，其中纵坐标是 AI 模型的参数量。在过去几年间，AI 大模型的参数量飞速增长，大约每 18 个月时间就会增长 40 倍。例如 2016 年，世界最好的大模型 ResNet-50，参数量约为 2000 万；而到 2020 年的 GPT-3 模型，参数量已达到 1750 亿；到今天，根据 OpenAI 透露的消息，MOE 混合专家系统这种架构的大模型参数量大约 1.7 万亿；2021 年初 Google 的 Switch Transformer 模型采用了相同的混合专家系统架构，参数量已达到 1.6 万亿。

近来，OpenAI 在训练的 GPT-5 模型，预计将达到 100 万亿参数。此外，Inception 等创业公司甚至认为未来模型的参数量将超过 GPT-5 模型 100 倍。

随着模型规模的指数级增长，未来可能需要使用成千甚至上万个 GPU。因此，背后的软件基础设施至关重要，不管是 AI 框架还是分布式系统，都将起到关键作用。

未来的 AI 训练，其核心在于扩展性和高效性。所谓扩展性，即从 1 个 GPU 扩展到 10 个、100 个、1000 个、1 万个甚至更多的 GPU，而在保证扩展性的同时能否始终保持百分之八九十的扩展性效率？这就是所谓的扩展性和高效性。

2. Why Do We Believe in Large Models?

我们投入了大量精力去完善基础设施，就是为未来 AI 大模型而准备的。从源头上说，我们之所以看好大模型的发展，原因主要有两点：其一，在模型更大的情况下，确实能够获得更优的性能，其 AI 智能程度也更高；其二，大模型的“大”字其实很难定义（今天，100 亿参数属于大模型，明天，可能 1 万亿参数才是大模型），但是有一点是肯定的，那就是未来的小模型也会呈现指数级上升。比如我们今天觉得 Bert 已经成为了一个小模型，两年之后可能 GPT-3 也变成了我们眼中的小模型；再过三年，可能今天的 GPT 子版本 ChatGPT 也变成小模型了。

如上图所示，横坐标为时间，纵坐标为模型参数量。例如 2015 年，中小型企业一般会使用 LeNet；到 2016 年，中小型企业已经能用 ResNet-50 了；到 2020 年的时候，中小型企业甚至可以使用 Bert、GPT-2 等模型；预计到明年，中小型企业会陆续使用 LLAMA 这类模型（参数量大约 700 亿）。如上文所说，小模型实际上也在指数级上升。

因此所有 AI 模型的训练和微调都需要分布式，因为至少需要几十甚至上百张 GPU。近期，英伟达的 GH 200 通过将 CPU 和 GPU 连接在一起，使用更高速的网络将其串联。未来，GPU 训练会更加便捷，就像今天的服务器一样，厂家自行采购一套实体服务器或云端服务器，即可实现大模型的训练。

大模型的训练往往需要很高的并行度，那么如何提高模型训练的效率呢？

首先可以关注云上计算的价格。如上图，横坐标是时间，纵坐标是单位美元能买到的浮点数计算次数（T flops 次数）。从图中可以看出，从 2015 年到 2020 年，该次数呈指数级上升了 45 倍。因此，绝对的算力实际上是在变便宜的，即单位美元买到的浮点数计算次数是呈指数上升的。因此，如果能将算力挖掘出来，预计到 2026 年，GPT-3 的训练成本将从现在的 300 万美金降低到 7 万美金左右。

如果进一步优化模型的并行训练效率，设计出更好的算法，使其收敛速度变得更快，这个训练成本可能会降低至两三万美金（甚至可能会降到 1 万美金以下）。

3. 大模型的挑战

然而现在，训练成本仍然是很高的，例如谷歌的 PaLM 模型，按照云上的价格，需要 900 万美金。另一方面，如果一个公司想从零开始训练一个大模型，至少需要投入五六十个人，带来了较高的人力成本，增加了模型落地的困难。

对于大部分的业务型企业，企业本身希望快速实现对大模型的训练、部署和微调，这些对大模型的落地是非常关键的因素，而上述现状对大模型工作在这类企业的开展就会变得非常困难。

4. Colossal-AI

基于上述原因，我们打造了 Colossal-AI 这个工具，希望帮助更多企业实现大模型的快速训练、微调和部署。

Colossal-AI 包括三个部分：

• 高效的内存管理系统

• N 维分布式系统

• 低延迟推理系统

利用以上三个部分，三位一体地实现大模型的高效训练。不管用户使用什么样的编程接口或者框架，通过 Colossal-AI 的内核都可以在 CPU、GPU、NPU 上获得最好的性能。

上图回顾了 Colossal-AI 开源社区发展的情况，横坐标是时间，纵坐标是 Github 星数，可以看出 Colossal-AI 目前的增长速度已经远远超过了同时期的 Spark、ClickHouse、MongoDB、Kafka 等。

其背后的原因，除了 Colossal-AI 自身优秀的开源设计，还在于这一波 AI 的浪潮预计比上一波大数据的浪潮更大，未来 AI 的价值可能会贯穿到企业的整个生命周期，因此未来的机会是非常大的，降低模型的训练难度和训练成本，是非常重要的一件事情。

上图是 Colossal-AI 的用户分布图，可以看出，目前在中国、美国、欧洲、印度、东南亚都有很多的用户。

同时，很荣幸，Colossal-AI 被收录到很多国际顶尖 AI 会议的官方教程，如今年的 Neural NIPS（之前叫 NIPS，该会议是目前 AI 的第一技术峰会，每年 OpenAI、谷歌等顶尖技术人员，包括像图灵奖得主等，都会到该会议上演讲），Colossal-AI 很荣幸成为该会议的官方教程。此外，Colossal-AI 曾被收录进了 AAAI、CVPR 等顶尖 AI 会议的官方教程。

目前，MLPerf 是衡量 AI 芯片的公认指标，是英伟达、英特尔、谷歌、阿里巴巴等公司共同发起的基准，现在称为 mlcomments。现在 MLPerf 的官方 Github 也依赖 Colossal-AI，因此 Colossal-AI 也很荣幸地成为了 MLPerf 标准的一部分。

此外，在全球 AI 生态系统中，目前 Colossal-AI 也处于非常理想的位置。例如，目前世界第一大 AI 生态 PyTorch，其官方生态直接依赖于 Colossal-AI 的分布式；此外，世界第二大 AI 生态 Hugging Face 的很多应用也是基于 Colossal-AI 的。目前开源 AI 大模型由 Facebook 所主导，LLAMA 和PyTorch 都是 Facebook 开发的，目前 Facebook 官方大模型 OPT 有一个链接就指向了 Colossal-AI。

02

N 维并行系统

未来大模型的训练，会涉及处理成百上千个 GPU 集群内部的关系问题，即如何把并行效率提升到极致。未来使用 AI 超级计算机（如使用 256 个 GH 200 组装而成），虽然其带宽很高，通讯传输效率很高，但仍然需要足够的并行度才能将算力发挥出来。

对于其他科技巨头，如华为，使用类似英伟达的 GH 200 这种超级计算机，其在并行计算方面，主要使用的方式包括数据并行、流水线并行和张量并行，而第一个实现这项技术的还是英伟达。

GPT-3 模型是 2020 年 6 月份训练出来的。在随后的 GTC 峰会上，英伟达黄总介绍了他们是如何训练 GPT-3 的，当时提出 3-D 并行这一概念。他们用了 3072 个 GPU（每个服务器有 8 个 GPU），每个服务器内使用 Tensor parallelism 进行张量并行，然后 64 个服务器形成一个小组，小组内部使用流水线并行，而 6 个小组之间使用数据并行，这样就是 8x64x6=3072。

这样的配置，主要考虑的核心点就是张量并行的通信代价太高，因此在服务器内部使用流入线并行的方式，通信代价相对较低。在服务器之间，为了扩展到更大，又引入了数据并行的方式。而未来 GH200 组成的超级计算机，其网络带宽可能会非常高，芯片之间可以直接交换，这样就可以直接使用张量并行的方式（虽然通信代价高，但如果网络传输非常好，则可以很好地解决这个问题）。

回到今天，对于张量并行，最开始实现的是英伟达的 microtrain，然而其并没有把张量并行优化到极致。

张量的本质就是矩阵，矩阵的切片是非常灵活且容易实现的。然而理论上每个机器都需要跟其他所有机器打交道，这里就会涉及到 global synchronization，而我们的想法是能不能用更多的 local synchronization 来取代 global synchronization，这样会获得指数级的速度增长。因此我们提到了 2 维张量并行、3 维张量并行，甚至 2.5 维张量并行，其实质就是上述思想。

第二点，回顾大模型的变化趋势会发现，2016 年的 Resnet 50 模型，有 2000 万参数，包含 50 层；到 2020 年，GPT-3 模型有将近 2000 亿参数，只包含不到 100 层；而今天的 GPT-4 模型也只有 120 层。模型的参数规模越来越大，而层数却没有大的变化，说明模型的层变得更宽了。在这样的趋势下，未来对张量并行的要求将会更高，因此张量并行未来仍有很大的优化空间。对于流水线并行，其好处在于通讯代价很低，但是它总是有 bubble 的开销，会限制并行的效率，其 pipeline parallelism 只能在一定规模内使用，无法无限扩展。

第三点，我们需要新的并行 sequence parallelism。以对话类大模型为例，比如 ChatGPT，其 sequence Lens（contact Lens）是 8000（Bert 是 128512）。GPT 的核心原理是使用上下文信息去预测下一个词，因此，contact lens（即 sequel less）越长，预测的下一个单词就越准确。有些创业公司甚至想把 contact lens 做到 100 万，然而这样会带来巨大的内存开销压力（后文会介绍这一部分的优化）。

国内的一些通用大模型公司在 AI 基础设施这部分的关注点主要在于从千卡扩展到万卡。然而随着 GPU 数量的扩展，系统稳定性会下降，即使 OpenAI、微软这样的公司，其万卡集群的利用率也只有 30% 左右。如果能把稳定性的问题解决好，下一步就是扩展性的问题。

从 100 卡扩展到 1000 卡，已经带来了很大的困难，其效率可能会从 70% 直接降到 30% 左右。而从 1 千卡到 1 万卡，如果没有足够的并行计算效率，其训练可能会变得更慢。

这就引出了另一个问题：要将 AI 训练的基础设施增长到万卡训练的规模，哪方面是欠缺的？首先 pipeline parallelism 只能在局部使用，因此无法扩展到万卡；如果再加上 tensor parallelism 以及 sequence parallelism，未来还需要进一步优化数据并行。

数据并行的好处是，base size 增大 10 倍，epoch 数尽量固定，则计算量也是恒定的。如果 epoch 数量能够固定下来，batch size 增大 10 倍，则循环次数就减少 10 倍；如果同时实现数据并行，其通信就是梯度加和求平均，这样单次通信量又减少 10 倍；因此数据并行在万卡训练中也会发挥出更大的作用。假如把 batch size 扩大 10 倍，在保证 accuracy 的基础上从千卡扩展到万卡，其通信量会减少 10 倍，因此会更加高效。

2017 年曾有一篇论文《training ImageNet in one hour》，使用一小时训练 ImageNet。这一篇论文可以认为是整个分布式AI训练的开山之作。如今我们可能会觉得他的技术很简单，但其是第一个使用百卡来训练深度神经网络的一个高效案例。

当时只用了数据并行，而设备也只是 Ethernet，甚至没有使用 Infinity band，但其扩展性也能超过 90%，这说明 data parallelism 的通信效率还是非常高的。因此，对于未来的万卡训练，这个 data parallelism 也会非常关键。此外，我们仍然需要少部分的 pipeline parallelism、比较高效的 tensor parallelism 以及 sequence parallelism，也需要更大的 data parallelism。

前文提到的 data parallelism，只要将 epoch 数恒定，将 batch size 增大 10 倍，循环数就会减少 10 倍，而通信就是梯度加和求平均，因此还是非常 robust 的。假如使用 1000 个机器进行训练，即使有 5-10 个机器崩溃，也不会影响 data pair reserve 结果的一致性，因为只要样本分布在统计学上接近正态分布或其他较理想的分布，那么使用 1000 个 T5 加权平均和用 995 个 T5 加权平均的结果不会相差太大。

实际上 DIP LSM 业务存在一些自身问题，当 batch side 开得过大时，循环次数反而增长更大（例如 by side 增大了 1000 倍，循环次数增长了 2000 倍）。这就需要在 1/1000 的循环内寻找局部最优解，所以变得更具挑战性。因为这种情况下很容易陷入 sharp minimum，导致 accuracy 下降严重，从而得不偿失。因此 data parallelism 的挑战就是要设计出更好的优化方法。

Colossal-AI 中 loss 和 lamb 也做了很好的实现：首先在保持 icon 的情况下，以 VIP 为例，把 batch size 从 128 扩到 32K，同时计算时间从 733 小时降低到 40 分钟，因此可以看出其扩展性也是线性的。

之前我们的一些方法曾帮助腾讯、索尼、谷歌、富士通等公司将音频拆下来，训练时间从一小时缩短到了一分钟，并且曾帮助谷歌把 batch 训练时间从 3 天缩至 76 分钟。Data parallelism 可能在未来的万卡训练中发挥出更大作用。

下面介绍 pipeline parallelism 和 tensor parallelism。

Megatron 一维的方法非常简洁，但是其通信代价较高，因为其本质是 global synchronization，每一层都做的话计算量是非常恐怖的；这一点不同于 data parallelism，后者是算完一层后直接同步去通信。

而 Tensor parallelism 则是算完这一层才进行下一层，因此没有 overlap 的机会，导致每个层内都使用 together or reduce 的代价变得过高。

因此前文提到使用更多的 local synchronization 取代 global synchronization。

以 2D 方法为例，将每个机器视为一个网格，每次训练都先和同列的机器打交道，这样只需要跟根号 p 个机器打交道即可。同理，3D 方法就是将其分成三个维度，每次只和同个维度的机器打交道。

从上图可以看出，2D 方法的效率还是挺高的，且规模越大，其效果越好。

上图可以看出，参数内存的开销是差不多的。在通信层面，bandwidth latency 开销，可以看出维度的增加确实提高了通信效率。

未来的长序列也非常有挑战性，现在训练大模型的内存开销已经很高了，假如训练一个大概 2000 亿参数的 AI 模型（如 GPT-3），假定使用 LP32，则需要 200 乘以 10 的 9 次方；假如每个参数是 LP32 化，就是 4 个 Byte，则参数首先就会耗费 1000G 的内存（梯度优化器同理）。因此，在模型开始训练之前就需要上 T 的内存，而 GPU 一般只有 100G 的内存。该问题的严重性在于，如果使用更长的序列，则序列中每两个 token 之间都要计算 attention score，会产生很多中间内存的需求。因此序列越长，内存开销就越大，当达到一定的临界值之后，将会呈指数级上升。

因此，直观上，需要将序列划分成若干个不同的部分，将每一部分分配给不同 GPU 上不同的服务器，进而控制其内存开销上限。然而，在不同的 GPU 存放不同的 token，每个 token 都需要和其他 token 计算 attention score，这样会带来 p2 的通信操作，造成非常高的通信代价。基于此，我们提出了 Ring Self-Attention 方法。

该方法基于百度的 ring or reduce，其核心思想如下：比如 50 个人，每个人抱了一大包零食，现在需要每个人都尝一下其他所有人的零食，这样将会需要 2500 次的重新传输，但是如果大家手拉手围成一个圈，每次从左手边获得一包零食，再将其传给右手边的一个人，这样只需要批检一次自操作即可。

Ring Self-Attention，每次从左边邻居获得一个 key，再将其传递给右边邻居，进而计算出一个 tension score，循环往复，这样只需批检一次重新操作，即可算出所有信息。

上图中展示了一些很直观的比较，可以看出，throughput 相比之前的方法要快 50% 左右，能够同时在同样设备下训练的 sequence 长度比之前增长 5%。

03

高效内存管理系统

现在的AI架构主要还是基于 GPU 服务器去做训练，每个 GPU 的内存非常有限，因为内存只有 100G 左右，而训练一个千亿参数的模型至少需要上 T 的内存。CPU 上也有很多内存，GPU 放不下就放 CPU，再放不下就放 NVMe。

然而这样带来的问题就是会严重降低效率，因为 GPU 内的计算远快于 GPU 和 CPU 之间，以及 CPU 和 NVMe 之间的传输。因此使用这种多级的内存就需要最小化 GPU 和 CPU 之间，以及与 CPU 和 NVMe 之间的数据移动。

我们的方案是将参数分成了块状，每次选取有用的块，暂时用不上的块先存放到 CPU 上，通过这种方式动态地进行存储，将效率提到最高。使用这种方式，只要有一块内存，都会放满最常用的参数或者梯度。

04

杰出性能和应用案例

与 PyTorch 相比，Colossal-AI 可以训练 20 倍大小的模型，在同样设备条件下可以获得 3 倍的加速。从下图的比较可以看出，在相同设备条件下，可训练的参数量呈指数级上升。

从 Throughput 的图可以看出，Colossal-AI 与 DeepSpeed 相比，throughput 也是更高的。例如，在 Stable Diffusion 上可以获得约 6 倍的加速。同样在 inference 上打 Colossal-AI 的话也可以获得 40% 左右的加速，而加速之后，其内存得到了大幅的下降，但是对图的质量是没有任何影响的。

从上图可以看出，Stable Diffusion 的出图效果还是比较理想的。首先跟 PyTorch DDP 相比，Colossal-AI 在推理上可以获得 40% 左右的加速，在训练上可以获得 7 倍左右的加速。在同样的设备条件下，可以存放 10 倍大小的模型。

最近 Colossal-AI 也对 LLAMA 模型进行了优化，特别是在训练环节，获得了 2 倍左右的加速。

以上是 Colossal-AI 训练平台的介绍，可以实现一站式的 AI 模型训练、微调和部署。如果感兴趣，欢迎下载并查看 Colossal-AI 的开源代码，同时欢迎关注公众号，并加入微信群。

05

Q&A

Q1：大模型场景中，不同模型的结构差别比较大。以 Transformer 结构为例，包括 encoder-decoder 联合架构，也有纯 encoder 或纯 decoder 的架构。那么在训练优化和推理优化的过程中，每一种模型可能都需要进行 specific 的优化；如何看待 overhead 呢？

A1：对于 overhead，常会听到抱怨模型太多了。我们这个优化每种优化只能对到某个模型上，然后做像 confusion 这种。这个情况其实是往好的方向发展，因为近期逐渐都收敛到了 LLAMA、GPT、Falcon 这几种常见的模型，上述模型彼此间的相似度较高。例如未来 GPT-4 将相当于 GPT 加上 MOE；而加上 MOE 之后，对算子的下方是没有任何影响的。正因为未来大模型的试错成本越来越高，反而会导致未来收敛成少数常见的大模型，这是有利的地方。回到这个问题，如何通用地解决这类问题，其实个人感觉并没有太好的办法，只能尽量将通用的需求做好，在此基础上做优化。并且更严重的一个问题就是现在国产芯片有很多力量，就是想想未来布局一些新的事情，这种条件下我们可能就需要更多的优化。我觉得短期之内假如说有一个新的芯片或者有一个新的模型，让它直接跑起来，还是有挑战性的，所以说可能还真的需要一些优化。

Q2：相较于应用层，在模型框架层优化过程中，遇到最大的困难点是什么？

A2：近期全球算力紧张，受限于算力，只能通过和企业合作，帮助企业训练模型来验证框架的稳定性，这样反而不利于这个框架的发展。因此，算力资源对模型框架的发展是非常关键的。

Q3：和 deep speed Metro LM 相比，Colossal-AI 最大的技术优势是什么？

A3：每个团队的技术背景都不太一样，例如 deep speed 主要 做sharding Zero 内存优化方面的工作，而 Megtron 是英伟达在芯片背景条件下进行的整体优化。潞晨科技以前是做高性能计算的，因此会认为 AI 模型的分布式训练本质上就是数据移动的优化，因此核心思路在于通信的优化，从而实现上百、上千个 GPU 具有良好的扩展性。由于不同团队的基础背景不一样，因此对同一问题的看法也会不一样。

以上就是本次分享的内容，谢谢大家。

分享嘉宾

INTRODUCTION

尤洋博 加坡国立大学 校长青年教授、潞晨科技 创始人兼董事长

尤洋，新加坡国立大学校长青年教授，加州大学伯克利分校计算机科学博士学位，导师是计算机科学部和 EECS 系的前任系主任 James Demmel 教授。尤洋的研究兴趣包括并行/分布式算法、高性能计算和机器学习。他当前研究的重点是扩展分布式系统或超级计算机上的深度神经网络训练。2017 年，他的团队打破了 ImageNet 训练速度的世界纪录，并被 NSF、ScienceDaily、Science NewsLine、i-programmer 等科技媒体报道。2019 年，他的团队打破了 BERT 训练速度的世界纪录。BERT 训练技术已被 Google、Microsoft、NVIDIA  等众多科技巨头使用。尤洋团队开发的 LARS 和 LAMB 优化器可在行业基准 MLPerf 中使用。他曾获得 IPDPS 2015 最佳论文奖（0.8%）、ICPP 2018 最佳论文奖（0.3%）、AAAI 2023 杰出论文奖（0.14%）、ACL 2023 杰出论文奖（0.8%）和 ACM/IEEE George 迈克尔高性能计算奖学金。他是 Siebel 学者，Lotfi A. Zadeh 奖获得者。尤洋被加州大学伯克利分校提名为 ACM 博士论文奖（2020 年毕业的 81 名伯克利 EECS 博士生中有 2 名）。他还入选了《福布斯》亚洲 30 Under 30 榜单（2021 年），并荣获 IEEE CS TCHPC 早期职业研究人员高性能计算卓越奖。欲了解更多信息，请查看他的实验室主页：NUS HPC-AI Lab。