分布式训练^1

随着大语言模型参数量和所需训练数据量的急速增长，单个机器上有限的资源已无法满足其训练的要求。需要设计分布式训练系统来解决海量的计算和内存资源需求问题。在分布式训练系统环境下，需要将一个模型训练任务拆分成多个子任务，并将子任务分发给多个计算设备，从而解决资源瓶颈。如何才能利用数万个计算加速芯片的集群，训练千亿甚至万亿参数规模的大语言模型？这其中涉及集群架构、并行策略、模型架构、内存优化、计算优化等一系列的技术。

本章将介绍分布式机器学习系统的基础概念、分布式训练的并行策略、分布式训练的集群架构，并以 DeepSpeed 为例，介绍如何在集群上训练大语言模型。

分布式训练概述

分布式训练（Distributed Training）是指将机器学习或深度学习模型训练任务分解成多个子任务，并在多个计算设备上并行训练。图4.1 给出了单个计算设备和多个计算设备的示例，这里计算设备可以是中央处理器（Central Processing Unit，CPU）、图形处理器（Graphics Processing Unit，GPU）、张量处理器（Tensor Processing Unit，TPU），也可以是神经网络处理器（Neural network Processing Unit，NPU）。由于同一个服务器内部的多个计算设备之间可能并不共享内存，因此无论这些计算设备是处于一个服务器还是多个服务器中，其系统架构都属于分布式系统范畴。一个模型训练任务往往会有大量的训练样本作为输入，可以利用一个计算设备完成，也可以将整个模型的训练任务拆分成多个子任务，分发给不同的计算设备，实现并行计算。此后，还需要对每个计算设备的输出进行合并，最终得到与单个计算设备等价的计算结果。由于每个计算设备只需要负责子任务，并且多个计算设备可以并行执行，因此其可以更快速地完成整体计算，并最终实现对整个计算过程的加速。

促使人们设计分布式训练系统的一个最重要的原因是单个计算设备的算力已经不足以支撑模型训练。图4.2 给出了机器学习模型对于算力的需求以及同期单个计算设备能够提供的算力。机器学习模型快速发展，从 2013 年 AlexNet 被提出开始，到 2022 年拥有 5400 亿个参数的 PaLM 模型被提出，再到 2024 年拥有 6710 亿个参数的 DeepSeek-V2 发布，机器学习模型以每 18 个月增长 56倍的速度发展。模型参数规模增大的同时，对训练数据量的要求也呈指数级增长，这更加剧了对算力的需求。然而，近几年，CPU 的算力增加已经远低于摩尔定律（Moore’s Law），虽然计算加速设备（如 GPU、TPU 等）为机器学习模型提供了大量的算力，但是其增长速度仍然没有突破每18 个月翻倍的摩尔定律。只有通过分布式训练系统才可以匹配模型不断增长的算力需求，满足机器学习模型的发展需要。

分布式训练的总体目标就是加快总的训练速度，减少模型训练的总体时间。总训练速度可以用式（4.1）简略估计：

其中，单设备计算速度主要由单块计算加速芯片的运算速度和数据 I/O 能力决定，对单设备训练效率进行优化，主要的技术手段有混合精度训练、算子融合、梯度累加等；在分布式训练系统中，随着计算设备数量的增加，理论上峰值计算速度会增加，然而受通信效率的影响，计算设备数量增多会造成加速比急速降低；多设备加速比是由计算和通信效率决定的，需要结合算法和网络拓扑结构进行优化，分布式训练并行策略的主要目标就是提升分布式训练系统中的多设备加速比。

大语言模型的参数量和所使用的数据量都非常大，因此都采用了分布式训练架构完成训练。文献 [13] 仅在 GPT-3 的训练过程中提到全部使用 NVIDIA V100 GPU，文献 [29] 介绍了 OPT 使用 992 块 NVIDIA A100 80GB GPU，采用全分片数据并行（Fully Sharded Data Parallel）[168] 以及Megatron-LM 张量并行（Tensor Parallelism）[169]，整体训练时间近两个月。BLOOM[31] 模型的研究人员则公开了更多在硬件和所采用的系统架构方面的细节。该模型的训练一共花费了 3.5 个月，使用 48 个计算节点。每个计算节点包含 8 块 NVIDIA A100 80GB GPU（总计 384 块 GPU），并且使用 4×NVLink 用于节点内部 GPU 之间的通信。节点之间采用 4 个 Omni-Path 100 Gbps 网卡构建的增强 8 维超立方体全局拓扑网络进行通信。文献 [34] 并没有给出 LLaMA 模型训练中所使用的集群的具体配置和网络拓扑结构，但是给出了不同参数规模的总 GPU 小时数。LLaMA 模型训练使用 NVIDIA A100 80GB GPU，LLaMA-7B 模型训练需要 82432 GPU 小时，LLaMA-13B 模型训练需要 135168 GPU 小时，LLaMA-33B 模型训练需要 530432 GPU 小时，而 LLaMA-65B 模型训练需要高达 1022362 GPU 小时。LLaMA 使用的训练数据量远超 OPT 和 BLOOM 模型，虽然模型参数量远小于上述两个模型，但是其所需计算量非常惊人。

通过使用分布式训练系统，大语言模型的训练周期可以从单计算设备花费几十年，缩短到使用数千个计算设备花费几十天。分布式训练系统需要克服计算墙、显存墙、通信墙等挑战，以确保集群内的所有资源得到充分利用，从而加速训练过程并缩短训练周期。

• 计算墙：单个计算设备所能提供的计算能力与大语言模型所需的总计算量之间存在巨大差异。2022 年 3 月发布的 NVIDIA H100 SXM 的单卡 FP16 算力只有 2000 TFLOPS（Floating Point Operations Per Second），而 GPT-3 需要 314 ZFLOPS 的总计算量，两者相差了 8 个数量级。
• 显存墙：单个计算设备无法完整存储一个大语言模型的参数。GPT-3 包含 1750 亿个参数，如果在推理阶段采用 FP32 格式进行存储，则需要 700GB 的计算设备内存空间，而 NVIDIA H100 GPU 只有 80GB 显存。
• 通信墙：分布式训练系统中各计算设备之间需要频繁地进行参数传输和同步。由于通信的延迟和带宽限制，这可能成为训练的瓶颈。在 GPT-3 的训练过程中，如果分布式系统中存在128 个模型副本，那么在每次迭代过程中至少需要传输 89.6TB 的梯度数据。截至 2023 年 8月，单个 InfiniBand 链路仅能提供不超过 800Gbps 的带宽。

计算墙和显存墙源于单计算设备的计算和存储能力有限，与模型所需庞大计算和存储需求存在矛盾。这个问题可以通过采用分布式训练的方法解决，但分布式训练又会面临通信墙的挑战。在多机多卡的训练中，这些问题逐渐显现。随着大语言模型参数的增大，对应的集群规模也随之增加，这些问题变更加突出。同时，当大型集群进行长时间训练时，设备故障可能会影响或中断训练，对分布式系统的问题处理也提出了很高的要求。

分布式训练的并行策略

分布式训练系统的目标是将单节点模型训练转换成等价的分布式并行模型训练。对于大语言模型来说，训练过程就是根据数据和损失函数，利用优化算法对神经网络模型参数进行更新的过程。单个计算设备模型训练系统的结构如图4.3 所示，其主要由数据和模型两个部分组成。训练过程由多个数据小批次（Mini-batch）完成。图中数据表示一个数据小批次。训练系统会利用数据小批次根据损失函数和优化算法计算梯度，从而对模型参数进行修正。针对大语言模型多层神经网络的执行过程，可以由一个计算图（Computational Graph）表示。这个图有多个相互连接的算子（Operator），每个算子实现一个神经网络层（Neural Network Layer），而参数则代表了这个层在训练中所更新的权重。

计算图的执行过程可以分为前向计算和反向计算两个阶段。前向计算的过程是将数据读入第一个算子，计算出相应的输出结构，然后重复这个前向计算过程，直到最后一个算子结束处理。反向计算的过程是根据损失函数和优化算法，对每个算子依次计算梯度，并利用梯度更新本地的参数。在反向计算结束后，该数据小批次的计算完成，系统就会读取下一个数据小批次，继续下一轮的模型参数更新。

根据单个计算设备模型训练系统的流程，可以看到，如果进行并行加速，可以从数据和模型两个维度进行考虑。可以对数据进行切分（Partition），并将同一个模型复制到多个设备上，并行执行不同的数据分片，这种方式通常被称为数据并行（Data Parallelism，DP）。还可以对模型进行划分，将模型中的算子分发到多个设备上分别完成处理，这种方式通常被称为模型并行（Model Parallelism，MP）。训练大语言模型时，往往需要同时对数据和模型进行切分，从而实现更高程度的并行，这种方式通常被称为混合并行（Hybrid Parallelism，HP）。

数据并行

在数据并行系统中，每个计算设备都有整个神经网络模型的模型副本（Model Replica），进行迭代时，每个计算设备只分配一个批次数据样本的子集，并根据该批次样本子集的数据进行网络模型的前向计算。假设一个批次的训练样本数为 N，使用 M 个计算设备并行计算，每个计算设备会分配到 N/M 个样本。前向计算完成后，每个计算设备都会根据本地样本计算损失误差，得到梯度 $G_i$（i 为加速卡编号），并将本地梯度 $G_i$ 进行广播。所有计算设备需要聚合其他加速卡给出的梯度值，然后使用平均梯度 $(\sum^{N}_{i=1} G_i)/N$ 对模型进行更新，完成该批次训练。图4.4 给出了由两个计算设备组成的数据并行训练系统样例。

数据并行训练系统可以通过增加计算设备，有效提升整体训练吞吐量，即每秒全局批次数（Global Batch Size Per Second）。与单个计算设备训练相比，其最主要的区别在于反向计算中的梯度需要在所有计算设备中进行同步，以保证每个计算设备上最终得到的是所有进程上梯度的平均值。常见的神经网络框架中都有数据并行方式的具体实现，包括 TensorFlow DistributedStrategy、PyTorch Distributed、Horovod DistributedOptimizer 等。由于基于 Transformer 结构的大语言模型中每个算子都依赖单个数据而非批次数据，因此数据并行并不会影响其计算逻辑。一般情况下，各训练设备中前向计算是独立的，不涉及同步问题。数据并行训练加速比最高，但要求每个设备上都备份一份模型，显存占用比较高。

使用 PyTorch DistributedDataParallel 实现单个服务器多加速卡训练的代码如下。首先，构造DistributedSampler 类，将数据集的样本随机打乱并分配到不同计算设备上：

利用 DistributedSampler 类构造的完整的训练程序样例 main.py 如下：

通过以下命令行启动上述程序：

模型并行

模型并行往往用于解决单节点内存不足的问题。以包含 1750 亿个参数的 GPT-3 模型为例，如果模型中每一个参数都使用 32 位浮点数表示，那么模型需要占用 700GB 内存。如果使用 16 位浮点数表示，那么每个模型副本需要占用 350GB 内存。2022 年 3 月 NVIDIA 发布的 H100 加速卡仅支持 80GB 显存，无法将整个模型完整放入其中。模型并行可以从计算图角度，用以下两种形式进行切分。

（1）按模型的层切分到不同设备，即层间并行或算子间并行（Inter-operator Parallelism），也称之为流水线并行（Pipeline Parallelism，PP）。
（2）将计算图层内的参数切分到不同设备，即层内并行或算子内并行（Intra-operator Parallelism），也称之为张量并行（Tensor Parallelism，TP）。两节点模型并行训练系统样例如图4.5 所示，图4.5(a)为流水线并行，模型的不同层被切分到不同的设备中；图4.5(b) 为张量并行，同一层中的不同参数被切分到不同的设备中进行计算。

流水线并行

流水线并行是一种并行计算策略，将模型的各个层分段处理，并将每个段分布在不同的计算设备上，使得前后阶段能够流水式、分批工作。流水线并行通常应用于大语言模型的并行系统中，以有效解决单个计算设备内存不足的问题。图4.6 给出了一个由四个计算设备组成的流水线并行系统，包含前向计算和后向计算。其中 F1、F2、F3、F4 分别代表四个前向路径，位于不同的设备上；而 B4、B3、B2、B1 则代表逆序的后向路径，也分别位于四个不同的设备上。从图 4.6 中可以看出，计算图中的下游设备（Downstream Device）需要长时间持续处于空闲状态，等待上游设备（Upstream Device）计算完成，才能开始计算自身的任务。这种情况导致设备的平均使用率大幅降低，形成了模型并行气泡（Model Parallelism Bubble），也称为流水线气泡（Pipeline Bubble）。

朴素流水线策略所产生的并行气泡，使得系统无法充分利用计算资源，降低了系统整体的计算效率。为了减少并行气泡，文献 [170] 提出了 GPipe 方法，将小批次（Mini-batch）进一步划分成更小的微批次（Micro-batch），利用流水线并行方法，每次处理一个微批次的数据。在当前阶段计算完成得到结果后，将该微批次的结果发送给下游设备，同时开始处理后一个微批次的数据，这样可以在一定程度上减少并行气泡。图4.7 给出了 GPipe 策略流水线并行样例。前向 F1 计算被拆解为 F11、F12、F13、F14，在计算设备 1 中计算完成 F11 后，会在计算设备 2 中进行 F21 计算，同时在计算设备 1 中并行计算 F12。相比于最原始的流水线并行方法，GPipe 流水线方法可以有效减少并行气泡。

虽然 GPipe 策略可以减少一定的并行气泡，但是只有当一个小批次中所有的前向计算都完成时，才能执行后向计算。因此，还是会产生很多并行气泡，从而降低系统的并行效率。MegatronLM[171] 采了 1F1B 流水线并行策略，即一个前向通道和一个后向通道。1F1B 流水线并行策略引入了任务调度机制，使得下游设备能够在等待上游计算的同时执行其他可并行的任务，从而提高设备的利用率。1F1B 给出了非交错式和交错式两种调度模式，如图4.8 所示。

1F1B 非交错式调度模式可分为三个阶段。首先是热身阶段，在计算设备中进行不同数量的前向计算。接下来的阶段是前向-后向阶段，计算设备按顺序执行一次前向计算，然后进行一次后向计算。最后一个阶段是后向阶段，计算设备完成最后一次后向计算。相比于 GPipe 策略，1F1B 非交错式调度模式在节省内存方面表现得更好。然而，它需要与 GPipe 策略一样的时间来完成一轮计算。

1F1B 交错式调度模式要求微批次的数量是流水线阶段的整数倍。每个设备不仅负责连续多个层的计算，还可以处理多个层的子集，这些子集被称为模型块。具体而言，在之前的模式中，设备 1 可能负责层 1∼4，设备 2 负责层 5∼8，依此类推。在新的模式下，设备 1 可以处理层 1、2、9、10，设备 2 处理层 3、4、11、12，依此类推。在这种模式下，每个设备在流水线中被分配到多个阶段。例如，设备 1 可能参与热身阶段、前向计算阶段和后向计算阶段的某些子集任务。每个设备可以并行执行不同阶段的计算任务，从而更好地利用流水线并行的优势。这种模式不仅在内存消耗方面表现出色，还能提高计算效率，使大型模型的并行系统能够更高效地完成计算任务。

PyTorch 中也包含了实现流水线的 API 函数 Pipe，具体实现参考“torch.distributed.pipeline.sync.Pipe”类。可以使用这个 API 构造一个模型，其包含两个线性层，分别放置在两个计算设备中的样例如下：

张量并行

张量并行需要根据模型的具体结构和算子类型，解决如何将参数切分到不同设备，以及如何保证切分后的数学一致性这两个问题。大语言模型都是以 Transformer 结构为基础，Transformer 结构主要由嵌入式表示（Embedding）、矩阵乘（MatMul）和交叉熵损失（Cross Entropy Loss）计算构成。这三种类型的算子有较大的差异，需要设计对应的张量并行策略[169] 才可以实现将参数切分到不同的设备。

对于嵌入式表示算子，如果总的词表数非常大，会导致单计算设备显存无法容纳 Embedding层参数。举例来说，如果词表数量是 64000，嵌入式表示维度为 5120，类型采用 32 位精度浮点数，那么整层参数需要的显存大约为 64000 × 5120 × 4/1024/1024 = 1250MB，反向梯度同样需要1250MB 显存，仅仅存储就需要将近 2.5GB。对于嵌入表示层的参数，可以按照词维度切分，每个计算设备只存储部分词向量，然后通过汇总各个设备上的部分词向量，得到完整的词向量。图4.9给出了单节点 Embedding 和两节点 Embedding 张量并行的示意图。在单节点上，执行 Embedding操作，bz 是批次大小（batch size），Embedding 的参数大小为 [word_size, hidden_size]，计算得到[bz, hidden_size] 张量。图4.9 中 Embedding 张量并行示例将 Embedding 参数沿 word_size 维度切分为两块，每块大小为 [word_size/2, hidden_size]，分别存储在两个设备上。当每个节点查询各自的词表时，如果无法查到，则该词的表示为 0，各设备查询后得到 [bz, hidden_size] 结果张量，最后通过 AllReduce_Sum 通信①，跨设备求和，得到完整的全量结果。可以看出，这里的输出结果和单计算设备执行的结果一致。

矩阵乘的张量并行要充分利用矩阵的分块乘法原理。举例来说，要实现如下矩阵乘法 Y =XA，其中 X 是维度为 M × N 的输入矩阵，A 是维度为 N × K 的参数矩阵，Y 是结果矩阵，维度为 M × K。如果参数矩阵 A 非常大，甚至超出单张卡的显存容量，那么可以把参数矩阵 A 切分到多张卡上，并通过集合通信汇集结果，保证最终结果在数学计算上等价于单计算设备的计算结果。参数矩阵 A 存在以下两种切分方式。

（1）参数矩阵 A 按列切块，将矩阵 A 按列切成

（2）参数矩阵 A 按行切块，将矩阵 A 按行切成

图4.10 给出了参数矩阵按列切分的示例，参数矩阵 A 分别将 A1, A2 放置在两个计算设备上。两个计算设备分别计算 Y1 = XA1 和 Y2 = XA2。计算完成后，多计算设备间进行通信，从而获取其他计算设备上的计算结果，并拼接在一起得到最终的结果矩阵 Y ，该结果在数学上与单计算设备在计算结果上完全等价。

图4.11 给出了参数矩阵按行切分的示例，为了满足矩阵乘法规则，输入矩阵 X 需要按列切分X = [X1|X2]。同时，将矩阵分块，分别放置在两个计算设备上，每个计算设备分别计算 Y1 = X1A1和 Y2 = X2A2。计算完成后，多个计算设备间通信获取其他卡上的计算结果，可以得到最终的结果矩阵 Y 。同样，这种切分方式，既可以保证数学上的计算等价性，解决单计算设备显存无法容纳的问题，又可以保证单计算设备通过拆分的方式装下参数 A。

Transformer 中的 FFN 结构均包含两层全连接（Fully Connected，FC）层，即存在两个矩阵乘，这两个矩阵乘分别采用上述两种切分方式，如图4.12 所示。对第一个 FC 层的参数矩阵按列切块，对第二个 FC 层的参数矩阵按行切块。这样，第一个 FC 层的输出恰好满足第二个 FC 层的数据输入要求（按列切分），因此可以省去第一个 FC 层后的汇总通信操作。多头自注意力机制的张量并行与 FFN 类似，因为具有多个独立的头，所以相较于 FFN 更容易实现并行，其矩阵切分方式如图4.13 所示。具体可以参考文献 [169]。

分类网络最后一层一般会选用 Softmax 和 Cross_entropy 算子来计算交叉熵损失。如果类别数量非常大，则会导致单计算设备内存无法存储和计算 logit 矩阵。针对这一类算子，可以按照类别维度切分，同时通过中间结果通信，得到最终的全局交叉熵损失。首先计算的是 Softmax 值，公式如下：

其中，p 表示张量并行的设备号。得到 Softmax 计算结果之后，同时对标签 Target 按类别切分，每个设备得到部分损失，最后进行一次通信，得到所有类别的损失。整个过程，只需要进行三次小量的通信，就可以完成交叉熵损失的计算。

PyTorch 提供了细粒度张量级别的并行 API——DistributedTensor。也提供了粗粒度模型层面的API 对“nn.Module”进行张量并行。通过以下几行代码就可以实现对一个大的张量进行分片：

对于像“nn.Linear”这样已经有“torch.Tensor”作为参数的模块，也提供了模块级 API“distribute_module”在模型层面进行张量并行，参考代码如下：

混合并行

混合并行将多种并行策略如数据并行、流水线并行和张量并行等混合使用。通过结合不同的并行策略，混合并行可以充分发挥各种并行策略的优点，最大限度地提高计算性能和效率。针对千亿规模的大语言模型，通常，在每个服务器内部使用张量并行策略，由于该策略涉及的网络通信量较大，因此需要利用服务器内部的不同计算设备之间的高速通信带宽。通过流水线并行，将模型的不同层划分为多个阶段，每个阶段由不同的机器负责计算。这样可以充分利用多台机器的计算能力，并通过机器之间的高速通信传递计算结果和中间数据，以提高整体的计算速度和效率。最后，在外层叠加数据并行策略，以增加并发数量，加快整体训练速度。通过数据并行，将训练数据分发到多组服务器上进行并行处理，每组服务器处理不同的数据批次。这样可以充分利用多台服务器的计算资源，并增加训练的并发度，从而加快整体训练速度。

BLOOM 使用 Megatron-DeepSpeed[134] 框架进行训练，主要包含两个部分：Megatron-LM 提供张量并行能力和数据加载原语；DeepSpeed[172] 提供 ZeRO 优化器、模型流水线及常规的分布式训练组件。通过这种方式可以实现数据、张量和流水线三维并行，BLOOM 模型训练时采用的并行计算结构如图4.14 所示。BLOOM 模型训练使用由 48 个 NVIDIA DGX-A100 服务器组成的集群，每个 DGX-A100 服务器包含 8 块 NVIDIA A100 80GB GPU，总计包含 384 块。BLOOM训练采用的策略是先将集群分为 48 个一组，进行数据并行。接下来，模型整体被分为 12 个阶段，进行流水线并行。每个阶段的模型被划分到 4 块 GPU 中，进行张量并行。同时，BLOOM使用了 ZeRO（零冗余优化器）[173] 进一步降低模型对显存的占用。通过上述步骤可以实现数百个GPU 的高效并行计算。

计算设备内存优化

当前，大语言模型训练通常采用 Adam 优化算法，除了需要每个参数梯度，还需要一阶动量（Momentum）和二阶动量（Variance）。虽然 Adam 优化算法相较 SGD 算法效果更好也更稳定，但是对计算设备内存的占用显著增大。为了降低内存占用，大多数系统采用混合精度训练（Mixed Precision Training）方式，即同时存在 FP32（32 位浮点数）与FP16（16 位浮点数）或者 BF16（BFloat16）格式的数值。FP32、FP16 和 BF16 的表示如图4.15 所示。FP32 中第 31 位为符号位，第 30 位∼第23 位用于表示指数，第 22 位∼第 0 位用于表示尾数。FP16 中第 15 位为符号位，第 14 位∼第 10位用于表示指数，第 9 位∼第 0 位用于表示尾数。BF16 中第 15 位为符号位，第 14 位∼第 7 位用于表示指数，第 6 位∼第 0 位用于表示尾数。由于 FP16 的值区间比 FP32 的值区间小很多，所以在计算过程中很容易出现上溢出和下溢出。BF16 相较于 FP16 以精度换取更大的值区间范围。由于 FP16 和 BF16 相较 FP32 精度低，训练过程中可能会出现梯度消失和模型不稳定的问题，因此，需要使用一些技术解决这些问题，例如动态损失缩放（Dynamic Loss Scaling）和混合精度优化器（Mixed Precision Optimizer）等。

混合精度优化的过程如图4.16 所示。Adam 优化器状态包括采用 FP32 保存的模型参数备份，一阶动量和二阶动量也都采用 FP32 格式存储。假设模型参数量为 Φ，模型参数和梯度都是用 FP16格式存储，则共需要 2Φ + 2Φ + (4Φ + 4Φ + 4Φ) = 16Φ 字节存储。其中，Adam 状态占比 75%。动态损失缩放在反向传播前，将损失变化（dLoss）手动增大 2K 倍，因此反向传播时得到的激活函数梯度不会溢出；反向传播后，将权重梯度缩小 2K 倍，恢复正常值。举例来说，有 75 亿个参数的模型，如果用 FP16 格式，只需要 15GB 计算设备内存，但是在训练阶段，模型状态实际上需要耗费 120GB 内存。计算卡内存占用中除了模型状态，还有剩余状态（Residual States），包括激活值（Activation）、各种临时缓冲区（Buffer）及无法使用的显存碎片（Fragmentation）等。可以使用激活值检查点（Activation Checkpointing）方式使激活值内存占用大幅减少，因此如何减少模型状态尤其是 Adam 优化器状态是解决内存占用问题的关键。

零冗余优化器（Zero Redundancy Data Parallelism，ZeRO）的目标是针对模型状态的存储进行去除冗余的优化[173–175]。ZeRO 使用分区的方法，即将模型状态量分割成多个分区，每个计算设备只保存其中的一部分。这样整个训练系统内只需要维护一份模型状态，减少了内存消耗和通信开销。具体来说，如图4.17 所示，ZeRO 包含以下三种方法。

（1）对 Adam 优化器状态进行分区，图4.17 中的 Pos 部分。模型参数和梯度依然是每个计算设备保存一份。此时，每个计算设备所需内存是 $4Φ +\frac{12Φ}{N}$ 字节，其中 N 是计算设备总数。当 N 比较大时，每个计算设备占用内存趋向于 4ΦB，也就是 16ΦB 的 $\frac{1}{4}$。
（2）对模型梯度进行分区，图4.17 中的 $P_{os+g}$ 部分。模型参数依然是每个计算设备保存一份。此时，每个计算设备所需内存是 $2Φ +\frac{2Φ+12Φ}{N}$ 字节。当 N 比较大时，每个计算设备占用内存趋向于 2ΦB，也就是 16ΦB 的 1/8。
（3）对模型参数进行分区，图4.17 中的 $P_{os+g+p}$ 部分。此时，每个计算设备所需内存是 $\frac{16Φ}{N} B$。当 N 比较大时，每个计算设备占用内存趋向于 0。

在 DeepSpeed 框架中，Pos 对应 Zero-1，$P_{os+g}$ 对应 Zero-2，$P_{os+g+p}$ 对应 Zero-3。文献 [175]中也对 ZeRO 优化方法所带来的通信量增加的情况进行了分析，Zero-1 和 Zero-2 对整体通信量没有影响，虽然对通信有一定延迟影响，但是整体性能受到的影响很小。Zero-3 所需的通信量则是正常通信量的 1.5 倍。

PyTorch中也实现了ZeRO优化方法，可以使用ZeroRedundancyOptimizer调用，也可与“torch.nn.parallel.DistributedDataParallel”结合使用，以减少每个计算设备的内存峰值消耗。使用 ZeroRedundancyOptimizer 的参考代码如下所示：

执行上述代码，可以得到如下输出：

可以看到，每次迭代之后，无论是否使用 ZeroRedundancyOptimizer，模型参数都使用同样的内存。在启用 ZeroRedundancyOptimizer 封装 Adam 优化器后，优化器的 step() 操作的内存峰值消耗是Adam 内存消耗的一半。

分布式训练的集群架构

分布式训练需要使用由多台服务器组成的计算集群（Computing Cluster），而集群的架构也需要根据分布式系统、大语言模型结构、优化算法等综合因素进行设计。分布式训练集群属于高性能计算集群（High Performance Computing Cluster，HPC），其目标是提供海量的计算能力。在由高速网络组成的高性能计算上构建分布式训练系统，主要有两种常见架构：参数服务器架构和去中心化架构。

本章介绍高性能计算集群的典型硬件组成，并在此基础上介绍分布式训练系统所采用的参数服务器架构和去中心化架构。

高性能计算集群的典型硬件组成

典型的用于分布式训练的高性能计算集群的硬件组成如图4.18 所示。整个计算集群包含大量带有计算加速设备的服务器。每个服务器中往往有多个计算加速设备（通常为 2∼16 个）。多个服务器会被放置在一个机柜（Rack）中，服务器通过架顶交换机（Top of Rack Switch，ToR）连接网络。在架顶交换机满载的情况下，可以通过在架顶交换机间增加骨干交换机（Spine Switch）接入新的机柜。这种连接服务器的拓扑结构往往是一个多层树（Multi-Level Tree）。

在多层树结构集群中跨机柜通信（Cross-Rack Communication）往往会有网络瓶颈。以包含 1750亿个参数的 GPT-3 模型为例，每一个参数使用 32 位浮点数表示，在每一轮训练迭代中，每个模型副本会生成 700GB 的本地梯度数据。假如采用包含 1024 卡的计算集群，包含 128 个模型副本，那么至少需要传输 89.6TB（700GB×128 = 89.6TB）的梯度数据。这会造成严重的网络通信瓶颈。因此，针对大语言模型分布式训练，通常采用胖树[176]（Fat-Tree）拓扑结构，试图实现网络带宽的无收敛。此外，采用 InfiniBand（IB）技术搭建高速网络，单个 InfiniBand 链路可以提供 200Gbps或者 400Gbps 带宽。NVIDIA 的 DGX 服务器提供单机 1.6Tbps（200Gbps×8）网络带宽，HGX 服务器网络带宽更是可以达到 3.2Tbps（400Gbps×8）。

单个服务器通常由 2∼16 个计算加速设备组成，这些计算加速设备之间的通信带宽也是影响分布式训练的重要因素。如果这些计算加速设备通过服务器 PCIe 总线互联，则会造成服务器内部计算加速设备之间的通信瓶颈。PCIe 5.0 总线也只能提供 128GB/s 的带宽，而 NVIDIA H100 采用的 HBM 可以提供 3350GB/s 的带宽。因此，服务器内部通常采用异构网络架构。NVIDIA HGXH100 8-GPU 服务器采用 NVLink 和 NVSwitch（NVLink 交换机）技术，如图4.19 所示。每块 H100GPU 都有多个 NVLink 端口，并连接到所有（4 个）NVSwitch 上。每个 NVSwitch 都是一个完全无阻塞的交换机，完全连接所有（8 块）H100 计算加速卡。NVSwitch 的这种完全连接的拓扑结构，使得服务器内任何 H100 加速卡之间都可以达到 900GB/s 的双向通信速度。

参数服务器架构

参数服务器（Parameter Server，PS）架构的分布式训练系统中有两种服务器角色：训练服务器和参数服务器。参数服务器需要提供充足的内存资源和通信资源，训练服务器需要提供大量的计算资源。图4.20 为参数服务器的分布式训练集群的示意图。该集群包括两个训练服务器和两个参数服务器。假设有一个可分为两个参数分区的模型，每个分区由一个参数服务器负责参数同步。在训练过程中，每个训练服务器都拥有完整的模型，将分配到此服务器的训练数据集切片（Dataset Shard）并进行计算，将得到的梯度推送到相应的参数服务器。参数服务器会等待两个训练服务器都完成梯度推送，再计算平均梯度并更新参数。之后，参数服务器会通知训练服务器拉取最新的参数，并开始下一轮训练迭代。

参数服务器架构的分布式训练过程可以细分为同步训练和异步训练两种模式。

• 同步训练：训练服务器在完成一个小批次的训练后，将梯度推送给参数服务器。参数服务器在收到所有训练服务器的梯度后，进行梯度聚合和参数更新。
• 异步训练：训练服务器在完成一个小批次的训练后，将梯度推送给参数服务器。参数服务器不再等待接收所有训练服务器的梯度，而是直接基于已收到的梯度进行参数更新。

在同步训练的过程中，参数服务器会等待所有训练服务器完成当前小批次的训练，有诸多的等待或同步机制，导致整个训练速度较慢。异步训练去除了训练过程中的等待机制，训练服务器可以独立进行参数更新，极大地加快了训练速度。引入异步更新的机制会导致训练效果有所波动。应根据具体情况和需求选择适合的训练模式。

去中心化架构

去中心化（Decentralized Network）架构采用集合通信实现分布式训练系统。在去中心化架构中，没有中央服务器或控制节点，而是由节点之间进行直接通信和协调。这种架构的好处是可以减少通信瓶颈，提高系统的可扩展性。由于节点之间可以并行地训练和通信，去中心化架构可以显著降低通信开销，并减少通信墙的影响。在分布式训练过程中，节点之间需要周期性地交换参数更新和梯度信息。可以通过集合通信（Collective Communication，CC）技术实现分布式训练，常用通信原语包括 Broadcast、Scatter、Reduce、All Reduce、Gather、All Gather、Reduce Scatter、All to All 等。4.2 节介绍的大语言模型训练所使用的分布式训练并行策略，大多使用去中心化架构，并利用集合通信实现。

下面介绍一些常见的集合通信原语。

（1）Broadcast：主节点把自身的数据发送到集群中的其他节点。Broadcast 在分布式训练系统中常用于网络参数的初始化。如图4.21 所示，计算设备 1 对大小为 1 × N 的张量进行广播，最终每张卡输出均为 [1 × N] 的矩阵。

（2）Scatter：主节点对数据进行划分并散布至其他指定的节点。Scatter 与 Broadcast 非常相似，不同的是，Scatter 是将数据的不同部分按需发送给所有的进程。如图4.22 所示，计算设备 1 将大小为 1 × N 的张量分为 4 份后发送到不同节点。

（3）Reduce：是一系列简单运算操作的统称，将不同节点上的计算结果进行聚合（Aggregation），可以细分为 Sum、Min、Max、Prod、Lor 等类型的归约操作。如图4.23 所示，Reduce Sum 操作将所有计算设备上的数据汇聚到计算设备 1，并执行求和操作。

（4）All Reduce：在所有的节点上都应用同样的 Reduce 操作。可以细分为 Sum、Min、Max、Prod、Lor 等类型的归约操作。All Reduce 操作可通过单节点上的“Reduce + Broadcast”操作完成。如图4.24 所示，All Reduce Sum 操作将所有计算设备上的数据汇聚到各个计算设备中，并执行求和操作。

（5）Gather：将多个节点上的数据收集到单个节点上，可以将 Gather 理解为反向的 Scatter。如图4.25 所示，Gather 操作将所有计算设备上的数据收集到计算设备 1 中。

（6）All Gather：每个节点都收集所有其他节点上的数据，All Gather 相当于一个 Gather 操作之后跟着一个 Broadcast 操作。如图4.26 所示，All Gather 操作将所有计算设备上的数据收集到每个计算设备中。

（7）Reduce Scatter：将每个节点中的张量切分为多个块，每个块被分配给不同的节点。接收到的块会在每个节点上进行特定的操作，例如求和、取平均值等。如图4.27 所示，每个计算设备都将其中的张量切分为 4 块，并分发到 4 个不同的计算设备中，每个计算设备分别对接收的分块进行特定操作。

（8）All to All：将每个节点的张量切分为多个块，每个块分别发送给不同的节点。如图4.28 所示，每个计算设备都将其中的张量切分为 4 块，并分发到 4 个不同的计算设备中。

分布式集群中的网络硬件多种多样，包括以太网、InfiniBand 网络等。PyTorch 等深度学习框架通常不直接操作硬件，而是使用通信库。常用的通信库包括 MPI、GLOO、NCCL 等，可以根据具体情况进行选择和配置。MPI（Message Passing Interface）是一种广泛使用的并行计算通信库，常用于在多个进程之间进行通信和协调。GLOO 是 Facebook 推出的一个类似 MPI 的集合通信库（Collective Communications Library），也大体遵照 MPI 提供的接口规定，实现了包括点对点通信、集合通信等相关接口，支持在 CPU 和 GPU 上的分布式训练。NCCL（NVIDIA Collective Communications Library）是 NVIDIA 开发的高性能 GPU 间通信库，专门用于在多个 GPU 之间进行快速通信和同步，因为 NCCL 是 NVIDIA 基于自身硬件定制的，能做到更有针对性且更便于优化，故在 NVIDIA硬件上，NCCL 的效果往往比其他通信库更好。GLOO、MPI 和 NCCL 在 CPU 和 GPU 环境下对通信原语的支持情况如表4.1 所示。在进行分布式训练时，根据所使用的硬件环境和需求，选择适当的通信库可以充分发挥硬件的优势并提高分布式训练的性能和效率。一般而言，如果在 CPU 集群上进行训练，则可选择使用 MPI 或 GLOO 作为通信库；而如果在 GPU 集群上进行训练，则可以选择 NCCL 作为通信库。

以 PyTorch 为例，介绍如何使用上述通信原语完成多计算设备间通信。先使用“torch.distributed”初始化分布式环境：

接下来使用“torch.multiprocessing”开启多个进程，本例中共开启了 4 个进程：

每个新开启的进程都会调用“init_process”，接下来调用用户指定的函数“func”。这里以 All Reduce 为例：

根据 All Reduce 通信原语，在所有的节点上都应用同样的 Reduce 操作，可以得到如下输出：

DeepSpeed 实践

DeepSpeed[172] 是一个由 Microsoft 公司开发的开源深度学习优化库，旨在提高大语言模型训练的效率和可扩展性，使研究人员和工程师能够更快地迭代和探索新的深度学习模型和算法。它采用了多种技术手段来加速训练，包括模型并行化、梯度累积、动态精度缩放、本地模式混合精度等。此外，DeepSpeed 还提供了一些辅助工具，例如分布式训练管理、内存优化和模型压缩，以帮助开发者更好地管理和优化大规模深度学习训练任务。DeepSpeed 是基于 PyTorch 构建的，因此将现有的 PyTorch 训练代码迁移到 DeepSpeed 上通常只需要进行简单的修改。这使得开发者可以快速利用 DeepSpeed 的优化功能来加速训练任务。DeepSpeed 已经在许多大规模深度学习项目中得到了应用，包括语言模型、图像分类、目标检测等领域。大语言模型 BLOOM[31]（1750 亿个参数）和 MT-NLG[134]（5400 亿个参数）都采用 DeepSpeed 框架完成训练。

DeepSpeed 的主要优势在于支持大规模神经网络模型、提供了更多的优化策略和工具。DeepSpeed 通过实现三种并行方法的灵活组合，即 ZeRO 支持的数据并行、流水线并行和张量并行，可以应对不同工作负载的需求。特别是通过 3D 并行性的支持，DeepSpeed 可以处理具有万亿个参数的超大规模模型。DeepSpeed 还引入了 ZeRO-Offload，使单个 GPU 能够训练比其显存容量大 10倍的模型。为了充分利用 CPU 和 GPU 的内存来训练大语言模型，DeepSpeed 还扩展了 ZeRO-2。此外，DeepSpeed 还提供了稀疏注意力核（Sparse Attention Kernel），支持处理包括文本、图像和语音等长序列输入的模型。DeepSpeed 还集成了 1 比特 Adam 算法（1-bit Adam），该算法可以只使用原始 Adam 算法 1/5 的通信量，达到与 Adam 类似的收敛率，显著提高分布式训练的效率，降低通信开销。

DeepSpeed 的 3D 并行充分利用硬件架构特性，综合考虑了显存效率和计算效率。4.3 节介绍了分布式集群的硬件架构，截至 2023 年 9 月，分布式训练集群通常采用 NVIDIA DGX/HGX 节点，利用胖树网络拓扑结构构建计算集群。因此，每个节点内部 8 个计算加速设备之间具有非常高的通信带宽，节点之间的通信带宽则相对较低。由于张量并行是分布式训练策略中通信开销最大的，因此优先考虑将张量并行计算组放置在节点内以利用更大的节点内带宽。当张量并行组不能占满节点内的所有计算节点时，选择将数据并行组放置在节点内，否则就使用跨节点进行数据并行。流水线并行的通信量最低，因此可以使用跨节点的方式调度流水线的各个阶段，降低通信带宽的要求。每个数据并行组需要通信的梯度量随着流水线和模型并行的规模线性减小，因此总通信量少于单纯使用数据并行。此外，每个数据并行组会在局部的一小部分计算节点内部独立通信，组间通信可以并行。过减少通信量和增加局部性与并行性，数据并行通信的有效带宽有效增大。

图4.29 给出了 DeepSpeed 3D 并行策略示意图。图中给出了 32 个计算设备进行 3D 并行的例子。神经网络的各层分为 4 个流水线阶段。每个流水线阶段中的层在 4 个张量并行计算设备之间进一步划分。最后，每个流水线阶段有两个数据并行实例，使用 ZeRO 内存优化在这 2 个副本之间划分优化器状态量。

DeepSpeed 软件架构如图4.30 所示，主要包含以下三部分。

（1）API：DeepSpeed 提供了易于使用的 API 接口，简化了训练模型和推断的过程。用户只需调用几个 API 接口即可完成任务。通过“initialize”接口可以初始化引擎，并在参数中配置训练参数、优化技术等。这些配置参数通常保存在名为“ds_config.json”的文件中。
（2）RunTime：RunTime 是 DeepSpeed 的核心运行时组件，使用 Python 语言实现，负责管理、执行和优化性能。它承担了将训练任务部署到分布式设备的功能，包括数据分区、模型分区、系统优化、微调、故障检测及检查点的保存和加载等任务。
（3）Ops：Ops 是 DeepSpeed 的底层内核组件，使用 C++ 和 CUDA 实现。它优化计算和通信过程，提供了一系列底层操作，包括 Ultrafast Transformer Kernels、Fuse LAN Kernels、Customary Deals 等。Ops 的目标是通过高效的计算和通信加速深度学习训练过程。

基础概念

DeepSpeed 提供了分布式计算框架，首先需要明确几个重要的基础概念：主节点、节点编号、全局进程编号、局部进程编号和全局总进程数。DeepSpeed 主节点（master_ip+master_port）负责协调所有其他节点和进程的工作，由主节点所在服务器的 IP 地址和主节点进程的端口号来确定主节点。主节点还负责监控系统状态、处理任务分配、结果汇总等任务，因此是整个系统的关键部分。节点编号（node_rank）是系统中每个节点的唯一标识符，用于区分不同计算机之间的通信。全局进程编号（rank）是整个系统中的每个进程的唯一标识符，用于区分不同进程之间的通信。局部进程编号（local_rank）是单个节点内的每个进程的唯一标识符，用于区分同一节点内的不同进程之间的通信。全局总进程数（world_size）是整个系统中运行的所有进程的总数，用于确定可以并行完成多少工作及完成任务所需的资源数量。

在网络通信策略方面，DeepSpeed 提供了 MPI、GLOO、NCCL 等选项，可以根据具体情况进行选择和配置。在 DeepSpeed 配置文件中，在 optimizer 部分配置通信策略，以下是使用 1-Bit Adam优化器的配置样例，配置中使用了 NCCL 通信库：

DeepSpeed 中也支持多种类型 ZeRO 的分片机制，包括 ZeRO-0、ZeRO-1、ZeRO-2、ZeRO-3 以及
ZeRO-Infinity。ZeRO-0 禁用所有类型的分片，仅将 DeepSpeed 当作分布式数据并行使用；ZeRO-1对优化器状态进行分片，占用内存为原始的 1/4，通信容量与数据并行性相同；ZeRO-2 对优化器状态和梯度进行分片，占用内存为原始的 1/8，通信容量与数据并行性相同；ZeRO-3 对优化器状态、梯度及模型参数进行分片，内存减少与数据并行度和复杂度成线性关系，同时通信容量是数据并行性的 1.5 倍；ZeRO-Infinity 是 ZeRO-3 的拓展，允许通过使用 NVMe 固态硬盘扩展 GPU 和CPU 内存来训练大语言模型。

以下是 DeepSpeed 使用 ZeRO-3 配置参数的样例：

如果希望在 ZeRO-3 的基础上继续使用 ZeRO-Infinity 将优化器状态和计算转移到 CPU 中，则可以在配置文件中按照如下方式配置：

甚至可以进一步将模型参数也装载到 CPU 内存中，在配置文件中按照如下方式配置：

如果希望将更多的内存装载到 NVMe 中，则可以在配置文件中按照如下方式配置：

LLaMA 分布式训练实践

LLaMA 模型是目前最流行、性能最强大的开源模型之一，基于 LLaMA 构造的模型生态可以覆盖绝大部分模型使用场景。在设置完必要的数据和环境配置后，本节将逐步演示如何使用 DeepSpeed 框训练 LLaMA 模型。

DeepSpeed 可以很好地兼容 PyTorch 和 CUDA 的大多数版本，其安装过程通常无须指定特殊配置选项，直接通过 pip 命令完成。

训练数据配置

使用 PyTorch 和 transformers 库来设置预训练模型的数据加载器，以实现在单机或多机分布式训练环境中对数据的加载和采样。需要导入的模块如下。

• DataLoader 是 PyTorch 提供的工具，用于从数据集加载数据到模型进行训练或评估。
• RandomSampler 和 SequentialSampler 是 PyTorch 提供的两种采样器。RandomSampler 随机采样数据，而 SequentialSampler 顺序采样数据。
• DistributedSampler 是用于分布式训练的数据采样器。
• default_data_collator 是 transformers 库提供的默认数据收集器，用于将多个样本整合为一个批量数据。
• create_pretrain_dataset 是一个自定义函数，用于创建预训练数据集。

通过检查 args.local_rank 是否为 −1，代码会选择使用普通的采样器（单机）还是分布式采样器（多机）。DistributedSampler 确保在分布式训练环境中，每个进程或节点都能获得数据的一个不重复的子集，这使得分布式训练变为可能。而在单机环境中，使用常规的随机或顺序采样器即可。具体代码如下所示：

模型载入

使用 transformers 库加载和配置 LLaMA 模型及其相关的词元分析器。从 transformers 库中导入 LLaMA 模型、相应的词元分析器和模型配置后，使用 from_pretrained 方法加载预训练的 LLaMA模型、词元分析器和配置。为了确保词元分析器可以处理各种文本的长度，还需要进行填充设置。如果词元分析器还没有指定填充符号，则将其设置为 [PAD]，并确定填充行为发生在句子的右侧。此外，为了保证模型能够正确地处理句子结束和填充，还为模型配置设置了结束符号和填充符号的 ID。最后，为了优化模型在硬件上的性能，还需要调整模型的词汇表嵌入大小，使其成为 8 的倍数。通过这些步骤，可以成功地加载并配置 LLaMA 模型，为后续的训练任务做好准备。具体代码如下：

优化器设置

DeepSpeed 库提供了高效的优化器算法，如 DeepSpeedCPUAdam 和 FusedAdam，这些算法经
过特殊优化以提高在大规模数据和模型上的训练速度。优化器配置主要包含以下几个方面。

（1）参数分组：通过 get_optimizer_grouped_parameters 函数将模型参数分为两组，一组使用权重衰减，另一组则不使用。这种参数分组有助于正则化模型，防止过拟合，并允许对特定参数应用不同的学习设置。
（2）优化器选择：根据训练设置（如是否在 CPU 上进行模型参数卸载），可以选择使用DeepSpeedCPUAdam 或 FusedAdam 优化器。这两种优化器都是对经典的 Adam 优化器进行优化和改进的版本，为大规模训练提供了高效性能。
（3）学习率调度：不同于固定的学习率，学习率调度器在训练过程中动态调整学习率。例如，在训练初期快速提高学习率以加速收敛，在训练中后期逐渐降低学习率以获得更精细的优化。我们的配置考虑了预热步骤、训练的总步数及其他关键因素。

具体代码如下所示：

DeepSpeed 设置

在配置代码的开始，定义了两个关键参数 GLOBAL_BATCH_SIZE 和 MICRO_BATCH_ SIZE。GLOBAL_BATCH_SIZE 定义了全局的批次大小。这通常是所有 GPU 加起来的总批次大小。MICRO_BATCH_SIZE定义了每块 GPU 上的微批次大小。因为微批次处理每次只加载并处理一小部分数据，所以可以帮助大语言模型在有限的 GPU 内存中运行。训练配置函数 get_train_ds_config主要包括以下内容。

（1）ZeRO 优化配置：ZeRO 是 DeepSpeed 提供的一种优化策略，旨在减少训练中的冗余并加速模型的训练。其中的参数，如 offload_param 和 offload_optimizer，允许用户选择是否将模型参数或优化器状态卸载到 CPU。
（2）混合精度训练：通过设置 FP16 字段，使模型可以使用 16 位浮点数进行训练，加速训练过程并减少内存使用。
（3）梯度裁剪：通过 gradient_clipping 字段，可以防止训练过程中出现梯度爆炸问题。
（4）混合引擎配置：hybrid_engine 部分允许用户配置更高级的优化选项，如输出分词的最大数量和推理张量的大小。
（5）TensorBoard 配置：使用 DeepSpeed 时，可以通过配置选项直接集成 TensorBoard，从而更方便地跟踪训练过程。
（6）验证集配置函数 get_eval_ds_config：此函数提供了 DeepSpeed 的验证集。与训练配置相比，验证集配置更为简洁，只需要关注模型推理阶段。

具体代码如下所示：

DeepSpeed 初始化

设置 DeepSpeed 的配置参数后，可以利用 DeepSpeed 进行模型训练的初始化，初始化流程如下。

（1）确定运行的设备：首先，检查代码是否有指定的本地 GPU（通过 args.local_rank）。如果没有指定，则程序默认使用 CUDA 设备。否则，它会为进程设置指定的 GPU。
（2）初始化分布式后端：在分布式训练中，使用 deepspeed.init_distributed() 函数实现每个进程与其他进程的同步，初始化分布式环境。
（3）获取当前进程的全局排序：在分布式训练中，使用 torch.distributed.get_rank() 函数获得每个进程的唯一排序或 ID。
（4）设置 DeepSpeed 配置：根据用户参数（如是否进行 offload、使用哪个 Zero Stage 等）构建一个 DeepSpeed 配置字典，来决定训练设置。
（5）同步所有工作进程：使用 torch.distributed.barrier() 确保在进一步的初始化之前所有进程都已同步。
（6）DeepSpeed 初始化：这是最关键的一步。通过 deepspeed.initialize 函数，可以将模型、优化器、参数和先前构建的 DeepSpeed 配置传递给库，进行初始化。这个函数会返回一个已经根据DeepSpeed 配置进行了优化的模型和优化器。
（7）梯度检查点：对于特别大的模型，梯度检查点是一种节省显存的技巧，即只在需要时计算模型的中间梯度。如果用户启用了这个选项，则会调用 model.gradient_checkpointing_enable() 方法来实现相关功能。

具体代码如下所示：

模型训练

借助 DeepSpeed 框架实现对模型的训练，训练步骤大致分为以下几个阶段。

（1）训练前的准备：使用 print_rank_0 函数输出当前的训练状态。该函数确保只有指定的进程（通常是主进程）会打印消息，避免了多进程环境下的重复输出。在开始训练之前，对模型进行一次评估，计算模型的困惑度。
（2）训练循环：每个周期的开始，都会打印当前周期和总周期数。在每次迭代中，数据批次先被移动到相应的 GPU 设备，接着模型对这个批次进行前向传播计算损失。使用 model.backward(loss)计算梯度，并使用 model.step() 更新模型参数。对于主进程，还会使用 print_throughput 函数打印吞吐量，这有助于了解模型的训练速度和效率。
（3）保存模型：如果指定了输出目录，则模型的状态和配置将被保存。模型可以在不同的格式中保存，例如 HuggingFace 的模型格式或 DeepSpeed 的 Zero Stage 3 特定格式。save_hf_format函数用于保存模型为 HuggingFace 格式，这意味着训练后的模型可以使用 HuggingFace 的from_pretrained 方法直接加载。对于 Zero Stage 3，save_zero_three_model 函数负责保存，因为在这个阶段，每个 GPU 只保存了模型的一部分。

具体代码如下所示：

参考文献