WAIC 2023 | 新加坡国立大学尤洋教授 ：AI大模型的挑战与系统优化
机器之心报道
演讲：尤洋
编辑：微胖
在 WAIC 2023 AI 开发者论坛上，新加坡国立大学校长青年教授、潞晨科技创始人兼董事长尤洋发表了主题演讲 《Colossal-AI：AI 大模型的挑战与系统优化》。演讲中，尤洋首先介绍了大模型「巨人症」问题所在以及开源软件 Colossal-AI 社区目前发展情况。接下来，尤洋介绍了 Colossal-AI 背后的一些技术细节，主要是训练大模型的并行系统以及内存优化方面的努力。最后展示了 Colossal-AI 在产业应用上取得的成果。
以下为尤洋在 WAIC 2023 AI 开发者论坛上的演讲内容，机器之心进行了不改变原意的编辑、整理：
今天的演讲主要介绍我们公司做的开源软件 Colossal-AI 的一些技术原理和应用。
首先简单介绍一下我本人。我在加州大学伯克利分校获得博士学位，现在在新加坡国立大学任教，同时创立了潞晨科技。公司另外一位核心成员 James Demmel 教授是美国科学院工程院院士，也是加州大学伯克利分校前计算机系主任兼院长。
今天的演讲分四个部分。第一部分简单介绍大模型的挑战以及目前 Colossal-AI 社区的发展情况，接下来两个部分介绍一些技术细节，最后一部分介绍具体应用上的效果。
首先给大家展示一张图：横坐标是时间，纵坐标是 AI 模型的参数量。过去六年中，最好的 AI 模型参数量已经上升了 1 万倍左右。
比如，2016 年 ResNet-50 只有 2000 万参数，2020 年 GPT-3 已经达到 1750 亿参数的规模。据说 GPT-4 也是用的混合专家系统，跟谷歌 Switch Transformer、智源的「悟道」都是同一种技术。Switch Transformer 参数规模大概 1.6 万亿，据说 GPT-4 有 16 个专家（模型），每个专家（模型）有千亿左右（参数）。
所以说，过去六、七年 —— 从 ResNet-50 到 GPT-4—— 最好模型的参数量刚好大了 10 万倍左右。
但是，以 GPT-3 为例，模型构造没有到 100 层，ResNet-50 也是 50 层左右，层数基本上没变化，模型不是变得更深而是变得更宽，大了 1-10 万倍左右，也给 GPU 内存造成更大压力，但 GPU 内存每 18 月只增长 1.7 倍，这就需要对下一代人工智能基础设施进行优化或者重建。
所以，我们打造了 Colossal-AI 系统。这是 Colossal-AI 系统结构图，包括三个层次。
第一个层次是内存管理系统，因为大模型太吃内存。
第二部分是 N-Dim 并行技术（N 维并行技术）。据说 OpenAI 已经用 10 万张 GPU 卡训练大模型。前两天一家美国创业公司融资了 13 亿美金，背后基础设施据说已经有 2 万张 GPU 卡。未来，从 1 个 GPU 到 10、100、10000 个 GPU ，自动扩展效率会对训练系统产生根本性影响，所以我们打造了 N 维并行系统。
第三部分是低延迟的推理系统，也是目前 Colossal-AI 的主要模块。模型训练好后要服务用户，用户每调用一次模型就是做一次推理，这跟成本有直接关系。所以，推理的延迟要很低，成本要降到最低。
虽然 Colossal-AI 开源社区只推出了 20 个月左右，但发展速度非常快。（下图）横坐标是时间，纵坐标是 GitHub 上的星数，可以看出 Colossal-AI 增长速度远超于传统开源软件。
Colossal-AI 增速也远超与 Colossal-AI 类似软件，比如 DeepSpeed。
目前 Colossal-AI 用户遍布全球。中国、美国、欧洲、印度、东南亚都有很多用户。在全球 AI 生态系统中也都发挥了更重要的作用。
目前世界公认的第一大 AI 生态是 PyTorch ，它的生态图中有一个链接指向了 Colossal-AI。Lightning AI 是 PyTorch 第一大子生态，直接依赖于 Colossal-AI。
HuggingFace 世界公认的第二大 AI 生态，很多应用也是基于 Colossal-AI 做的。
现在开源生态做的最好的是 PyTorch 也由 Facebook 主导，我们的开源大模型很多基于 Facebook 的 LLaMA。
OPT 也是 Facebook 的一个很重要模型。OPT 官网截图显示，它也有一个链接指向了 Colossal-AI，OPT 用户也可以通过 Colossal-AI 进行优化。
第二部分讲一下技术细节。假设未来需要用 1 万或者 10 万张 GPU 卡训练大模型，怎么充分利用这些 GPU 卡呢？与二十年前相比，今天特别是下层系统软件面临的问题有很大不同。
20 年前，英特尔芯片快了 3 倍，因为都是串行代码，一行代码都不用改就可以直接快 3 倍。但在今天，提升算力主要靠并行计算，不管是单个 GPU 内的多线程还是多个 GPU，甚至于两三周之前英伟达 CEO 黄仁勋在演讲中说到，要把成百上千个 GPU 用高速网络连接在一起（现在已经是这样的状况了）。如何最大化算法或者说上层应用并行度，将对机器的性能发挥产生非常实质性影响。在同样条件下，好的分布式软件导致速度上出现 5-10 倍的差距也很正常，所以，并行系统非常重要。
目前，训练 AI 大模型的并行系统主要有三个。最开始也是黄仁勋在 GTC 中介绍的，包括数据并行、模型并行和流水线并行。
模型并行是指层内并行。刚才也说了，现在处在宽度学习时代，而不是深度学习时代。从 ResNet-50 到 GPT-3，层的数量基本没有变化，但参数量大了 1 万倍，变宽了很多。在层的宽度越来越大、每层计算量越来越大的情况下，张量并行有可能发挥更加重要的作用。
当然，张量并行最大的弊端是它的通信开销太大，这也是黄仁勋在 2021 年演讲里提到把所有张量并行都放在服务器内的原因，跨服务器的通信开销太大，得不偿失。所以，如何优化张量并行的通信就非常重要。我们又打造了二维张量并行、三维张量并行、2.5 倍张量并行，去最小化数据移动。
另一个是数据序列并行。我们前段时间看到一些工作，看能不能把 sequence 增长到 4K 的量级。比如对 ChatGPT 应用而言，sequence 越长，一定概率上，预测效果也会更好，因为上下文得到了更多信息。但是，Transformer 的架构决定了每个 token 都需要跟其他的 token 算 Attention score，内存开销非常大。数据序列并行主要优化的也是 sequence 划分之后，如何尽量减少它们的数据移动。
当然，最传统的还是数据并行。我会一一介绍这三种并行系统。
先介绍一下最简单的数据并行。最开始我们训练 ResNet、AlexNet 甚至 BERT，很大程度上可以把模型下放到单个芯片，只要单个芯片能把模型放下去，芯片和芯片之间做数据并行就可以了。
比如，将一堆数据分 10 份，每个 GPU 放十分之一。这样情况下，工作就很理想了 —— 只需要增大 batch size 10 倍，假定能把循环次数减少 10 倍，这样 batch size 线性增大、循环次数线性减少就变成了一个很理想的扩展性问题，这样我们可以保持 Epoch 不变，但总体运行时间显著减少。
数据并行的通信，其实就是每次循环的时候大概需要交换一下梯度就可以了。当然，梯度大小和参数大小是一样的，对传统 CNN 而言，它们参数也不太多，所以，通信也非常友好，数据并行之前也是最常用的一种方式。
但是，大模型时代面临的问题是循环次数逐渐减少情况下，收敛变得越来越困难 —— 假如之前用 1 万次循环做收敛，batch size 增加 10 倍就需要在 1000 次循环中完成收敛。所以，数据并行的难度在于它的 accuracy 有时候会掉很多。因此需要看能不能设计出更好的优化方法，去解决这个问题。
我们团队打造了 LARS 和 LAMB 的方法，也发表在 ICLR 论文上，现在引用次数也快 700 次。最近，我们基于 LAMB 方法做了另外的工作，今年年初也获得了 ACL 杰出论文，另一篇获得 AAAI 2023 杰出论文。其实都是基于这个工作线，优化训练的收敛性和效率。
包括谷歌，Facebook SEER，Deepmind BYOL 等使用了 LARS 和 LAMB 方法。我们也帮助谷歌把 BERT 训练时间从三天缩减到 76 分钟。
说完数据并行，再说一下模型并行。模型并行就像盖楼 —— 要盖一个很宽的楼，每层 1 万平，直接用 10 个工程队把它分成 10 份，每个工程队一份，同步盖完一层后，再盖下一层。所以，层内并行的问题是通信开销太大。
第三个是流水线并行。并行的字面意思可以理解为，算完这一层算下一层，第 N 层要依赖于第 N-1 层的结果，第 N+1 层也要依赖于第 N 层的结果，前后依赖。比如，现在要盖 1000 栋楼，有 20 个工程队。第一个工程队先盖第一栋楼第一层，盖好之后，第二个工程队才能入场盖第二层。这时，第一个工程队移步第二栋楼，开始盖第一层，直到盖完 20 栋楼的第一层时，所有 20 个工程队才能都进入施工现场。
流水线并行的层数和 GPU 数量之间关系，就像工程队数量（类似 GPU 数量）与楼层数（相当于流水线并行的层数）关系一样密切。数据流水线数和 GPU 之间比值越大，并行效率越高。20 个工程队盖 1000 栋楼，很多时候并行度可以达到 20，但用 20 个工程队去盖 30 栋楼，效率就不是很高。所以，流水线并行本质上还是需要将 batch size 扩得很大，因为不同流水线进行了不同计算，本质上它的并行也来源于数据并行。
这是三种传统并行方式。GPT-3 刚出来的时候，OpenAI 用了 1 万个 GPU 训练 GPT-3，但是英伟达优化之后发现 3072 个 GPU 就够：把 64 个服务器分成一组，每个服务器内有 8 个 GPU，8 个 GPU 做张量并行。64 个服务器之间就做流水线并行，每个组有 64 个服务器，组之间再做数据并行，因为流水线并行的传输代价非常小，只传层和层之间的信息。
我们之前的工作做了很多张量并行的优化。张量并行未来发展空间也是最大的，因为模型的层变得更宽，通信开销也非常高昂。而二维张量并行、三维张量并行的核心思想是用更多的局部序列化替换全局序列化 ，用更多局部通信换取全局通信，从而降低通信成本，提高效率。
这几种张量并行方式的训练效率差异非常大，特别是二维和三维张量并行。可以看一下最后两列（最右侧两列）：第一个是通信带宽上的代价，第二个是通信延迟的代价。二维、三维、2.5 维张量并行可以大幅度降低通信和内存开销。
数据序列并行非常重要。前段时间有一个美国创业公司发了一个大新闻，可以把单次数据序列的 input sequence length (输入序列长度) 做到 4K Token。这个事情很重要，这张图可以说明一些问题：
横坐标反映数据序列长度，纵坐标预测下一个词的精度 —— 数据序列越长我们预测的越准确，因为它有更多的上下文信息。
虽然不好的地方在于计算变得更加复杂，但长的 Sequences 是算法侧需求，下层的基础设施也应该尽可能去满足。
这两天我又看到一篇论文，有团队声称把数据序列甚至能做到 100 万（当然需要再求证）。数据序列变长的需求还是会长期存在，越长确实效果就越好。
这张图展示的是为什么变长之后，给下层实现、基础设施带来很多麻烦。
现在基于 Transformer 的架构，每个 Sequence 都有很多 token，每两个 token 之间都需要算 Attention score，导致内存开销非常大，特别是 Sequence 长度到一定级别后，内存开销会指数级上涨。
Transformer 结构内存压力本来很大。比如训练 GPT-3 ，如果按照 2000 亿参数量计算，假设用单精度，每个参数要占 4 个字节，仅参数就要占 800G 内存，梯度要占 800G 内存，等等。何况 Transformer Attention score 等更是指数级上升，内存压力非常大。因此，如果数据序列想变得更长，就需要并行，把不同 Token 划分到不同的 GPU 上。
另一个问题是如何减少 GPU 之间的数据移动。100 个 GPU，所有都需要跟其他 GPU 交换数据，这是 Transformer 的 Attention 机制决定的。我们打造了环状 self-attention 通信算法，尽量减少数据之间移动，每次循环只需要跟左右邻居打交道，把通信复杂度从 P 平方降低到 P-1（P，指 GPU 数量或服务器数量）。
时间关系，实验结果就不介绍了，感兴趣的话可以看一下我们的官网，Colossal-AI 做了很多优化。
比如接下来要介绍的内存优化。其实就是减少 GPU 和 CPU 之间、CPU 与 NVME 硬件之间的数据移动。这一点在未来非常重要。因为 GPU 内存有限，我们需要用 CPU 甚至 NVME 硬件，但它们之间的数据移动可能比计算慢上千倍，如何最小化它们之间的数据移动，以更低成本容纳大模型将非常重要。
快速过一下一些应用效果。在一些重要应用比如 Stable diffusion 上，Colossal-AI 取得了很好的加速。我们可以把内存开销降低 5.6 倍，硬件成本降低 46 倍。Colossal-AI 经过了很多业界检验。
用 Colossal-AI 加速过的效果图是没有任何损失的。
我们通过 Colossal-AI 做了 ColossalChat，是世界上首个开源的最接近 ChatGPT 原始技术方案，具备完整 RLHF 流程的低成本 ChatGPT 复现方案。只有 70 亿或 100 亿的参数模型可以达到更好效果。
这是原理图。
相对于业界标准，Colossal-AI 可以将推理上获得 30-50% 的提速，在训练上获得大概 7 倍加速。
包括可以把 ChatGPT 训练成本从 300 万美金降低到 140 万美金左右。
接下来是一些 demo 展示和与斯坦福羊驼的对比，以及一些更进一步分布式 PPO 和 evaluation 的工作。
当然，由于我们做的是计算优化，因此不会限制于某个特定行业或者模型，具备良好的通用性。例如对于生物医药行业的蛋白质预测模型 AlphaFold2，我们也可以优化提升约 10 倍的训练推理速度。