大模型并行训练指南：通俗理解Megatron-DeepSpeed之模型并行与数据并行

前言

本文可以看做是本文《千亿参数开源大模型 BLOOM 背后的技术，这是其英文原文》与相关论文的解读，但修正了原文中的部分细节错误，以及补充了大量的解释说明，使得其读起来一目了然、通俗易懂

第一部分 BLOOM与其背后的Megatron-DeepSpeed

1.1 BLOOM的训练细节：硬件/Checkpoints/数据集

BLOOM 的模型架构与 GPT3 非常相似，只是增加了一些改进，176B BLOOM 模型的训练于 2022 年 3 月至 7 月期间，耗时约 3.5 个月完成 (约 100 万计算时)，以下是其训练的一些细节信息

训练硬件

GPU: 384 张 NVIDIA A100 80GB GPU (48 个节点) + 32 张备用 GPU
每个节点 8 张 GPU，4 条 NVLink 卡间互联，4 条 OmniPath 链路
CPU: AMD EPYC 7543 32 核处理器
CPU 内存: 每个节点 512GB
GPU 显存: 每个节点 640GB
节点间连接: 使用 Omni-Path Architecture (OPA) 网卡，网络拓扑为无阻塞胖树
NCCL - 通信网络: 一个完全专用的子网
磁盘 IO 网络: GPFS 与其他节点和用户共享

Checkpoints

主 checkpoints
每个 checkpoint 含精度为 fp32 的优化器状态和精度为 bf16+fp32 的权重，占用存储空间为 2.3TB。如只保存 bf16 的权重，则仅占用 329GB 的存储空间。

数据集

41.5TB 经过大量去重和清洗的文本，包含 46 种语言，最终转换为 350B 个词元
模型的词汇表含 250,680 个词元
更详细信息，请参阅 The BigScience Corpus A 1.6TB Composite Multilingual Dataset

1.2 Megatron-DeepSpeed：

176B BLOOM 模型使用 Megatron-DeepSpeed 进行训练，它结合了两种主要技术:

Megatron-LM 是由 NVIDIA 应用深度学习研究团队开发的大型、强大的 transformer 模型框架，对应论文为：《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》，这是李沐对其的视频解读，这是对其的文字/代码解读
DeepSpeed 是微软开发的一个深度学习优化库，让分布式训练变得简单、高效且有效

DeepSpeed团队通过将“下面第一项与后面三项相结合”，开发了一种基于3D并行的实现，这就是Megatron-Deepspeed，它使得千亿级参数量以上的大规模语言模型比如BLOOM的分布式训练变得更简单、高效和有效

Megatron-LM中的张量并行(Tensor Parallelism，可以理解为模型并行的一种)
每个张量都被分成多个块，因此张量的每个分片都位于其指定的 GPU 上，而不是让整个张量驻留在单个 GPU 上。在处理过程中，每个分片在不同的 GPU 上分别并行处理，结果在步骤结束时同步。这就是所谓的水平并行，因为是做的水平拆分
零冗余优化器 (Zero Redundancy Optimizer，简称ZeRO，是微软DeepSpeed库的核心)
也执行与 TP 相类似的张量分片，但整个张量会及时重建以进行前向或反向计算，因此不需要修改模型。它还支持各种卸载技术以补偿有限的 GPU 内存
数据并行(Data Parallelism)
相同的设置和模型被复制多份，每份每次都被馈送不同的一份数据。处理是并行完成的，所有份在每个训练步结束时同步
管道并行(也称流水线并行，Pipeline Parallelism)
模型在多个 GPU 上垂直 (即按层) 拆分，因此只有一个或多个模型层放置在单个 GPU 上。每个 GPU 并行处理流水线的不同阶段，并处理 batch 的一部分数据

请注意，BLOOM团队BigScience版本的Megatron-DeepSpeed 基于原始 Megatron-DeepSpeed 代码库，但在其上添加了不少代码

下表列出了在训练 BLOOM 时各采用了两个框架的哪些组件

组件	DeepSpeed	Megatron-LM
ZeRO 数据并行	是
张量并行		是
流水线并行	是
BF16 优化器	是
CUDA 融合核函数		是
DataLoader		是

其中，Megatron-LM 和 DeepSpeed 都有流水线并行和 BF16 优化器实现，但我们使用 DeepSpeed 的实现，因为它们集成进了 ZeRO

第二部分张量并行(Tensor Parallelism，算模型并行的一种)

在张量并行 (TP) 中，每个 GPU 仅处理张量的一部分，并且仅当某些算子需要完整的张量时才触发聚合操作。

众所周知，Transformer 的主要模块就两块：一个自注意力层 + 残差连接，一个全连接层MLP + 残差连接

2019年，英伟达通过 Megatron-LM 论文(Efficient Large-Scale Language Model Training on GPU Clusters)，可以将其点积部分写为 Y = GeLU (XA) ，其中和是输入和输出向量，是权重矩阵。

如果以矩阵形式表示的话，很容易看出矩阵乘法可以如何在多个 GPU 之间拆分，如下图所示(记为图1)：

2.1 MLP的并行化：权重A矩阵竖着切 B矩阵横着切最后MERGE

Megatron-LM 论文作者有一个示意图(记为图2)：与矩阵相乘，得到的结果 X A 再与矩阵相乘(其中，矩阵按列竖着拆，这里是前向传播中的恒等运算符，后向传播中的 all reduce，而是前向传播中的 all reduce 和后向传播中的恒等式)

别小看上面这个示意图，其实里面的细节很多，值得反复揣摩，具体如下图所示(记为图3)

对于输入，它的行数是批量大小乘以序列长度，列数是隐藏层的宽度即
其隐藏层的模块里面其实就是两个全连接层
$\rightarrow$ 假定第一个隐藏层的权重是的话(行数为K，列数为K'，比如4行2列)，则先做矩阵乘法 $X\cdot A$ ，然后再接一个激活函数 $\sigma$ 比如GELU (GELU类似把ReLU的拐点顺滑下)
$\rightarrow$ 假定第二个隐藏层的权重是的话(行数K'，列数K，比如4行一列)，最终 $\sigma (X\cdot A)B = Y$
接下来，我们看下如何做拆分，好用上多个GPU做并行
如果是输入数据比较大，则优先选择做数据并行，即对输入做切分
如果是模型本身比较大，则优先选择做模型并行，即对矩阵做拆分，拆分方式分两种
$\rightarrow$ 第一种(对应上上图图1中的下半部分)：矩阵按行横着拆(则相应的按列竖着拆，导致的结果是两个GPU之间需要通讯)
$\rightarrow$ 第二种(对应上上图图1中的上半部分)：矩阵按列竖着拆，如来自图3的下图所示，一列蓝色，一列绿色(则相应的按行横着拆，或者在两个GPU上都得有一份才行，此时则无需任何额外的通信)
确定按第二种拆分方式之后(即矩阵按列竖着拆)，乘以得到一个大矩阵 $X\cdot A$ ，然后把矩阵按行切开(如来自图3的下图所示，第二个矩阵即矩阵B中，蓝色的放在GPU 0上，绿色的放在GPU 1上)
最终矩阵的列与矩阵的行做矩阵乘法，得到的大小与的大小一致，但结果只有(指的是GPU的数量)
换言之，通过执行矩阵乘法先得到到，然后再得到个输出向量 $Y_1 Y_2 \cdots Y_n$ ，它们可以独立输入GeLU：
最终merge上述第3步的结果最终得到完整的

通过上述操作，我们可以更新任意深度的 MLP，只需在每个拆列-拆行序列之后同步 GPU

2.2 多头注意力层的并行化：各个头各自计算

并行化多头注意力层甚至更简单，因为它们本来就是并行的，因为有多个独立的头！如下图所示(记为图4)

针对输入矩阵，其行数依然是批量大小乘以序列长度(假定批量大小为1)，列数是，在自注意力机制中(如果你忘了什么是自注意力机制，请详看这篇Transformer笔记的第三部分)，输入会被复制成三份，分别对应为：的向量矩阵(类似三个分身)
至于多头注意力，头的维度为，假定，之后针对每个头中输入矩阵中各个单词的向量，会与各自上下文的向量做缩放点积然后做softmax得到一个注意力分数或权重，之后再与做加权和，得到一个 $l \times k/h$ 的输出
最后再与一个 $k\times k$ 的投影做相乘，得到一个 $l \times k$ 的结果

第二个头的计算过程类似
如此，你会发现，每个头的计算是各自独立并行的，互不影响，那意味着一个头可以放在GPU 0上(蓝色表示)，另一个头可以放在GPU 1上(绿色表示)，最后all reduce每个头的结果

整个过程还可以如下图所示(记为图5)

需要特别考虑的是：

由于前向和后向传播中每层都有两个 all reduce，因此 TP 需要设备间有非常快速的互联。因此，除非你有一个非常快的网络，否则不建议跨多个节点进行 TP
我们训练 BLOOM 的硬件配置中，节点间的速度比 PCIe 慢很多。实际上，如果节点有 4 个 GPU，则最高 TP 度设为 4 比较好。如果需要 TP 度为 8，则需要使用至少有 8 个 GPU 的节点
该组件由 Megatron-LM 实现。Megatron-LM 最近扩展了张量并行能力，新增了序列并行的能力，用于难以使用前述切分算法的算子，如 LayerNorm(该论文Reducing Activation Recomputation in Large Transformer Models 提供了此技术的详细信息)
不过，由于序列并行是在训练 BLOOM 之后开发的，所以 BLOOM 训练时并未采用此技术

2.3 针对输入与输出的并行化

接下来，我们看下针对输入与输出的并行化，如下图所示(记为图6)

对于输入，是一个批量大小乘以序列长度的矩阵，里面存放着一行一行的句子，而embedding层是一个行数是vocab(相当于是一个词表)、列数是K的词表
整个字典可以横着切(比如上半部分放到GPU 0上标蓝表示，下半部分放到GPU 1上标绿表示)，通过查表得到一个 $b\times l \times k$ 的输出
对于输出，行数是 $b\times l$ ，列数是K，经过词表之后，得到一个 $b\times l \times v$ 的输出，其左半部分可以放到GPU 0上，右半部分可以放到GPU 1上
对于 $b\times l \times v$ 的输出的每一行可以横向加起来，但V可能会比较大比如几万的大小，当然，每个GPU可以各自算自己的部分

第三部分流水线并行Pipeline Parallelism(模型并行的另一种)

朴素流水线并行 (naive PP) 是将模型各层分组分布在多个 GPU 上，并简单地将数据从 GPU 移动到 GPU，就好像它是一个大型复合 GPU 一样。该机制相对简单 - 将所需层用 .to() 方法绑到相应设备，现在只要数据进出这些层，这些层就会将数据切换到与该层相同的设备，其余部分保持不变

这其实就是垂直模型并行(类似画大多数模型的拓扑图，垂直切分模型各层的)，例如，下图显示一个 8 层模型(记为图7)：

===================  ===================
|  0 | 1 | 2 | 3  |  |  4 | 5 | 6 | 7  |
===================  ===================
        GPU0                 GPU1

我们将它垂直切成 2 部分，将层 0-3 放置在 GPU0 上，将层 4-7 放置在 GPU1 上

现在，当数据从第 0 层传到第 1 层、第 1 层传到第 2 层以及第 2 层传到第 3 层时，这就跟单 GPU 上的普通前向传播一样
但是当数据需要从第 3 层传到第 4 层时，它需要从 GPU 0 传输到 GPU1，这会引入通信开销
$\rightarrow$ 如果参与的 GPU 位于同一计算节点 (例如同一台物理机器) 上，则传输非常快
$\rightarrow$ 但如果 GPU 位于不同的计算节点 (例如多台机器) 上，通信开销可能会大得多
然后第 4 到 5 到 6 到 7 层又像普通模型一样，当第 7 层完成时，我们通常需要将数据发送回标签所在的第 0 层 (或者将标签发送到最后一层)
之后，便可以计算损失，然后使用优化器来进行更新参数了

该方法为什么被称为朴素流水线并行呢，它又有什么缺陷呢？

主要是因为该方案在任意给定时刻除了一个 GPU 之外的其他所有 GPU 都是空闲的。因此，如果使用 4 个 GPU，则几乎等同于将单个 GPU 的内存量翻两番，而其他资源 (如计算) 相当于没用上。另外还需要加上在设备之间复制数据的开销
所以 4 张使用朴素流水线并行的 6GB 卡将能够容纳与 1 张 24GB 卡相同大小的模型，而后者训练得更快，因为它没有数据传输开销。但是，比如说，如果你有 40GB 卡，但需要跑 45GB 模型，你可以使用 4x 40GB 卡(也就刚刚够用，因为还有梯度和优化器状态需要显存)
共享嵌入可能需要在 GPU 之间来回复制。使用的流水线并行 (PP) 与上述朴素 PP 几乎相同，但它解决了 GPU 闲置问题，方法是将传入的 batch 分块为 micros batch 并人工创建流水线，从而允许不同的 GPU 同时参与计算过程

下图来自于 GPipe 论文 (记为图8)，图上把网络分成4块，每一块放在一个GPU上，不同的颜色表示不同的GPU)，于是就有了 F0、F1、F2、F3 这 4 个管级的前向路径，然后是 B3、B2、B1、B0 的逆序后向路径

其b部分表示朴素 PP 方案
为一般的模型并行的运算模式，在每个时间点只有一台设备在处理计算逻辑，完成计算后将结果发送给下一台设备
c部分是 PP 方法
PP 引入了一个新的超参数来调整，称为 块 (chunks)。它定义了通过同一管级按顺序发送多少数据块。例如，在图8.(c)的部分，你可以看到 chunks = 4。GPU 0 在 chunk 0、1、2 和 3 (F0,0、F0,1、F0,2、F0,3) 上执行相同的前向路径，然后等待，等其他 GPU 完成工作后，GPU 0 会再次开始工作，为块 3、2、1 和 0 (B0,3、B0,2、B0,1、B0,0) 执行后向路径
请注意，从概念上讲，这与梯度累积 (gradient accumulation steps，GAS) 的意思相同。PyTorch 叫它 块，而 DeepSpeed 叫它 GAS。

因为 块，PP 引入了 micro-batches (MBS) 的概念。DP 将全局 batch size 拆分为小 batch size，因此如果 DP 度为 4，则全局 batch size 1024 将拆分为 4 个小 batch size，每个小 batch size 为 256 (1024/4)。而如果 块 (或 GAS) 的数量为 32，我们最终得到的 micro batch size 为 8 (256/32)。每个管级一次处理一个 micro batch
计算 DP + PP 设置的全局批量大小的公式为: mbs*chunks*dp_degree (8*32*4=1024)

GPipe实验发现，通过将mini-batch进一步划分成更小的micro-batch，同时利用pipipline方案，每次处理一个micro-batch的数据，得到结果后，将该micro-batch的结果发送给下游设备，同时开始处理后一个 micro-batch的数据，通过这套方案减小设备中的Bubble(设备空闲的时间称为 Bubble)

第四部分数据并行与ZeRO

4.1 数据并行

简言之，模型小、数据大则适合采用数据并行的训练方式。此时，每个 GPU 都会复制一份模型的参数(模型被完全复制到每个 GPU，然后在每次迭代后所有模型相互同步各自的状态)，我们只需要把训练数据均分给多个不同的 GPU，然后让每个 GPU 作为一个计算节点独立的完成前向和反向传播运算

数据并行(DistributedDataParallel，简称DDP，这是相应的 PyTorch 文档)，不仅通信量较小，而且可以很方便的做通信计算重叠，因此可以取得最好的加速比

4.2 ZeRO：数据并行下，对冗余空间进行深度优化

4.2.1 大规模训练中的显存占用问题

在模型并行、数据并行、流水钱并行这三种并行训练方案中，数据并行因其易用性，得到了最为广泛的应用。然而，数据并行会产生大量冗余 Model States 的空间占用。ZeRO 的本质，是在数据并行的基础上，对冗余空间占用进行深度优化

包括ZeRO在我司金融审稿、英文审稿等项目中也会经常用到，特此也分享如下

比如，24年5月上旬，我司金融微调审稿项目组的乐事遇到一个问题，之前基于qwen微调时最大上下文长度为4k多就OOM，后面我加了一张卡微调时，使用deepspeed的zero2和zero3，量化4bit，甚至offload能微调的最大长度大概6k多
按照他的理解跟多卡没关系，因为都是batchsize为1，按理来说在使用deepspeed zero3+量化的情况下能微调6k，如果开启offload可以把优化器状态，梯度以及模型权重参数下沉到cpu内存上面，减少gpu显存开销，这样就可以继续加大最大长度
不过实验室看日志是开启了，但效果好像没起作用，可能需要再研究下zero3开启offload的参数试试了
对此，我司英文审稿项目组的青睐表示，他在做审稿GPT的的时候也遇到了这个问题，当时没有用s2atten，48G的显卡放不下，于是他找了2张48G和很大的cpu内存，使用 zero3 + cpu offload 依然会报OOM，然后查了相关资料
才了解到，zero3 + cpu offload 本质上还是数据并行，如果一个层的计算量大于卡的极限，不管多少张卡都是OOM的，这个时候就要用到张量并行或者序列并行了

大规模训练中的显存占用可以分为 Model States 与 Activation 两部分，而 ZeRO 就是为了解决 Model States 而诞生的一项技术。那模型在训练过程中 Model States 是由什么组成的呢？

Optimizer States
Optimizer States 是 Optimizer 在进行梯度更新时所需要用到的数据，例如 SGD 中的 Momentum 以及使用混合精度训练时的 Float32 Master Parameters
Gradient
在反向传播后所产生的梯度信息，其决定了参数的更新方向。
Model Parameter
模型参数，也就是我们在整个过程中通过数据“学习”的信息

在传统数据并行下，每个进程都使用同样参数来进行训练。每个进程也会持有对 Optimizer States 的完整拷贝，同样占用了大量显存。在混合精度场景下，以参数量为 Ψ 的模型和 Adam optimzier 为例，Adam 需要保存：

Float16 的参数和梯度的备份。这两项分别消耗了 2Ψ 和 2Ψ Bytes 内存(1 Float16 = 2 Bytes)
Float32的参数、Momentum、Variance备份，对应到 3 份 4Ψ 的内存占用(1 Float32 = 4 Bytes)

最终需要 2Ψ + 2Ψ + 12Ψ = 16Ψ bytes 的显存。一个 7.5B 参数量的模型，就需要至少 120GB 的显存空间才能装下这些 Model States。当数据并行时，这些重复的 Model States 会在 N 个 GPU 上复制 N 份

4.2.2 ZeRO数据并行的三种形态

2020年，微软DeepSpeed团队通过论文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》提出Zero Redundancy Optimizer(简称ZeRO)，但优化是永恒的话题，所以DeepSpeed团队在过去几年发表了三篇ZeRO相关的论文，提出了去除冗余参数、引入CPU和内存、引入NVMe等方法，从始至终都围绕着一个目标：将显存优化进行到底，使用 ZeRO 后，各个进程之后只保存完整状态的 1/GPUs，互不重叠，不再存在冗余

下图图9很好地描述了 ZeRO 数据并行 (来自这篇博文)，ZeRO 有三个不同级别，分别对应对 Model States 不同程度的分割 (Paritition)，图中的 $P_{os}$ 、 $P_{os+g}$ 、 $P_{os+g+p}$ 分别代表 ZeRO 1 、 ZeRO 2 、 ZeRO 3

ZeRO-1：分割 Optimizer States
ZeRO-2：分割 Optimizer States 与 Gradients
ZeRO-3：分割 Optimizer States、Gradients 与 Parameters

论文中的这个图并不好懂，为方便大家更好的理解，我借用李沐画的下图(记为图10)给大家再完整解释一下

假设用两块卡做数据并行，且就一个层，一台机器存：fp32的w以及adam里面的状态，以及一个fp16的w，再加一个fp16的已经算好的梯度
怎么做切分呢？很简单，上图图10左上角划线的部分存在GPU 0上，右上角划线的部分存在GPU 1上
此时存储的开销是 $12\psi /Nd$ ，此即为，表示机器的个数
现在，对梯度也得做同样的分割，还是两块卡，上图图10左下角左侧划线的部分存GPU 0上，右下角右侧划线的部分存GPU 1上
这就叫，能够把存梯度的 $2\psi$ ，降低到 $2\psi /Nd$

由于图10左下角右侧那块的梯度，GPU 0不维护它，(按箭头所指)便可以把它发送到能维护它的GPU1 做累加
类似的，图10右下角左侧那块的梯度，GPU 1不维护它，(按箭头所指)便可以把它发送到能维护它的GPU 0做累加
如此每个GPU都能获取到全局的梯度
接下来，要对状态进行更新，更新完状态之后，每个卡把自己维护的那一更新好的给发出去(且为方便大家文图对应，故又把上图图10 又重新贴下)
言下之意，即是上图图10总共六个框
第一行两个框，A B C D
第二行两个框，a b c d
第三行也两个框

第一行第一个框的左侧部分A直接copy给第二行第一个框的左侧部分a (↓)，然后第一行第一个框的左侧部分A发送给第二行第二个框的左侧部分c (即按箭头↘所示)
第一行第二个框的右侧部分D直接copy给第二行第二个框的右侧部分d (↓)，然后第一行第二个框的右侧部分D发送给第二行第一个框的右侧部分b (即按箭头↙所示)
对于图10 第二行那两个框而言，GPU 0和GPU 1各自维护各自的那一块
a b c d
但是要计算的时候
GPU 0去问GPU 1，让GPU 1的右侧部分d发送给GPU 0的右侧部分b(即按箭头↩所示)
GPU 1去问GPU 0，让GPU 0的左侧部分a发送给GPU 1的左侧部分c(即按箭头↪所示)
这一系列操作即称为，如此把计算的复杂度从 $2\psi$ ，降低到 $2\psi /Nd$

ZeRO 1 和 ZeRO 2 并不会带来额外的通讯，但 ZeRO3 每一次要算的时候，都得去别的机器拿回来，相当于带来了额外的通讯(增加了50%)

待更..

第五部分从DP+PP、DP+PP+TP到ZeRO DP+PP+TP

5.1 DP+PP

DeepSpeed 流水线并行教程中有一张图演示了如何将 DP 与 PP 结合起来，如下所示。

这里重要的是要了解 DP rank 0 是看不见 GPU2 的， DP rank 1 是看不到 GPU3 的。对于 DP 而言，只有 GPU 0 和 1，并向它们馈送数据。GPU 0 使用 PP “秘密地” 将它的一些负载卸载到 GPU2。同样地， GPU1 也会得到 GPU3 的帮助

由于每个维度至少需要 2 个 GPU，因此这儿至少需要 4 个 GPU

5.2 DP+PP+TP

为了更高效地训练，可以将 PP、TP 和 DP 相结合，称为 3D 并行，如下图所示(此图来自博文3D 并行: 扩展到万亿参数模型)

由于每个维度至少需要 2 个 GPU，因此在这里你至少需要 8 个 GPU 才能实现完整的 3D 并行

5.3 ZeRO DP+PP+TP

DeepSpeed 的主要功能之一是 ZeRO，它是 DP 的超级可伸缩增强版，我们在 ZeRO 数据并行一节中已经讨论过了。通常它是一个独立的功能，不需要 PP 或 TP。但它也可以与 PP、TP 结合使用

当 ZeRO-DP 与 PP (以及 TP) 结合时，它通常只启用 ZeRO 1，它只对优化器状态进行分片。ZeRO 2 还会对梯度进行分片，ZeRO 3 也对模型权重进行分片
虽然理论上可以将 ZeRO 2 与流水线并行一起使用，但它会对性能产生不良影响。每个 micro batch 都需要一个额外的 reduce-scatter 通信来在分片之前聚合梯度，这会增加潜在的显著通信开销。根据流水线并行的性质，我们会使用小的 micro batch ，并把重点放在算术强度 (micro batch size) 与最小化流水线气泡 (micro batch 的数量) 两者间折衷。因此，增加的通信开销会损害流水线并行
此外，由于 PP，层数已经比正常情况下少，因此并不会节省很多内存。PP 已经将梯度大小减少了 1/PP，因此在此基础之上的梯度分片和纯 DP 相比节省不了多少内存
ZeRO 阶段 3 也可用于训练这种规模的模型，但是，它需要的通信量比 DeepSpeed 3D 并行更多。一年前，在对我们的环境进行仔细评估后，我们发现 Megatron-DeepSpeed 3D 并行性表现最佳。此后，ZeRO 阶段 3 的性能有了显著提高，如果我们今天要对其进行重新评估，也许我们会选择阶段 3

第六部分 BF16Optimizer/CUDA 融合核函数/数据集/嵌入LayerNorm/位置编码

6.1 优先使用BF16 Optimizer，而非FP16

用 FP16 训练巨型 LLM 模型是一个禁忌

我们已经通过花费几个月的时间训练 104B 模型自证了这一点，你可以从 Tensorboard 发现，彻头彻尾地失败了。在与不断发散的 lm-loss 作斗争的过程中，我们学到了很多：

我们也从 Megatron-LM 和 DeepSpeed 团队那里得到了相同的建议，在他们训得 530B 模型后。最近发布的 OPT-175B 也报告说他们在 FP16 上训练得非常艰难

所以早在一月份，我们就知道我们要在支持 BF16 格式的 A100 上进行训练。Olatunji Ruwase 开发了一个用来训练 BLOOM 的 BF16Optimizer，如果您不熟悉这种数据格式，请查看它的位布局

BF16 格式的关键是它的指数位数与 FP32 相同，因此不会溢出，但 FP16 经常溢出

$\rightarrow$ FP16 的最大数值范围为 64k，您只能进行较小数的乘法。例如你可以做 250*250=62500，但如果你尝试 255*255=65025，你就会溢出，这是导致训练出现问题的主要原因。这意味着你的权重必须保持很小。一种称为损失缩放(loss scaling)的技术有助于缓解这个问题，但是当模型变得非常大时，FP16 较小的数值范围仍然是一个问题
$\rightarrow$ BF16 没有这个问题，你可以很容易地做 10_000*10_000=100_000_000

有意思是的，我司七月审稿项目组在用一万多条paper-review数据微调llama2 7B时也遇到了这个问题，即
“原本的LongQLoRA源码训练参数指定使用fp16数据类型，但是训练可能会很不稳定
具体而言，可能会出现数值溢出问题，初期就出现loss严重震荡，甚至高达上百，并在数步后返回模型已收敛(但实际并未)的提示而中断训练，故要设置使用bf16数据类型进行训练，loss就能稳定从4.多开始收敛”
更多详见此文《七月论文审稿GPT第2版：用一万多条paper-review数据集微调LLaMA2 7B最终反超GPT4》的5.1节LLaMA2 7b chat + LongQLoRA训练
当然，由于 BF16 和 FP16 的大小相同，均为 2 个字节，因此，没有免费的午餐，当使用 BF16 时，代价就是它的精度非常差。然而，你应该还记得我们在训练时采用的随机梯度下降法及其变体，该方法有点像蹒跚而行，如果你这步没有找到完美的方向其实没关系，你会在接下来的步骤中纠正自己
无论使用 BF16 还是 FP16，都有一个权重副本始终在 FP32 中 —— 这是由优化器更新的内容。因此 16 位格式仅用于计算，优化器以全精度更新 FP32 权重，然后将它们转换为 16 位格式以用于下一次迭代

所有 PyTorch 组件都已更新，以确保它们在 FP32 中执行任何累加，因此不会发生精度损失
一个关键问题是梯度累积，它是流水线并行的主要特征之一，因为每个 micro batch 处理的梯度都会累积。在 FP32 中实现梯度累积以保证训练的精确性至关重要，这正是 BF16Optimizer 所做的

除了其他改进之外，我们认为使用 BF16 混合精度训练将潜在的噩梦变成了一个相对平稳的过程，这可以从以下 lm 损失图中看出:

6.2 CUDA 融合核函数

GPU 主要做两件事。它可以将数据写到显存或从显存读数据，并对这些数据执行计算。当 GPU 忙于读写数据时， GPU 的计算单元就会空闲。如果我们想有效地利用 GPU，我们希望将空闲时间降至最低

核函数是一组实现特定 PyTorch 操作的指令。例如，当你调用 torch.add 时，它会通过一个 PyTorch 调度器，它会根据输入张量及其他变量的取值来决定它应该运行哪些代码，最后运行它。CUDA 核函数使用 CUDA 来实现这些代码，因此只能在 NVIDIA GPU 上运行

比如举个例子，当使用 GPU 计算 c = torch.add (a, b); e = torch.max ([c,d]) 时
$\rightarrow$ 一般情况下，PyTorch 将执行的操作是启动两个单独的核函数，一个执行 a 和 b 的加法，另一个执行取 c 和 d 两者的最大值。在这种情况下，GPU 从其显存中获取 a 和 b，执行加法运算，然后将结果写回显存(一次对a b的读取一次对c的写入)。然后它获取 c 和 d 并执行 max 操作，然后再次将结果写回显存(再一次对c d的读取一次对max[c d]结果的写入)
$\rightarrow$ 但如果我们要融合这两个操作，即将它们放入一个 “融合核函数” 中，然后启动那个内核，我们不会将中间结果 c 写到显存中，而是将其保留在 GPU 寄存器中(仅一次对 a b的读取)，并且仅需要获取 d 来完成最后的计算(一次对d的读取一次对max[c d]结果的写入)。这节省了大量开销并防止 GPU 空闲，因此整个操作会更加高效
融合核函数就是这样。它们主要将多个离散的计算和进出显存的数据移动替换为有很少数据移动的融合计算。此外，一些融合核函数会对操作进行数学变换，以便可以更快地执行某些计算组合

July注：这其实有点类似flash attention的思想，详见此文：通透理解FlashAttention与FlashAttention2：全面降低显存读写、加快计算速度
为了快速高效地训练 BLOOM，有必要使用 Megatron-LM 提供的几个自定义 CUDA 融合核函数。特别地，有一个 LayerNorm 的融合核函数以及用于融合缩放、掩码和 softmax 这些操作的各种组合的核函数。Bias Add 也通过 PyTorch 的 JIT 功能与 GeLU 融合。这些操作都是瓶颈在内存的，因此将它们融合在一起以达到最大化每次显存读取后的计算量非常重要。因此，例如，在执行瓶颈在内存的 GeLU 操作时同时执行 Bias Add，运行时间并不会增加。这些核函数都可以在 Megatron-LM repository 代码库中找到

6.3 数据集

Megatron-LM 的另一个重要特性是高效的数据加载器。在首次训练启动前，每个数据集中的每个样本都被分成固定序列长度(BLOOM 为 2048)的样本，并创建索引以对每个样本进行编号。基于训练超参，我们会确定每个数据集所需要参与的 epoch 数，并基于此创建一个有序的样本索引列表，然后打乱它

举个例子，如果一个数据集中有 10 个样本并应参与 2 个 epoch 的训练，则

系统首先按 [0, ..., 9, 0, ..., 9] 顺序排好样本索引，然后打乱该顺序为数据集创建最终的全局顺序
请注意，这意味着训练不会简单地遍历整个数据集然后重复，你有可能在看到另一个样本之前看到同一个样本两次，但在训练结束时模型将只看到每个样本两次。这有助于确保整个训练过程中的训练曲线平滑。这些索引，包括每个样本在原始数据集中的偏移量，被保存到一个文件中，以避免每次开始训练时都重新计算它们
最后，可以将其中几个数据集以不同的权重混合到训练最终使用的数据中

6.4 嵌入LayerNorm

在我们努力阻止 104B 模型发散的过程中，我们发现在第一个层词嵌入层之后添加一个额外的 LayerNorm 可以使训练更加稳定

该洞察来自对 bitsandbytes 的实验，bitsandbytes 有一个 StableEmbedding 操作，它是一个带有 LayerNorm 的普通嵌入，其使用均匀 xavier 函数来初始化

6.5 位置编码

基于论文 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation，我们还用 AliBi 替换了普通的位置嵌入，它允许外推比训练模型的输入序列更长的输入序列。因此，即使我们训练时使用长度为 2048 的序列，模型也可以在推理过程中处理更长的序列

第七部分其他事项：训练中的困难与重要资源

7.1 训练中的困难

随着架构、硬件和软件的就位，我们得以在 2022 年 3 月上旬开始训练。然而，从那时起，事情其实并非一帆风顺。在本节中，我们将讨论我们遇到的一些主要障碍

在训练开始之前，有很多问题需要弄清楚。特别是，我们发现了几个问题，这些问题只有在我们开始在 48 个节点上进行训练后才会出现，而不会在小规模时出现

例如，需要设 CUDA_LAUNCH_BLOCKING=1 来防止框架挂起，我们需要将优化器组分成更小的组，否则框架会再次挂起。你可以在训前编年史中详细了解这些内容
训练期间遇到的主要问题类型是硬件故障。由于这是一个拥有大约 400 个 GPU 的新集群，平均每周我们会遇到 1-2 个 GPU 故障。我们每 3 小时 (100 次迭代) 保存一个检查点。因此，我们每周因硬件崩溃平均损失 1.5 小时的训练成果。Jean Zay 系统管理员随后将更换有故障的 GPU 并恢复节点。与此同时，我们有备用节点可供使用。
我们还遇到过多次导致 5-10 小时停机的各种其他问题，其中一些与 PyTorch 中的死锁错误有关，另一些则是由于磁盘空间不足。如果您对具体细节有兴趣，请参阅训练编年史
在对训练这个模型进行可行性分析时，所有这些停机时间都被计划在内了，我们也据此选择了合适的模型大小和我们希望模型消耗的数据量。因此，即使存在这些停机问题，我们还是成功地在预计时间内完成了训练。如前所述，它需要大约 100 万个计算时才能完成
另一个问题是 SLURM 并非设计为供一组人使用。SLURM 作业由单个用户拥有，如果他们不在身边，则该组的其他成员无法对正在运行的作业执行任何操作。我们制定了一个终止方案，允许组中的其他用户终止当前进程，而不需要启动该进程的用户在场。这在 90% 的问题上都很有效。如果 SLURM 设计者读到这篇文章，请添加一个 Unix 组的概念，这样一个 SLURM 作业就可以由一个组拥有
由于训练是全天候 24/7 进行的，我们需要有人随叫随到，但由于我们在欧洲和加拿大西海岸都有人，因此不需要有人携带传呼机，我们能很好地互相备份。当然，周末的训练也得有人看着。我们自动化了大部分事情，包括自动从硬件崩溃中恢复，但有时仍需要人工干预

7.2 重要链接

7.3 论文与文章

我们不可能在本文中详细解释所有内容，因此如果此处介绍的技术激起你的好奇心，使你想了解更多信息，请阅读以下论文:

Megatron-LM:

DeepSpeed:

Megatron-LM 和 Deepspeeed 联合:

Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model.

ALiBi:

Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
What Language Model to Train if You Have One Million GPU Hours? - 你会在那里找到最终选择 ALiBi 的实验

BitsNBytes:

8-bit Optimizers via Block-wise Quantization (使用了该论文中的嵌入 LaynerNorm，但是论文的其他部分及其技术也很妙，没用 8 位优化器的唯一原因是我们已经使用 DeepSpeed-ZeRO 节省了优化器内存)