基于昇腾910B训练万亿参数的语言模型

   

Abstract

在本工作中，作者开发了一个系统，该系统在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]训练。但训练大型模型需高端集群，与成本效益目标相悖。此外，模型规模扩大增加部署成本，与成本最小化目标冲突。如何平衡各项挑战，是当前面临的关键问题。

面对训练阶段高昂的计算成本，作者追求实用高效的模型。为此，提出在多个领域训练模型，并借助持续学习模式，实现向任意领域扩展的潜力，虽受限于计算资源，但目标清晰，追求实际应用的高效性。

PanGu模型，坐拥万亿参数，在训练中融合多元数据。然而，部署时携带庞大参数对每应用并非必要，甚至难以实现。为此，该模型能灵活分组与分离参数，适配各种训练与部署需求，优势显著。

2.2.1 Overview

盘古运用自回归语言建模，搭载堆叠Transformer解码层与独特Query层。其架构灵活多变，M层全局共享，N层（含Query层）按需激活。每层RRE配备G组K位专家，每组专家数量各异。这种创新设计赋予盘古三种工作模式，展现其卓越性能与广泛适应性。

图1：盘古架构，融合密集与稀疏Transformer层。下部M层为跨领域共享密集层，上部N层通过随机路由专家（RRE）激活稀疏前馈部分。各领域标记独特嵌入，实现高效跨领域学习与处理，引领架构创新潮流。

• 混合模式：当，且时，模型包含稀疏RRE层和密集层。
• 密集模式：当或时，架构将简化为密集的盘古-模型。
• 稀疏模式：当且时，架构将是一个稀疏模型。

在万亿参数建模实践中，作者巧妙运用混合配置，将共享参数置于输入层附近，专家参数则置于输出层顶端，形成层次分明的结构。经过基准测试，此策略实现了最低的语言建模困惑度。作者认为，底层捕捉一般知识，而顶层则专注于具体知识的抽象学习。在标记嵌入层，作者根据领域差异选择不同嵌入矩阵，确保模型的高效性与专业性。

2.2.2 Random Routed Experts

在顶层结构中，作者运用多条件激活的前馈子层（专家）替代传统层，遵循专家混合（MoE）模式，实现高效模型设计。

在设计MoE架构时，如何有效将标记分配给专家是核心挑战。PanGu-采用的随机路由专家（RRE）机制，灵感源自哈希层，通过两级路由策略实现。首先，标记按域映射至候选专家组；随后，基于随机初始化的标记-专家路由图，从中挑选特定专家处理。此策略确保每层独立初始化映射，以实现计算资源的均衡分配。

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倍的吞吐量提升.

 

-对此，您有什么看法见解？-

-欢迎在评论区留言探讨和分享。-