利兹大学徐洁：面向大规模深度神经网络的分布式训练系统 | CCCF精选

大家好，很荣幸受邀参加今年的中国计算机大会（CNCC2023），并在此进行特邀报告，也很高兴能与众多国内外计算机领域的研究者、从业者、兴趣爱好者一起分享关于分布式系统的一得之见。当前，大模型的发展如火如荼，我们已经在各种研究中听到了很多关于大模型的算法及其在各种场景下的应用。今天，我想要和大家分享讨论的是关于大模型训练框架层面的系统软件的研究与最新进展。

在过去的二十多年里，以数据为驱动力的机器学习和深度学习技术取得了显著的发展。2022年11月，ChatGPT凭借其卓越表现崭露头角，随着其用户群体的不断扩大，大规模语言模型(Large Language Model，LLM，简称大模型)以前所未有的速度迅速崛起，推动了通用人工智能技术在各行各业的应用，并协助科学研究取得重大突破。大模型不仅在处理规则化和重复性的任务上表现出色，如开放性问答、角色扮演、多轮对话、图片或视频生成、自动驾驶等，还具有发掘特定领域中新知识的潜力，同时帮助科学研究者打破固定思维，加速新一轮的技术创新与工业革命。可以毫不夸张地说，继万物相连的移动互联网之后，大模型将成为下一代以数据为中心的通用人工智能技术发展的转折点与里程碑，推动社会进入下一个信息技术发展浪潮。

传统的机器学习算法，如线性回归、聚类等，通常依赖于人工从特定任务的数据集中抽取特定特征构建模型，以解决特定的下游任务。然而，大模型的工作机制与此不同。在这里，我们主要关注两个重要方面：首先，大模型是基于深度神经网络(Deep Neural Network，DNN)的架构，采用端到端的学习方式进行数据建模。这意味着它能自动构建数据特征，无需人工介入。其次，大模型被视为一种通用的基础模型，它构建于海量数据集之上，涵盖了几乎所有类别的互联网语料数据。这使得大模型能够实现任务无关的通用表征能力，从而在各类下游任务中表现出色。

大模型之所以能够展现出令人惊艳的学习能力，主要依赖于两个“大”：第一是模型参数规模大。目前主流的大模型，如Meta LLaMA，其参数规模达到十亿至百亿。深度学习领域内的研究显示，所谓大模型的涌现能力，是当模型参数增加到一定程度后突然表现出的前所未有的能力。因此，在百亿参数之后继续扩大模型规模，是否会再次引发模型性能的质变？此问题值得探索和期待。但就目前研究进展来看，大模型领域尚未探索到Transformer机制的能力上限，模型参数量规模仍有继续增大的趋势。第二“大”是训练数据集规模之大，通常达到万亿词元(tokens)。如何高效地完成大模型的训练，不仅给数据和算法研究者带来了新的挑战，同时也（或者更多地）给机器学习框架层的系统软件带来新的考验。我们仍需要不断优化训练方法和相关技术，以便更经济有效地利用大规模数据完成大模型的训练。

图1 模型参数量变化与硬件存储容量的关系

硬件配置的升级速度与作业计算量的激增速度之间的失衡，极大地增加了模型开发成本。我们总结了近五年主流大模型的模型参数量和硬件存储单元容量的变化，如图1所示。硬件主要包括英伟达GPU和谷歌TPU。从图1可以看出，在过去五年时间里，模型参数规模增长超过了100倍，但硬件存储单元容量仅增加了约6.7倍，这种巨大的增速差异导致模型训练时依赖的资源数大幅增加。目前，从零开始完整训练一次大模型的成本开销少则十万美元，多则千万美元。如此昂贵的开销已经远远超过了大多数研究者所能承受的范围。

接下来，我从分布式系统软件的角度出发，分析如何解决大模型在训练阶段所面临的问题与挑战，以及我们的研究成果。

随着模型参数量和训练数据集呈指数级增长，模型训练所需的存储资源也呈指数扩增。这使得硬件存储资源，如GPU内置内存，成为了限制模型扩增规模的主要性能瓶颈。因此，我们需要采取特定的并行策略，从系统软件层面对此进行支持。

我们将现有的分布式模型训练方法归纳为水平扩展和垂直扩展两大类别。水平扩展是指在数量维度对同构硬件进行扩展，例如：通过购置更多GPU增加系统整体的可用内存容量，从而支持更多参数和更大数据集的模型训练。垂直扩展是指对异构资源进行融合，即利用其他硬件资源进行时间与空间的置换，以此缓解GPU的内存压力。

水平扩展的常见策略包括数据并行(Data Parallelism，DP)、张量并行(Tensor Parallelism，TP)、流水线并行(Pipeline Parallelism，PP)和序列并行(Sequence Parallelism，SP)等。在此，我们简要同步一些名词的含义：Batch是指优化器在进行一次参数更新前所处理的样本数；Mini Batch是指当前一次输入的样本数。

数据并行

由于设备内存容量的限制，我们通常无法一次性加载整个Batch样本数据。目前主流的机器学习框架支持将Batch切分为多个独立的Mini Batch，然后以梯度累积的方式串行或并行执行，以此实现Batch的训练效果。然而，串行执行Mini Batch的方式会降低模型训练的吞吐率，因此数据并行应运而生。

数据并行将训练样本在Batch维度进行切分，并采用并行执行的方式对多个Mini Batch同时进行模型训练。当使用数据并行时，每组模型独立维护一份完整的模型参数、梯度等。每份Mini Batch作为输入被发送到一份模型副本上进行梯度求解。在优化器更新参数前，采用参数服务器或分布式通信算子库进行梯度同步，其核心逻辑表示为all_reduce(grads, group=dp_group)，如图2所示。

图2 梯度同步的两种方式

数据并行策略有效地解决了数据集大的问题，同时也是提高系统吞吐率最直接有效的方式。用户仅需要扩增设备数量即可近线性地提升训练吞吐率。随着机器学习框架的不断发展，数据并行已经成为标准功能，且用户仅需要修改一行代码，即可将数据并行策略应用于各类模型的训练过程中。然而，数据并行策略也存在潜在问题，即模型副本的冗余存储，这意味着随着设备数量的增加，冗余数据导致存储效率越来越低。因此在性能和成本投入之间，研究者需要做好取舍平衡。

张量并行

模型并行主要解决模型参数规模过大的问题。其包含两种切分策略，即水平切分（层内张量切分）和垂直切分（以层为单位的层间切分），两种切分策略相互正交。参与模型并行的实例共同维护一整份模型副本。

图3 张量并行

张量并行策略将层内张量进行切分并存储在不同设备上，通过引入通信语义，协同完成相应的张量计算，如图3所示。以三维张量运算为例，在进行A^b×m×nB^1×n×k=C^b×m×k的矩阵运算时，可以将张量A^b×m×n切分为（A^b1×m×n，A^b2×m×n），并分别存储在设备1和设备2中；矩阵B则被复制两份并分别存储于设备1和设备2中。当执行乘法运算时，每个设备独立完成计算量的一部分，最后再将各设备的输出结果进行拼接，即A^b×m×nB^1×n×k=（A^b1×m×n，A^b2×m×n）B^1×n×k=（A^b1×m×nB^1×n×k，A^b2×m×nB^1×n×k）=（C ^b1×m×k，C ^b2×m×k）=C ^b×m×k。此样例中矩阵B产生了冗余存储，当矩阵B的规模增长到也需要切分时，感兴趣的读者可以推导一下切分规则和所需的通信元语。张量并行策略可以增加可训练模型的宽度。

流水线并行

流水线并行策略是一种基于层间切分的模型并行策略，其核心思想是将神经网络模型按层进行切分，每个设备仅存储并维护部分层的参数。通过层间通信，在进行前向或反向传播过程中，将上一层的输出传递给下一层作为输入，如图4所示。这种策略允许增加可训练模型的层数，进一步拓展了模型深度。配合张量并行策略，用户可以更加灵活地扩展和优化模型结构。

图4 流水线并行

尽管流水线并行策略具有诸多优势，但仍然存在一些潜在的问题。在模型训练过程的初期或后期，下一阶段的运算需要等待上一阶段的任务完成才可以开始，导致在资源利用方面出现空闲，这些空闲间隙被形象地称为“气泡”（bubbles）。气泡的数量是评测流水线并行策略性能的重要指标，已经有许多工作和研究致力于通过任务编排减少气泡的数量，进一步提高训练效率。

序列并行

随着时序模型架构从LSTM演变为Transformer，数据处理方式也从串行转至并行，极大地增加了模型训练的吞吐率。然而，自注意力机制在计算注意力分数时，需要对输入的词元进行两两分数计算，即临时矩阵大小与K²成正比，其中K为输入序列中词元的数量。换句话说，临时矩阵的大小会随着输入序列长度的增加而呈指数增长，极大地消耗了内存存储空间。因此，序列并行策略旨在突破机器学习框架对输入序列长度的限制。序列并行策略将输入序列切分为多个子序列，每个子序列被输入至不同的设备上进行处理，类似数据并行。不同之处在于序列并行需要通过引入通信语义协同完成注意力分数的计算，在逻辑层面保证长序列注意力计算的有效性，其核心逻辑如图5所示。然而目前主流机器学习框架对此功能的支持还不够完善，仍存在一些改进空间。

图5 序列并行

序列并行与其他并行策略，如数据并行、张量并行，流水线并行等相互正交，共同实现在不同维度上的扩展能力，使分布式机器学习框架能够满足各类算法研究的需求。在硬件资源充足的情况下，这些并行策略理论上都可以实现无限扩展。

其他——专家并行

随着模型参数量的增加，模型训练成本呈现出指数增长趋势。为了降低开销，算法研究人员试图利用稀疏化结构来扩增模型参数规模，在增加参数量的同时，降低资源消耗。混合专家模型(Mixture of Experts，MoE)是其中一种稀疏模型架构，即在模型结构中，引入专家模块，利用一组专家代替单一组件，以此提升模型的通用性。例如Switch Transformers将多层感知器(Multi-Layer Perceptron，MLP)视为一个专家，使用数个MLP替换单一的MLP扩增模型参数规模。

随着专家数量的增加，单个GPU内存容量无法一次性加载所有的专家模块。因此，专家并行策略被提出，即在专家模块的数量维度上进行并行切分，使不同专家运行在不同设备上，以此实现将更大规模的混合专家模型部署在分布式环境下。专家并行策略可以被视为张量并行策略的变体。

垂直扩展可以理解为时间与空间置换策略的集合。不同的垂直扩展方式都是在利用一种资源（如计算、网络通信等）对存储资源进行置换。常见策略包括激活检查点(Activation Checkpointing)、零冗余优化(ZeRO)和卸载(Offloading)等。

激活检查点

在神经网络训练过程中，前向传播产生的激活值默认会被临时保存下来，用于反向传播阶段中的梯度求解，如图6（a）所示。随着模型参数的增加，这些激活值临时占用的存储空间相当可观，例如，Softmax激活输出的大小为序列的平方，因此对临时激活值进行存储空间上的优化显得格外重要。

图6 激活检查点核心逻辑图

激活检查点是一种非线性的内存节省方案，其在模型训练前向传播过程中，通过释放绝大部分激活值减少对内存的消耗。在反向传播过程中，通过额外的前向计算，重新计算所需的激活值，如图6（b）所示。这种策略通过引入额外的前向计算节省了内存占用，因此，可以将其视为计算资源（即时间）置换存储资源（即空间）。

ZeRO-DP

在模型训练过程中，设备内存不仅需要存储模型参数，还需要维护梯度和优化器状态。以常见的混合精度Adam优化器为例，设备内存中总计存放静态变量为半精度浮点数（fp16）的模型参数和梯度，以及优化器状态中为单精度浮点数（fp32）的模型参数、动量和方差。除此之外，设备内存中还要存放激活值等临时变量。默认配置下的数据并行、模型参数、梯度和优化器状态都会冗余存储N份，导致内存的利用率较低。

微软团队率先提出零冗余优化器ZeRO，其核心思想是通过减少数据并行下的冗余存储，显著降低内存消耗，从而扩大了可训练模型参数规模。ZeRO将模型状态均等地划分为N份，参与数据并行的每个成员仅存储全量参数的1/N，其中N等于数据并行的度。以此优化后，所有参与数据并行的成员共同完成了整份模型状态的存储。根据切分数据的不同，ZeRO策略可以划分为三个等级：ZeRO-1是对优化器状态的切分；ZeRO-2是对优化器状态和梯度的切分；ZeRO-3是对优化器状态、梯度和模型参数的切分。ZeRO策略可以被视为通过网络资源/通信来置换存储资源。

卸载

卸载是一种常用且直接的缓解设备内存压力的方法，其通过利用其他层级的存储资源，如CPU DRAM或磁盘等，将设备内存中变量卸载到这些容量相对充裕的存储单元上。然而，由于不同层级的存储单元的读写性能差异较大，对于性能要求较高的应用，简单的卸载机制势必会对程序性能产生较大的影响。因此，在实施卸载策略时，需要考虑其对程序性能的影响，并采取适当的优化措施以平衡内存使用和程序性能的需求。

我们将卸载机制划分成动态与静态两大类。动态卸载指工作流始终存在于一种介质内，模型参数在程序运行过程中频繁地进行跨介质的传输操作，以此减少对设备内存的消耗，如图7（a）所示。静态卸载指程序参数，如模型参数、优化器状态等，被提前静态地分配到不同的存储单元上，且在程序运行期间，依赖临时变量进行跨介质的信息传递，其工作流跨越不同存储介质，如图7（b）所示。

图7 卸载机制的两种类型

接下来，我将介绍一下我们团队在分布式训练框架优化方面的部分工作。这项工作基于动态卸载机制，针对大规模神经网络模型进行设计，旨在在有限的设备资源情况下，提高可训练模型参数规模，同时保证训练吞吐性能。

大模型参数规模较大，无法一次性完整加载至单个GPU内存中。因此，我们在GPU端设计一个工作窗口，用于加载模型中的特定层参数，将其余层参数暂时维护在CPU内存内。通过将参数在GPU内存和CPU内存之间动态转移，我们能够在较少的设备资源上完成较大规模模型的训练，实现更经济、更有效的模型构建方式。

工作窗口机制可以自适应前向传播和反向传播中的卸载过程。其中，数据传输过程采用异步方式处理，这样可以增加数据传输操作和GPU内部运算的重叠时间，从而在一定程度上隐藏/缓解由数据移动而额外引入的通信开销。这种机制可以有效地优化内存使用和程序性能，从而提高训练效率。

工作窗口机制的基本流程如下：在前向传播阶段，执行工作窗口内每一层的运算前，首先调用每层的pre hook函数，触发模型参数的异步加载请求，将工作窗口外的下一层参数从CPU内存加载至GPU内存，如图8（a）中的①所示；然后，工作窗口内的每一层执行前向传播计算，如图8（a）中的②所示；在当前层完成计算后，调用post hook函数，将前向计算中产生的临时变量卸载至CPU内存，如图8（a）中的③所示。在此过程中，工作窗口内每层占用的GPU缓存区都将被新的模型层参数重复利用。

图8 基于工作窗口的卸载机制

在反向传播阶段，首先通过pre hook函数进行模型层参数的加载，如图8（b）中的①所示，即将参数从CPU内存转移到GPU内存中；不同于前向传播，此处依然是通过pre hook函数，将已经完成反向传播运算的模型层（如梯度）从GPU内存中卸载至CPU内存中，如图8（b）中的②所示，同时判断是否需要进行参数更新；如果需要，则触发CPU端的优化器进行参数更新，如图8（b）中的③所示；最后，执行工作窗口内下一层的反向传播运算，如图8（b）中的④所示。

值得注意的是，工作窗口内的“层”不一定是模型中完整的一层，也可以是张量并行下部分参数的切片。此方法与数据并行、张量并行、流水线并行以及激活检查点等策略是相互正交的。

为了避免或抵消因数据转移操作而额外引入的性能开销，我们提出了一系列的优化。这里简单介绍一下性能优化方面的优化器参数更新与反向传播运算的交叠。

通常，模型训练包含三个步骤：前向传播、反向传播和参数更新。这些过程在默认执行的情况下，会按照严格的顺序串行执行，没有重叠。然而，在动态卸载的机制下，优化器状态被放置在CPU端，即需要利用CPU核的硬件资源执行参数更新过程，因此，并行执行效率势必会低于在GPU端的执行。为了弥补由此带来的性能损失，我们在对优化器进行实例化时，独立维护每一层的参数，并且在每一层的反向运算中，加入post hook函数。对于已经完成反向传播运算且需要进行参数更新的模型层，立即执行参数更新操作，而不必等待模型其他层的反向运算完成，此过程如图9所示。通过此优化，我们进一步保证了卸载机制的性能。

图9 优化器与反向传播交叠示意图

分布式模型训练系统在大模型发展中扮演着重要的角色，但其基础组件仍在不断发展和完善。尽管已有工作在高性能指标方面进行了支持，目前许多工作正在高性价比指标和高可用性指标方面进行探索和研究，但该领域依然存在诸多挑战和问题。本文仅简要介绍了目前主流解决方案的核心思想以及个人的一些观点，希望各位读者批评指正，也希望这些内容能够对读者提供帮助。

（本文根据CNCC2023特邀报告整理而成）

作者：

整理：

CCF推荐