浪潮人工智能研究院

　　从 2018 年的 BERT 到 2020 年的 GPT-3，NLP 语言模型经历了爆发式的发展过程，其中 BERT 模型的参数量为 3.4 亿，而 GPT-3 的模型参数量达到了 1750 亿。2021 年 9 月，浪潮人工智能研究院发布了“源 1.0”，它是目前规模最大的中文 AI 单体模型，参数规模高达 2457 亿，训练采用的中文数据集达 5TB。“源 1.0”在语言智能方面表现优异，获得中文语言理解评测基准 CLUE 榜单的零样本学习和小样本学习两类总榜冠军。测试结果显示，人群能够准确分辨人与“源 1.0”作品差别的成功率低于 50%。

　　海量的参数带来了模型训练和部署上的巨大挑战。本文将聚焦“源 1.0”背后的计算挑战以及我们采取的训练方法。

　　1．“源 1.0”的模型结构

　　“源 1.0”是一个典型的语言模型。语言模型通俗来讲就是能够完成自然语言理解或者生成文本的神经网络模型。对于“源 1.0”，我们考虑语言模型（Language Model，LM）和前缀语言模型（Prefix Language Model，PLM）两种模型结构。如下图所示：

　　我们比较了 130 亿参数的 LM 和 PLM 在不同下游任务上的结果，注意到 LM 在 Zero-Shot 和 Few-Shot 上表现更好，而 PLM 在微调方面表现出色。微调通常会在大多数任务中带来更好的准确性，然而微调会消耗大量的计算资源，这是不经济的。所以我们选择 LM 作为“源 1.0”模型的基础模型结构。

　　2. 如何训练“源 1.0”

　　2. 1. 源 1.0 训练面对的挑战

　　“源 1.0”的训练需要面对的第一个挑战就是数据和计算量的挑战。

　　数据方面，如果把训练一个巨量模型的训练过程比作上异常战役的话，那么数据就是我们的弹药。数据量的多少，决定了我们可以训练模型的规模，以及最后的效果。针对这一方面，我们构建了一个全新的中文语料库，清洗后的高质量数据规模达到了 5TB，是目前规模最大的中文语料库。

　　算力方面，根据 OpenAI 提出的 PetaFlop/s-day 衡量标准，我们可以估算“源 1.0”训练的计算需求情况。根据 Wikipedia 提供的数据（https://en.wikipedia.org/wiki/OpenAI），GPT-3 的计算需求约为 3640 PetaFlop/s-day， 约等于 64 个 A100 GPU 训练 1 年时间。而“源 1.0”的计算需求达到了 4095 PetaFlop/s-day。

　　计算资源的巨大开销是限制研究人员研发具有数以千万计参数的 NLP 大模型的瓶颈。例如 GPT-3 是在由 10000 个 GPU 所组成的集群上训练得到的。我们在设计“源 1.0”的模型结构时，考虑到了影响大规模分布式训练的关键因素，采用了专门的分布式训练策略，从而加速了模型的训练过程。

　　在模型训练时一般最常用的是采用数据并行分布式计算策略，但这只能满足小模型的训练需求。对于巨量模型来说，由于其模型参数量过大，远远超过常用计算设备比如 GPU 卡的显存容量，因此需要专门的算法设计来解决巨量模型训练的显存占用问题，同时还需要兼顾训练过程中的 GPU 计算性能的利用率。

　　2. 2“源1. 0”的训练策略

为了解决显存不足的问题，我们采用了张量并行、流水并行、数据并行相结合的并行策略，实现了在 2128 个 GPU 上部署“源 1.0”，并完成了 1800 亿 tokens 的训练。

　　一、张量并行

　　针对单个 GPU 设备不能完整的承载模型训练，一个解决方案就是张量并行+数据并行的 2D 并行策略。具体来说，使用多个 GPU 设备为 1 组，比如单个服务器内的 8 个 GPU 为 1 组，组内使用张量并行策略对模型进行拆分，组间（服务器间）采用数据并行。

对于张量并行部分，NVIDIA 在 Megatron-LM 中提出了针对 Transformer 结构的张量并行解决方案。其思路是把每一个 block 的参数和计算都均匀的拆分到N个 GPU 设备上，从而实现每个 GPU 设备都承担这一 block 的参数量和计算量的1/N效果。图 3 展示了对 Transformer 结构中的 MLP 层和 self-attention 层进行张量并行拆分计算的过程示意图。

　　在训练过程中，tensor 经过每一层的时候，计算量与通信数据量之比如下：

　　其中，S为输入序列的长度，h为隐藏层的大小（hidden size）。

　　二、流水并行

　　对于具有数千亿参数的语言模型，这些参数很难被存放在单个节点中。流水线并行将 LM 的层序列在多个节点之间进行分割，以解决存储空间不足的问题，如图 5 所示。每个节点都是流水线中的一个阶段，它接受前一阶段的输出并将结果过发送到下一阶段。如果前一个相邻节点的输出尚未就绪，则当前节点将处于空闲状态。节点的空闲时间被称为流水线气泡（pipline bubble）。为了提高流水行并行的性能，我们必须尽可能减少在气泡上花费的时间。定义流水线中气泡的理想时间占比为如下形式：

　　根据这一公式，流水线气泡的耗时随着层数L的增加而增加，随着微批次大小（micro-batch-size）的增加而减小。当的时候，流水并行过程中的流水线气泡对训练性能的影响几乎可以忽略。

　　与此同时，在流水并行过程中，节点间的计算量与通信数据量之比为：

　　根据上面的公式，流水线中节点的计算效率与h和S呈线性关系，这与张量并行类似。

　　三、数据并行

　　采用数据并行时，全局批次大小（global batch size）按照流水线分组进行分割。每个流水线组都包含模型的一个副本，数据在组内按照局部批次规模送入模型副本。数据并行时的计算量与通信数据量的比值可用如下公式近似：

　　当时，上面公式可以进一步简化成：

　　根据这一公式，我们可以看出数据并行的计算效率与全局批次大小B和序列长度S呈正比关系。由于模型对内存的需求与S的平方成正比，与B成线性关系，因此增加全局批次大小可以更有效的提升数据并行的效率。

　　当全局批次大小过大的时候，模型很容易出现不收敛的问题，为了保证模型训练过程的稳定性，我们将全局批次大小限制在了个 token 内。

　　根据以上的理论分析，我们确定了设计“源 1.0”巨量模型结构的基本原则：

　　· 尽可能增加序列长度，因为它有利于张量并行、流水线并行和数据并行。由于内存占用与序列长度的平方成正比，因此有必要在反向传播时重新计算激活函数，以节省内存开销。

　　· 语言模型中层数太多会对性能产生负面影响，因为这会增加在流水线气泡上的时间消耗。

　　· 增加隐藏层大小可以提高张量并行和流水线并行的性能。

　　· 增加节点中的微批次大小可以提高流水线并行效率，增加全局批次大小可以提升数据并行的效率。

　　在这一设计原则的基础上，我们设计的“源 1.0”的模型结构以及分布式策略的设置如下表所示：

　　结合模型结构的特性以及我们使用集群的硬件特性，我们如下的节点配置和分布式策略选择：

　　· “源 1.0”模型在训练过程中共使用了 2128 个 GPU

　　· 模型分成了 7 组，每组 38 台 AI 服务器，里面放置一个完整的“源 1.0”模型，7 组之间采用数据并行

　　· 每组的 38 个服务器，采用流水并行每个服务器放置1/38 的模型（2 个 Transformer Layer），一共 76 层

　　· 在每台服务器内采用张量并行，按照 Transformer 结构的每一层进行均匀切分

　　模型收敛曲线如下图：

      关于“源 1.0”的更多信息，大家可以参照浪潮人工智能研究院发布在 arxiv 上的论文：https://arxiv.org/abs/2110.04725。

　　另外，我们也整理了更详细的《NLP 模型训练解决方案白皮书》。感兴趣的读者可以关注“浪潮 AIHPC”公众号，回复“NLP 白皮书”即可下载。