深度学习在不同的领域都有重要用途，包括计算机视觉、自然语言处理、推荐系统等。近年来，基于Transformer结构的大规模预训练模型在自然语言处理和计算机视觉领域的一系列下游任务中取得了最优的表现。从预训练模型的发展趋势来看，模型的规模与模型的效果呈现正相关的关系已经是学界的共识。随着预训练模型的规模的急剧增长，大规模模训练需要越来越大量的计算和内存能力，这无疑促进了人工智能和高性能计算的融合。
然而，由于缺乏软硬件特性的协同设计，在超级计算机系统上高效训练大规模预训练模型仍然存在来自并行策略、数据存储、数据精度与负载均衡等多个方面的巨大挑战。
针对上述问题，我们提出了首个基于国产超算的百万亿参数超大预训练模型训练系统BaGuaLu。该系统拥有可以以超过1 EFLOPS的混合精度性能训练十万亿参数的模型，并且支持训练高达百万亿规模参数量模型的训练。
具体而言，BaGuaLu系统在并行策略、参数存储、数据精度和负载均衡四个方面进行了创新。在并行策略方面，针对国产超算的网络拓扑结构，我们采用了混合数据并行混合数据与MoE并行的策略，在超节点内进行数据并行，在跨超节点做MoE并行，合理利用了网络带宽，提升了训练效率。
在参数存储方面，我们将优化器的参数拆分到数据并行组的每个节点中，更新时每个节点仅更新其存储的参数，并通过通信操作应用到数据并行组中的其他节点中，这样我们便实现了高效不重复的参数存储，将原本占用75%内存空间的优化器参数降低到了几乎可以忽略的程度。
在数据精度方面，我们采取了分层混合精度方法，针对发现不同层对精度的要求不一致的情况，我们对FFN层采用完全半精度，注意力层半精度计算、单精度更新，其它层用的是单精度，在保证模型收敛的情况下提升了训练性能。在负载均衡方面，朴素的MoE模型存在负载不均衡的问题，我们提出名为SWIPE的负载均衡策略：如果某个专家的输入量超过某个阈值时，就把它过多的输入分给其它的专家，这样就完全消除了负载不均衡的问题。
实验结果表明这一简单策略下模型训练仍然能够正常收敛。
我们基于BaGuaLu系统训练了三个多模态预训练模型，模型架构如下左图所示，训练数据来自最大的中文多模态数据集M6-Corpus，参数规模分别为从万亿到百万亿。实验结果显示，在1.93 T参数规模的时候，计算性能达到了1.18 EFLOPS，当模型增大到14.5 T的时候，混合精度性能达到1.00 EFLOPS。
这足以证明我们的训练系统的高效性。我们对1.93 T的模型进行了更长时间的训练，从上图中可以看到，模型表现出了明显的收敛趋势，且最终接近收敛。在强弱扩展性方面，我们分别根据节点规模增长扩展专家个数，或者随着节点规模的增长保持总batch size不变进行实验，实验结果表明BaGuaLu系统具有良好的强扩展性与接近线性的弱扩展性。