BaGuaLu: 基于国产超算的百万亿参数超大预训练模型训练方法 | 论文荐读

作者：

Zixuan Ma, Jiaao He, Jiezhong Qiu, Huanqi Cao, Yuanwei Wang, Zhenbo Sun, Liyan Zheng, Haojie Wang, Shizhi Tang, Tianyu Zheng, Junyang Lin, Guanyu Feng, Zeqiang Huang, Jie Gao, Aohan Zeng, JianWei Zhang, Runxin Zhong, Tianhui Shi, Sha Liu, Weimin Zheng, Jie Tang, Hongxia Yang, Xin Liu, Jidong Zhai, and Wenguang Chen. BAGUALU: Targeting Brain Scale Pretrained Models with over 37 Million Cores. In Proceedings of the 27th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming (PPoPP'22).

Paper: 

https://www.aminer.cn/pub/6228239b6750f804ca0b65f9

深度学习（Deep Learning）在不同的领域都有重要用途，包括计算机视觉（Computer Vision）、自然语言处理（Natural Language Processing）、推荐系统（Recommender System）等。近年来，基于 Transformer 结构的大规模预训练模型自然语言处理和计算机视觉领域的一系列下游任务中取得了最优的表现。从预训练模型的发展趋势来看，模型的规模与模型的效果呈现正相关的关系已经是学界的共识。随着预训练模型的规模的急剧增长，大规模模训练需要越来越大量的计算和内存能力，这无疑促进了人工智能和高性能计算的融合。然而，由于缺乏软硬件特性的协同设计，在超级计算机系统上高效训练大规模预训练模型仍然存在来自并行策略、数据存储、数据精度与负载均衡等多个方面的巨大挑战。如何在特定硬件架构下超大模型选取最高效的并行方式？在单精度下，万亿模型的参数高达 4 T，再加上模型的梯度、优化器的更新参数与计算中间值，总共需要 768 块 V100 才能存储，如此大的参数量如何进行高效划分？在训练中，更低的精度需要的内存量更小、性能更好，然而，低精度会带来模型准确度下降，如何选择合适的数据精度使得模型仍然能够正常收敛？此外，MoE 模型的专家存在严重的负载不均衡问题，如何改善这一现象提升训练效率？

 

针对上述问题，我们提出了首个基于国产超算的百万亿参数超大预训练模型训练系统  BaGuaLu。该系统拥有可以以超过 1 EFLOPS 的混合精度性能训练十万亿参数的模型，并且支持训练高达百万亿规模参数量模型的训练（174 T）。具体而言，BaGuaLu 系统在如下四个方面进行了创新

 

• 在并行策略方面，针对国产超算的网络拓扑结构，我们采用了混合数据并行混合数据与 MoE 并行的策略，在超节点内进行数据并行，在跨超节点做 MoE 并行，合理利用了网络带宽，提升了训练效率。

  
  

• 在参数存储方面，我们将优化器的参数拆分到数据并行组的每个节点中，更新时每个节点仅更新其存储的参数，并通过通信操作应用到数据并行组中的其他节点中，这样我们便实现了高效不重复的参数存储，将原本占用 75% 内存空间的优化器参数降低到了几乎可以忽略的程度。

  
  

• 在数据精度方面，我们采取了分层混合精度方法，针对发现不同层对精度的要求不一致的情况，我们对 FFN 层采用完全半精度，注意力层半精度计算、单精度更新，其它层用的是单精度，在保证模型收敛的情况下提升了训练性能。

  
  

• 在负载均衡方面，朴素的 MoE 模型存在负载不均衡的问题，我们提出名为 SWIPE 的负载均衡策略：如果某个专家的输入量超过某个阈值时，就把它过多的输入分给其它的专家，这样就完全消除了负载不均衡的问题。实验结果表明这一简单策略下模型训练仍然能够正常收敛。

  
  

我们基于 BaGuaLu 系统训练了三个多模态预训练模型，模型架构如下左图所示，训练数据来自最大的中文多模态数据集 M6-Corpus，参数规模分别为从万亿到百万亿（1.93 T、14.5 T、174 T）。如下右图所示：

实验结果显示，在 1.93 T 参数规模的时候，计算性能达到了1.18 EFLOPS，当模型增大到 14.5 T的时候，混合精度性能达到1.00  EFLOPS。这足以证明我们的训练系统的高效性。我们对 1.93 T 的模型进行了更长时间的训练，从上图中可以看到，模型表现出了明显的收敛趋势，且最终接近收敛。在强弱扩展性方面，我们分别根据节点规模增长扩展专家个数，或者随着节点规模的增长保持总 batch size 不变进行实验，实验结果表明 BaGuaLu 系统具有良好的强扩展性与接近线性的弱扩展性。

点击【阅读原文】查看paper