在WAIC2023AI开发者论坛上，新加坡国立大学校长青年教授、潞晨科技创始人兼董事长尤洋发表了主题演讲《Colossal-AI：AI大模型的挑战与系统优化》。演讲中，尤洋首先介绍了大模型「巨人症」问题所在以及开源软件Colossal-AI社区目前发展情况。接下来，尤洋介绍了Colossal-AI背后的一些技术细节，主要是训练大模型的并行系统以及内存优化方面的努力。最后展示了Colossal-AI在产业应用上取得的成果。

　　以下为尤洋在WAIC2023AI开发者论坛上的演讲内容，机器之心进行了不改变原意的编辑、整理：

　　今天的演讲主要介绍我们公司做的开源软件Colossal-AI的一些技术原理和应用。

　　首先简单介绍一下我本人。我在加州大学伯克利分校获得博士学位，现在在新加坡国立大学任教，同时创立了潞晨科技。公司另外一位核心成员JamesDemmel教授是美国科学院工程院院士，也是加州大学伯克利分校前计算机系主任兼院长。

　　今天的演讲分四个部分。第一部分简单介绍大模型的挑战以及目前Colossal-AI社区的发展情况，接下来两个部分介绍一些技术细节，最后一部分介绍具体应用上的效果。

　　首先给大家展示一张图：横坐标是时间，纵坐标是AI模型的参数量。过去六年中，最好的AI模型参数量已经上升了1万倍左右。

　　比如，2016年ResNet-50只有2000万参数，2020年GPT-3已经达到1750亿参数的规模。据说GPT-4也是用的混合专家系统，跟谷歌SwitchTransformer、智源的「悟道」都是同一种技术。SwitchTransformer参数规模大概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。LightningAI是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算Attentionscore，内存开销非常大。数据序列并行主要优化的也是sequence划分之后，如何尽量减少它们的数据移动。

　　当然，最传统的还是数据并行。我会一一介绍这三种并行系统。

　　先介绍一下最简单的数据并行。最开始我们训练ResNet、AlexNet甚至BERT，很大程度上可以把模型下放到单个芯片，只要单个芯片能把模型放下去，芯片和芯片之间做数据并行就可以了。

　　比如，将一堆数据分10份，每个GPU放十分之一。这样情况下，工作就很理想了——只需要增大batchsize10倍，假定能把循环次数减少10倍，这样batchsize线性增大、循环次数线性减少就变成了一个很理想的扩展性问题，这样我们可以保持Epoch不变，但总体运行时间显著减少。

　　数据并行的通信，其实就是每次循环的时候大概需要交换一下梯度就可以了。当然，梯度大小和参数大小是一样的，对传统CNN而言，它们参数也不太多，所以，通信也非常友好，数据并行之前也是最常用的一种方式。

　　但是，大模型时代面临的问题是循环次数逐渐减少情况下，收敛变得越来越困难——假如之前用1万次循环做收敛，batchsize增加10倍就需要在1000次循环中完成收敛。所以，数据并行的难度在于它的accuracy有时候会掉很多。因此需要看能不能设计出更好的优化方法，去解决这个问题。

　　我们团队打造了LARS和LAMB的方法，也发表在ICLR论文上，现在引用次数也快700次。最近，我们基于LAMB方法做了另外的工作，今年年初也获得了ACL杰出论文，另一篇获得AAAI2023杰出论文。其实都是基于这个工作线，优化训练的收敛性和效率。

　　包括谷歌，FacebookSEER，DeepmindBYOL等使用了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栋楼，效率就不是很高。所以，流水线并行本质上还是需要将batchsize扩得很大，因为不同流水线进行了不同计算，本质上它的并行也来源于数据并行。

　　这是三种传统并行方式。GPT-3刚出来的时候，OpenAI用了1万个GPU训练GPT-3，但是英伟达优化之后发现3072个GPU就够：把64个服务器分成一组，每个服务器内有8个GPU，8个GPU做张量并行。64个服务器之间就做流水线并行，每个组有64个服务器，组之间再做数据并行，因为流水线并行的传输代价非常小，只传层和层之间的信息。

　　我们之前的工作做了很多张量并行的优化。张量并行未来发展空间也是最大的，因为模型的层变得更宽，通信开销也非常高昂。而二维张量并行、三维张量并行的核心思想是用更多的局部序列化替换全局序列化，用更多局部通信换取全局通信，从而降低通信成本，提高效率。

　　这几种张量并行方式的训练效率差异非常大，特别是二维和三维张量并行。可以看一下最后两列（最右侧两列）：第一个是通信带宽上的代价，第二个是通信延迟的代价。二维、三维、2.5维张量并行可以大幅度降低通信和内存开销。

　　数据序列并行非常重要。前段时间有一个美国创业公司发了一个大新闻，可以把单次数据序列的inputsequencelength(输入序列长度)做到4KToken。这个事情很重要，这张图可以说明一些问题：

　　横坐标反映数据序列长度，纵坐标预测下一个词的精度——数据序列越长我们预测的越准确，因为它有更多的上下文信息。

　　虽然不好的地方在于计算变得更加复杂，但长的Sequences是算法侧需求，下层的基础设施也应该尽可能去满足。

　　这两天我又看到一篇论文，有团队声称把数据序列甚至能做到100万（当然需要再求证）。数据序列变长的需求还是会长期存在，越长确实效果就越好。

　　这张图展示的是为什么变长之后，给下层实现、基础设施带来很多麻烦。

　　现在基于Transformer的架构，每个Sequence都有很多token，每两个token之间都需要算Attentionscore，导致内存开销非常大，特别是Sequence长度到一定级别后，内存开销会指数级上涨。

　　Transformer结构内存压力本来很大。比如训练GPT-3，如果按照2000亿参数量计算，假设用单精度，每个参数要占4个字节，仅参数就要占800G内存，梯度要占800G内存，等等。何况TransformerAttentionscore等更是指数级上升，内存压力非常大。因此，如果数据序列想变得更长，就需要并行，把不同Token划分到不同的GPU上。

　　另一个问题是如何减少GPU之间的数据移动。100个GPU，所有都需要跟其他GPU交换数据，这是Transformer的Attention机制决定的。我们打造了环状self-attention通信算法，尽量减少数据之间移动，每次循环只需要跟左右邻居打交道，把通信复杂度从P平方降低到P-1（P，指GPU数量或服务器数量）。

　　时间关系，实验结果就不介绍了，感兴趣的话可以看一下我们的官网，Colossal-AI做了很多优化。

　　比如接下来要介绍的内存优化。其实就是减少GPU和CPU之间、CPU与NVME硬件之间的数据移动。这一点在未来非常重要。因为GPU内存有限，我们需要用CPU甚至NVME硬件，但它们之间的数据移动可能比计算慢上千倍，如何最小化它们之间的数据移动，以更低成本容纳大模型将非常重要。

　　快速过一下一些应用效果。在一些重要应用比如Stablediffusion上，Colossal-AI取得了很好的加速。我们可以把内存开销降低5.6倍，硬件成本降低46倍。Colossal-AI经过了很多业界检验。