「八卦炉」炼丹规模直逼人脑！清华、阿里等搞了个174万亿参数大模型

---

  新智元报道  

编辑：David 好困

【新智元导读】手头只有鸡蛋和葱花，如何做出满汉全席？

只要不差钱，人人都爱大模型！

 

这不，Meta就计划把CV、NLP、语音识别等技术所需的模型都做到「数万亿」的参数量。

 

年初的时候就联合英伟达推出了一个全新的超算——「人工智能研究超级集群」（ AI Research SuperCluster，RSC）。

 

既然Meta都这么搞了，是不是拿超算「炼丹」很好用？

 

答案是，并不……尤其是当你手里只有14纳米架构的芯片时。

 

为此，清华大学、阿里达摩院等机构的研究人员提出一种全新的思路——「八卦炉」（BaGuaLu）。

 

这也是第一个针对在超算，也就是新一代「神威」上，训练参数量达到「百万亿级」模型的工作。

 

论文链接： https://tinyurl.com/bdzc4d2j

 

「八卦炉」通过结合特定硬件的节点内优化和混合并行策略，在前所未有的大型模型上实现了体面的性能和可扩展性。

 

结果显示，「八卦炉」可以训练14.5万亿个参数的模型，使用混合精度的性能超过1 EFLOPS，并且有能力训练174万亿个参数的模型。

 

显然，这已经相当于人脑中突触的数量了。

 

简化的模型计算过程

 

有趣的是，团队怕老外看不懂「谐音梗」，贴心地在论文的注脚中写道：「八卦炉是中国古代神话中的一种神奇炉子，可以生产药物。」

 

看来，「炼丹」这个词着实是不好翻译啊。

实现方法和结果

 

硬件的节点内优化

 

这部分主要是让内存带宽可以接近理论值。

 

以矩阵乘法为例，单精度和半精度的GEMM，可以分别获得了89.2%和85.8%的峰值性能。

 

使用DMA和RMA从CPE访问内存

 

为了将「八卦炉」应用到整个超算系统，研究团队设计了一个高效的混合策略，包括混合并行策略MoDa、负载平衡策略SWIPE和内存高效优化器ParO。

 

其中，混合MoE并行和数据并行策略(MoDa) 可以将模型训练扩展到大脑规模，并具有良好的计算效率。

 

SWIPE则通过将输入项目重新分配给专家，并严格限制每个专家必须接受同等数量的输入项目，从而有效地区分负载不平衡。

 

 

对于每个参数，系统必须存储参数本身、其梯度和优化器中的变量，这对内存来说是一个巨大的挑战。

 

而基于分区的并行优化器(ParO) 用Reduce-Scatter和All-Gather取代了数据并行组中的All-Reduce，将优化状态在工作者之间进行分割。

 

由于在每个等级上只保存了元数据和优化器数据，使得从优化器中恢复模型参数的完整副本成为可能。

 

 

混合精度训练

 

SW26010-Pro支持例如FP64、FP32、FP16和BF16等不同类型的浮点计算。 其中，FP64和FP32的吞吐量为14.03 TFLOPS，FP16和BF16的吞吐量为55.30 TFLOPS。因此，用混合精度进行训练会带来明显的性能改善。

 

然而，现有的方法大多数现有的工作是基于GPU的，并没有在这种大规模的模型训练上得到验证，直接套在「神威」上显然是不现实的。 例如，NVIDIA APEX的训练有四个优化级别：

 

O0：FP32训练。

O1：在诸如GEMM或卷积的运算中使用FP16。

O2：将权重和输入数据丢给FP16，并保持一个FP32主权重用于优化。

O3：FP16训练。

 

通常，在NVIDIA GPU上训练常规模型时，O1和O2级训练可以收敛到与O0相同的损失。然而，在「神威」上直接使用APEX方法时，O1和O2级优化都不起作用。

 

因此，研究团队将所有的层分为不同的类别。针对每个类别进行调整优化水平，并选择最佳的优化级别，这样就可以同时获得高性能和快速收敛。

 

「八卦炉」上的混合精度训练，在嵌入层使用O0，在注意力层使用O1，在FFN（前馈）层使用O2

 

为了保证数值的稳定性，避免可能的溢出和下溢。部分特定的运算符，如还原、exp、sqrt、gelu、softmax和layer_norm，则使用FP64运算。

 

在All-Reduce通信中，则采用了在线平均算法，而不是简单的求和除法算法来提高数值稳定性。

 

结果评估

 

研究团队评估了MoDa-1.93T、MoDa14.5T和MoDa-174T等模型的单精度和混合精度的性能。

 

MoDa-1.93T在单精度和混合精度中分别达到647 PFLOPS和1.180 EFLOPS。

 

MoDa-14.5T在单精度和混合精度方面分别达到了525 PFLOPS和1.002 EFLOPS。

 

MoDa-174T模型是所提出的模型中最大的一个，有173.9万亿个参数。训练过程中，在单精度和混合精度下分别达到198PFLOPS和230PFLOPS。

 

14纳米芯片，也能造E级超算

承载这个「人脑规模」模型计算任务的，是新一代的「太湖之光」神威超算。

 

此次新一代超算使用的处理器是2016年老一代超算的SW26010处理器的扩展版本：SW26010 Pro。

 

 

260核的SW26010处理器使用中芯国际28纳米工艺蚀刻而成，即使在当时也算不上是尖端芯片。SW26010-Pro处理器使用14纳米工艺，同样称不上尖端。

 

早在2021年2月，外界广泛猜测如何使用14纳米芯片建造E级超算系统，当时认为，如果要让发热量保持在低水平，同时保持其处理器的时钟速度不变，将芯片上的计算单元数量增加一倍，将矢量宽度增加到512位，并将机柜数量增加一倍，以达到FP64精度下的E级峰值理论性能。

 

事实证明，SW26010 Pro处理器将计算单元增加了50%，将矢量宽度提高到512位，也许还提高了时钟速度，并提高了节点和机柜的数量，使其达到E级性能。

 

SW26010-Pro的计算引擎架构

 

处理器中有六块核心组，每个核心组都有一个用于管理Linux线程管理处理元件（MPE）和一个由计算处理元件（CPE）组成的8乘8网格，有256 KB的二级缓存。

 

每个CPE有4个逻辑块，在一组单元上可以支持FP64和FP32，在另一组单元上支持FP16和BF16。SW26010-Pro中的每个核心组都有一个DDR4内存控制器和16GB内存，内存带宽为51.4GB/秒，整个设备有96GB的主内存和307.2GB/秒的带宽。

 

6个CPE通过环形互连连接，并有2个网络接口，使用专有的互连方式将它们与外部连接。

 

SW26010-Pro在FP64或FP32精度下的额定速度为14.03 petaflops，在BF16或FP16精度下为55.3 petaflops。

 

新超算系统一个机柜里有1024个节点（四个超级节点），总共40个机柜，从而达到125.4 petaflops的峰值理论性能。

 

据了解，SW26010-Pro处理器与SW26010的插槽是兼容的，所以系统设置应该是一样的。经测试，新系统的最大配置访问了107520个节点（每个节点搭载一个SW26010-Pro），总共有4193万个核心。这相当于105个机柜，是前一代的2.6倍。

 

新一代「神威」的架构

 

SW26010的额定峰值为3.06 teraflops，运行频率为1.45GHz。

 

做个计算，如果芯片制程从28纳米缩小到14纳米，然后将核心数量增加50%，再将带宽加倍，达到512位，保持时钟速度不变，这样下来，系统算力近似增加200%，达到9.2 teraflops。

 

但实际上SW26010-Pro芯片可提供14.03 teraflops的性能，按照这个结果推算，时钟速度提高了52.7%至2.22GHz，才能达到这个水平。

 

总体而言，在「八卦炉」上测试的105个机柜系统及其107250个SW26010-Pro处理器的峰值理论性能为1.51 exaflops。

 

未来OceanLight系统可能扩展到160个机柜，即163840个节点，峰值FP64和FP32性能略低于2.3 exaflops。

能否超越美国？

如果真的实现了160机柜的规模，该系统就可以超越今天美国橡树岭国家实验室正在优化中的 1.5 exaflops 「Frontier」的性能，以及2 exaflops 的「Aurora」超级计算机的峰值理论性能。

 

甚至可能比预计2023年进入劳伦斯利弗莫尔国家实验室的「El Capitan」超级计算机还要强，它的理论峰值性能在2.2~2.3 exaflops 左右。

 

 

诚然，采用相对落后制程的一个不可避免的缺点就是发热问题。SW26010-Pro芯片可能会很热，会增加电力和散热成本。但如果中芯国际能够在14纳米工艺上获得良好的良率，那么制造成本的降低可以在一定程度上弥补这部分费用。

 

比起散热和能耗成本问题，对于中国来说，拥有本土自产的零部件比解决能效问题更重要得多。

 

目前，中国的E级超算已经诞生一年了，美国至少现在还没有。

 

由于芯片制程上的优势，美国的E级超算会更节能，而且大概率在Frontier 和El Capitan和Aurora中诞生。但一切总还要等到诞生之后再说。

 

从某种意义上说，这种竞争会推动计算架构的进步。

 

话说回来，美国这几台怪兽级超算的问世，也就在不远的未来了。对面的「大招」如何放，我们拭目以待。

通信作者

 

陈文光

 

清华大学计算机系教授，主要研究领域为操作系统、程序设计语言与并行计算，国家自然科学基金杰出青年基金获得者。获国家科技进步二等奖一次，部级科技一等奖两次。

 

现为中国计算机学会杰出会员和杰出讲者，副秘书长，青年科技论坛荣誉委员；ACM中国理事会主席，ACM中国操作系统分会ChinaSys主席，ACM通讯中文版主编。

 

翟季冬

 

 

清华大学计算机科学与技术系副教授。曾任斯坦福大学客座教授（2015-2016），2013年MSRA（微软亚洲研究院）访问学者。现任ACM SIGHPC中国区秘书长。

 

他目前的研究兴趣包括并行计算、编译器、编程语言和性能评估。

 

他是清华学生集群团队的顾问。他带领的团队在SC、ISC 和 ASC的学生超级计算挑战中获得了11项国际冠军。

 

刘鑫

 

 

之江实验室智能超算研究中心研究员。国家并行计算机工程技术研究中心研究员，博士生导师，神威系列智能计算机常务副总设计师，长期从事超级计算机体系结构、并行算法及应用支撑软件的研究工作。

 

先后主持「神威·蓝光」、「神威·太湖之光」、神威E级原型机应用系统研制工作，负责核高基「面向数据中心（云平台）和集群计算的智能处理单元」软件系统研制，在科学计算和人工智能领域的并行算法、应用支撑软件和软硬件协同设计工作中成效显著。

 

带领团队研发的神威量子模拟器，入围2021年度「戈登·贝尔奖」。曾获省部级科技进步一等奖一项、二等奖三项、「清华大学-浪潮集团计算地球青年人才」，「CCF-IEEE CS青年科学家奖」。

 

杨红霞

 

美国杜克大学博士学位，现任阿里巴巴资深算法专家，带领团队开发基于计算平台和搜索推荐的智能算法，稳定的支持了阿里巴巴搜索、广告等30几个核心BU和其业务场景。

 

在顶级统计和机器学习国际学术期刊会议发表论文50余篇，美国专利9项，任职Applied Stochastic Models in Business and Industry副主编，International Statistical Institute理事等。

 

唐杰

 

 

清华大学计算机科学与技术系教授。ACM Fellow、IEEE Fellow。

 

研究兴趣包括人工智能、数据挖掘、社交网络、机器学习和知识图谱，重点是为信息和社交网络挖掘设计新算法。

 

曾获SIGKDD时间检验奖、第二届国家科学技术奖、 NSFC杰出青年学者、 英国皇家学会-牛顿高级奖学金和SIGKDD服务奖。 

参考资料：

https://www.theregister.com/2022/03/11/china_exascale_sunway_ai/

https://www.nextplatform.com/2021/10/26/china-has-already-reached-exascale-on-two-separate-systems/

https://www.nextplatform.com/2022/03/11/pondering-the-cpu-inside-chinas-sunway-oceanlight-supercomputer/