在本工作中，作者开发了一个系统，该系统在Ascend 910 AI处理器集群和MindSpore框架上训练了一个万亿参数的语言模型，并提出了一个含有1.085T参数的语言模型，名为PanGu-。从PanGu-[1]继承了参数，作者将密集的Transformer模型扩展为稀疏模型，采用_Random Routed Experts_（RRE）方法，并通过使用_Expert Computation and Storage Separation_（ECSS）高效地在329B个 Token 上训练模型。这导致了通过异构计算使训练吞吐量提高了6.3倍。作者的实验发现表明，PanGu-在多种中文NLP下游任务的零样本学习中提供了最先进的表现。此外，在开放领域对话、问答、机器翻译和代码生成应用数据上的微调也展示了其强大的能力。
1 Introduction
大型语言模型（LLM）[2, 3, 1, 4, 5, 6, 7, 8, 9, 10等]在自然语言理解、生成和推理等领域展现了前所未有的能力和潜力。通过利用大量文本数据，语言模型的性能随着计算预算和模型参数的增加而提升，展示了强大的零样本/少样本学习能力甚至涌现能力[4, 11]。自GPT-3[2]以来，已有多个拥有数百亿参数的大型语言模型发布，包括但不限于Megatron-Turing NLG[12]、PanGu-[1]、ERNIE 3.0 Titan[8]、Gopher[5]、PalLM[4]、OPT[6]、Bloom[10]和GLM-130B[9]。研究行人开始构建拥有超过一万亿参数的更大语言模型。这通常是通过利用稀疏激活模型，如Mixture-of-Experts（MoE）[13]来实现的。在现有的万亿参数模型中，有几项值得注意的工作，如Switch-C[14]、GLaM[15]、MoE-1.1T[16]、Wu Dao 2.0[17]和M6-10T[18]。然而，只有少数发布了在广泛任务上的全面评估结果并同时实现了预期的性能。根据作者的经验，主要的困难在于扩展效率。
最近关于语言模型扩展规律的研究[19, 20, 21]表明，为了达到最佳性能，有必要使用足够的训练数据和相应的计算预算来训练LLM。因此，这项工作的主要动机之一是设计一个可扩展的模型架构和一个高效的分布式训练系统，能够以高训练吞吐量消费数据。
模型扩展。大型语言模型的性能预期会随着模型大小的增加而提升。与训练密集型Transformer模型的昂贵计算成本相比，如Mixture-of-Experts (MoE) [13, 14, 15, 22]这样的稀疏架构被认为是一个吸引人的选择，可以在不引起计算成本线性增长的情况下扩展模型大小。然而，MoE模型存在诸如工作负载不平衡和全局通信延迟等问题。此外，如何将现有密集模型与MoE结合，以及在每个层级分配多少专家仍然是待解决的问题。因此，设计一个既具有高性能又高效训练的万亿参数稀疏模型是一个重要但具有挑战性的任务。 系统扩展。如DeepSpeed 4等框架已被提出以支持训练万亿参数模型。实际上，主要的障碍往往在于有限的计算预算，或者更具体地说，能够使用的加速设备（例如GPU、NPU、TPU）的数量。通过利用张量并行性[23]、流水线并行性[24]、零冗余优化器[25]和重物质化[26]等技术，实践者可以在数千个加速设备上以可行的批量大小的训练万亿参数模型。另外，实践者可以通过利用异构计算技术，例如将部分计算卸载到宿主设备[27]，来减少计算资源的数量。然而，当前技术不可避免地因宿主设备与设备间缓慢的带宽以及与加速设备相比CPU较弱的计算能力而阻碍了训练吞吐量，这阻止了为大型语言模型提供合理数量的数据以及实现最佳性能。因此，如何在有限的计算预算下高效地扩展系统性能对大型语言模型的性能至关重要。
在本工作中，作者介绍了PanGu-，一个拥有1.085万亿参数的稀疏架构大型语言模型。作者在MindSpore 5框架下开发PanGu-模型，并在仅包含512个Ascend 910 AI加速器的集群上对其进行训练[28]，在100天内使用了3290亿个 Token 。PanGu-从PanGu-[1]继承了参数，并采用Transformer解码器架构，通过随机路由专家（RRE）进行扩展。与传统的MoE不同，RRE采用两级路由。在第一级，专家按领域或任务分组；在第二级， Token 随机且均匀地映射到每个组中的专家，而不使用MoE中的任何可学习门控函数。通过RRE的设计，可以轻松地从PanGu-中提取子模型，用于各种下游应用，如对话、翻译、代码生成或一般的自然语言理解。为了使训练系统高效且可扩展，作者提出了专家计算与存储分离（ECSS）机制，该机制在512个Ascend 910加速器的集群上训练1.085万亿的PanGu-时，实现了69905 tokens/s的观察吞吐量，并且大幅减少了主机到设备和设备到主机的通信以及优化器更新计算。总体而言，与具有相同超参数但采用MoE架构的模型相比，训练吞吐量提高了6.3倍。通过消耗超过40种自然语言和编程语言的3290亿个 Token ，PanGu-在中文领域的子模型在零样本设置下，没有进行任何多任务微调或指令调整，就在六个类别的16个下游任务上显著优于包括拥有13B参数的PanGu-和拥有260B参数的ERNIE 3.0 Titan[8]在内的先前SOTA模型。作者还对PanGu-在对话、机器翻译和代码生成等应用领域进行了微调后的性能进行了测试。PanGu-在相应领域优于SOTA模型。
技术报告的其余部分组织如下。第2部分介绍PanGu-模型的设计理念和架构。第3部分介绍数据集的收集和组织。第4部分描述系统设计和加速技术。第5部分呈现PanGu-模型的实验结果。
Design Principles
PanGu-旨在实现以下目标：
性能：在多个领域和任务上达到最先进的NLP性能。
效率：在适度的集群上以最大系统性能训练万亿参数模型。
可用性：可扩展到各种领域或任务，无需从零开始重新训练模型。
部署：易于定制，并在各种实际环境中部署。
同时实现上述所有目标极具挑战性。考虑到第一个目标，一个能够在多个领域泛化和表现良好的语言模型应具有非常大的参数数量，并根据规模法则[19, 20, 21]在大数据上进行训练。然而，训练如此大的模型也意味着需要一个高端集群，这在某种程度上与第二个目标相矛盾。而且模型规模的扩大也导致部署训练后的模型成本增加，这与第四个目标有关。
考虑到训练阶段产生的高计算成本，作者希望得到的模型在实际应用中既可用又高效。鉴于这一目标，作者 Proposal 在多个领域训练模型，并使其在持续学习模式下进一步可扩展到任意数量的领域，受限于计算资源。
在训练阶段，万亿参数的PanGu-模型接收来自多个领域的数据。然而，在部署阶段，对于每个应用都托管万亿参数模型往往是不必要的，甚至是不可能的。因此，一个允许根据各种训练和部署设置对其参数进行分组和分离的模型具有显著优势。
2.2.1 Overview
盘古-采用了自回归语言建模，具有堆叠的 Transformer 解码层和一个顶部的 Query 层。盘古-架构提供了灵活的设计。底部的M层在所有领域全局共享，顶部的N层（包括 Query 层）根据输入数据的领域稀疏激活。在每个RRE层中，总共有G组中的K个专家，每组中的专家数量可以不同。这种灵活的设计提供了三种模式。
图1：盘古-架构。该架构由密集的 Transformer 层和稀疏的 Transformer 层混合组成。下部的M层是跨不同领域共享的密集层。上部的N个 Transformer 层的前馈部分通过随机路由专家（RRE）稀疏激活。来自不同领域的标记具有不同的嵌入。
混合模式：当，且时，模型包含稀疏RRE层和密集层。 密集模式：当或时，架构将简化为密集的盘古-模型。 稀疏模式：当且时，架构将是一个稀疏模型。
在这个万亿参数建模实践中，作者使用混合配置，将共享参数放置在靠近输入层（底部），所有稀疏激活的专家参数放置在靠近输出层（顶部）。在模型设计阶段，作者在较小规模的模型上对各种专家放置策略进行了基准测试，所选策略获得了最低的语言建模困惑度。作者的假设是底层倾向于学习一般知识，而具体知识在更高层次的抽象中，更适合由顶层学习。在标记嵌入层，作者选择为不同领域使用不同的嵌入矩阵。
2.2.2 Random Routed Experts
在顶部N层中，作者用多个条件激活的前馈子层（专家）替换每个前馈子层，遵循专家混合（MoE）范式。
在设计MoE架构时，一个关键问题是如何将标记路由到专家。对于PanGu-，作者提出了一个随机路由专家（RRE）机制，该机制受[29]中提出的哈希层的启发。具体来说，RRE通过ID以两级方式路由标记。在第一级，标记通过域映射到一组候选专家，然后在第二级，根据标记-专家路由图从这组专家中选择一个专家来处理标记。路由图是随机初始化的，每层都有独立初始化的映射以平衡计算。
RRE与常用的可学习路由器相比具有几个优点。
图2：PanGu-中的随机路由专家（RRE）。标记首先通过域路由到一个专家组，然后在该域中随机路由到其中一个专家。模型中没有可学习的路由器。
在训练期间，PanGu-允许在不影响其他专家的情况下，增加、修改或移除特定域的专家。这个属性使得PanGu-在解决常见的大灾难性遗忘问题上非常灵活，这对于终身学习或连续学习至关重要。* 在大多数实际部署环境中，部署具有一万亿个参数的模型是不必要或不可行的。PanGu-允许根据实际需求提取特定域的子模型，并且只部署子模型。子模型可能包含数十亿个参数，但仍然保持原始模型在目标域上的预测能力。使用这种提取和部署操作，作者可以轻松地为多个工业应用部署模型。 所有传统的MoE模型都依赖于全局通信集体操作，以在不同设备上居住的专家之间移动数据。由于作者提出的两级路由，不同域的专家不交换标记，全局通信限制在每个域内进行。因此，昂贵全局全局操作减少到了分组全局操作，节省了大量通信量并减少了端到端训练延迟。 可学习路由器需要更多的计算，并且可能遭受专家之间的负载不平衡问题，这通常使得训练过程变得不稳定。RRE避免了以上所有问题，因为不需要引入额外的参数，随机初始化的路由表有助于在专家上平衡负载。
RRE需要一个在预训练之前初始化的路由图，算法1描述了作者如何构建路由表。
算法1随机路由专家（RRE）机制中的路由表构建过程。
Collection
为了更好地展示PanGu-模型能够高效且独立地从多个领域学习的能力，作者收集了40个领域的数据集，其中四个主要领域（中文、英文、中英双语和代码）拥有大量数据。其余领域的数据量较小，包括26种其他单语自然语言、6种编程语言以及来自金融、健康、法律和诗歌领域文本数据。
对于中文文本，作者收集了WuDaoCorpora 2.0 [30]（包含200GB）和CLUECorpus2020 [31]（包含100GB）。对于英文文本，作者收集了Pile数据集[32]（包含800GB）和C4数据集[3]（包含750GB）。对于代码，作者使用了在PanGu-Coder [33]中使用的Python代码（147GB），以及来自GHTorrent [34]的Java代码（161GB），这些代码通过文件大小（小于1MB）、每行平均字符数（小于200）、每行最大字符数（小于1000）和它们的可编译性进行了过滤。然后，这些收集的英文、中文和代码文本数据被采样并分配到四个主要领域。最终，作者在四个主要领域获得了超过3000亿个标记。四个主要领域的数据分布和数据源的详细统计在表1中呈现。
对于剩余的36个领域，26个单语领域的数据主要来自CCAligned [35]和CCMatrix [36]。与上述代码领域类似，6个编程语言领域的数据通过GHTorrent [34]收集，并通过类似的方式进行过滤。金融领域的数据通过标签从WuDaoCorpora 2.0 [30]中过滤得到。健康领域的数据来自Chinese MedDialog Dataset [37]。法律领域的数据是从CAIL2018 [38]中采样的。诗歌领域的数据集来自Wernero-Poetery 6。最终，作者为这36个领域采样了超过250亿个标记。
Format
对于四个主要领域，每个领域都可以适应不同的下游任务。为了更好地支持特定领域的下游任务，本文针对不同领域采用不同的数据格式。对于中文和英文领域，在每个训练样本的末尾插入标记，以表示训练文本的结束。
对于双语领域，根据训练样本的来源（来自中文数据集或英文数据集），在训练样本的开头插入或标记，并在每个训练样本的末尾插入标记。
对于代码领域，根据训练样本的编程语言类型，在训练样本的开头插入或标记，并在每个训练样本的末尾插入标记。
对于其余36个领域，26个单语领域的金融、健康、法律和诗歌领域的数据格式与中文和英文领域相同，6个编程语言领域的的数据格式与代码领域相同。
对于一个格式化数据集，假设它包含n个训练样本。为了充分利用Ascend 910集群的计算能力，并在预训练阶段加速训练，作者将数据集中的所有样本拼接成一个序列，然后根据固定长度（1024）从拼接的序列中截取训练实例，如图6所示。在微调阶段，对于格式化数据集中的每个训练样本，如果长度小于固定长度，作者使用特殊标记将样本填充到固定长度。如果长度大于固定长度，则截断多余的部分。图7展示了这个过程。与PanGu-模型不同，PanGu-模型的每个训练样本包含两个字段：训练实例的标记ID输入序列及其领域ID。领域ID表示训练实例属于哪个领域。PanGu-模型的RRE层通过领域ID决定将训练标记路由到哪个专家。
4 System
盘古-Σ是使用MindSpore 1.6框架7实现的，并在512个Ascend 910加速器（也称为Ascend 910 NPU）上进行训练。训练一个万亿参数的语言模型提出了多个挑战。首先，它在训练过程中需要巨大的内存。尽管稀疏架构可以有效节省计算量，但它并没有减少内存消耗，作者仍然需要在加速器内存中存储所有参数和优化状态。假设使用Adam优化器[39]和混合精度训练[40]，一个1T模型通常仅用于参数、梯度和优化器状态就需要总共16TB的内存。在训练过程中，模型还需要额外的内存用于输入数据、网络激活、通信缓冲区和临时变量。作者估计，使用合理的批处理大小训练一个拥有1万亿参数的盘古-Σ模型需要超过32TB的内存，并需要超过1000个Ascend 910加速器或配备32GB高带宽内存（HBM）的NVIDIA V100 GPU。
作者没有将大量硬件资源用于扩大模型规模，而是致力于使用512个Ascend加速器合理大小的集群来训练盘古-Σ。为此，作者采用了异构训练，并将优化器状态卸载到CPU[27]。启用异构训练后，所有优化器状态从加速器移动到具有750GB主机内存和鲲鹏920 CPU 8的主机，作者可以将整个训练过程放入集群中。其次，在启用原始优化器卸载后，系统吞吐量是不可接受的。根本原因仍然是大量的参数。梯度和更新的参数需要通过缓慢的主机到设备和设备到主机的通信进行交换，CPU需要遍历所有参数并更新它们。为了提高训练吞吐量，作者利用了盘古-Σ架构的稀疏特性。由于盘古-Σ使用稀疏架构，其大部分参数是条件激活的，优化器在一个迭代中只需要更新部分专家。因此，作者提出了图8所示的专家计算与存储分离（ECSS）方法。
在专家计算与存储分离中，作者将专家视为知识数据库，用于存储不同任务或领域的特定知识。在每次迭代中，具有特定领域的不同 Token ID会稀疏激活专家。在MindSpore中，作者使用查找操作符来选择被激活的专家的部分，并在反向计算中稀疏更新它们的参数。在优化器CPU卸载计算中，MindSpore将FP16参数从主机CPU复制到NPU，在NPU上计算梯度，将FP16梯度从NPU移动到CPU，并在主机CPU上计算优化器状态和更新参数。当专家稀疏比率较低，如时，计算成本仅为全模型的大约10%。除了采用Ascend-KunPeng稀疏异构计算的ECSS之外，作者还采用了MindSpore和CANN 9提供的其他并行训练和加速技术。作者对所有的注意力和前馈层使用8路模型并行，64路专家并行（不复制），以及非专家部分使用64路数据并行。为了进一步优化内存占用，还采用了重新物质化[26]和优化器并行[25]来减少峰值内存消耗。作者还使用FastGelu和融合的LayerNorm来加速逐点计算。通过结合所有这些技术，与原始的PanGu-异构训练相比，作者实现了6.3倍的吞吐量提升.
参考
[1].PanGu-: Towards Trillion Parameter Language Model with Sparse Heterogeneous Computing.
