本文可以看做是本文《千亿参数开源大模型 BLOOM 背后的技术,这是其英文原文》与相关论文的解读,但修正了部分细节错误,以及补充了大量的解释说明,使得其读起来一目了然、通俗易懂
BLOOM 的模型架构与 GPT3 非常相似,只是增加了一些改进,176B BLOOM 模型的训练于 2022 年 3 月至 7 月期间,耗时约 3.5 个月完成 (约 100 万计算时),以下是其训练的一些细节信息
训练硬件
Checkpoints
数据集
176B BLOOM 模型使用 Megatron-DeepSpeed 进行训练,它结合了两种主要技术:
DeepSpeed 是微软开发的一个深度学习优化库,让分布式训练变得简单、高效且有效
DeepSpeed团队通过将下面第一项与后面三项相结合
开发了一种基于3D并行的实现,这就是Megatron-Deepspeed,它使得千亿级参数量以上的大规模语言模型比如BLOOM的分布式训练变得更简单、高效和有效
请注意,BLOOM团队BigScience版本的Megatron-DeepSpeed 基于原始 Megatron-DeepSpeed 代码库,但在其上添加了不少代码
下表列出了我们在训练 BLOOM 时各采用了两个框架的哪些组件
组件 | DeepSpeed | Megatron-LM |
---|---|---|
ZeRO 数据并行 | 是 | |
张量并行 | 是 | |
流水线并行 | 是 | |
BF16 优化器 | 是 | |
CUDA 融合核函数 | 是 | |
DataLoader | 是 |
请注意,Megatron-LM 和 DeepSpeed 都有流水线并行和 BF16 优化器实现,但我们使用 DeepSpeed 的实现,因为它们集成进了 ZeRO
在张量并行 (TP) 中,每个 GPU 仅处理张量的一部分,并且仅当某些算子需要完整的张量时才触发聚合操作。
众所周知,Transformer 的主要模块就两块:一个自注意力层 + 残差连接,
一个全连接层MLP + 残差连接
2019年,英伟达通过 Megatron-LM 论文(Efficient Large-Scale Language Model Training on GPU Clusters),可以将其点积部分写为,其中
和
是输入和输出向量,
是权重矩阵。
如果以矩阵形式表示的话,很容易看出矩阵乘法可以如何在多个 GPU 之间拆分,如下图所示(记为图1):
别小看上面这个示意图,其实里面的细节很多,值得反复揣摩,具体如下图所示(记为图2)
:![\left[Y_{1}, Y_{2}\right]=\left[\operatorname{GeLU}\left(X A_{1}\right), \operatorname{GeLU}\left(X A_{2}\right)\right]](http://img.e-com-net.com/image/info8/fe15bd3e055f4985b22c5da2f3ce52dc.png)
拆列-拆行
序列之后同步 GPU。Megatron-LM 论文作者为此提供了一个不错的图示(记为图3): 这里 并行化多头注意力层甚至更简单,因为它们本来就是并行的,因为有多个独立的头!如下图所示(记为图4)
需要特别考虑的是:
- 由于前向和后向传播中每层都有两个 all reduce,因此 TP 需要设备间有非常快速的互联。因此,除非你有一个非常快的网络,否则不建议跨多个节点进行 TP。我们训练 BLOOM 的硬件配置中,节点间的速度比 PCIe 慢很多。实际上,如果节点有 4 个 GPU,则最高 TP 度设为 4 比较好。如果需要 TP 度为 8,则需要使用至少有 8 个 GPU 的节点
- 该组件由 Megatron-LM 实现。Megatron-LM 最近扩展了张量并行能力,新增了序列并行的能力,用于难以使用前述切分算法的算子,如 LayerNormReducing Activation Recomputation in Large Transformer Models 论文提供了此技术的详细信息。序列并行是在训练 BLOOM 之后开发的,所以 BLOOM 训练时并未采用此技术
接下来,我们看下针对输入与输出的并行化,如下图所示(记为图6)
朴素流水线并行 (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,这会引入通信开销。如果参与的 GPU 位于同一计算节点 (例如同一台物理机器) 上,则传输非常快,但如果 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 的逆序后向路径
块 (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) 执行后向路径块
,而 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 batchmbs*chunks*dp_degree
(8*32*4=1024
)简言之,模型小、数据大则适合采用数据并行的训练方式。此时,每个 GPU 都会复制一份模型的参数(模型被完全复制到每个 GPU,然后在每次迭代后所有模型相互同步各自的状态),我们只需要把训练数据均分给多个不同的 GPU,然后让每个 GPU 作为一个计算节点独立的完成前向和反向传播运算
数据并行(DistributedDataParallel,简称
DDP,这是相应的 PyTorch 文档),不仅通信量较小,而且可以很方便的做通信计算重叠,因此可以取得最好的加速比
在模型并行、数据并行、流水钱并行这三种并行训练方案中,数据并行因其易用性,得到了最为广泛的应用。然而,数据并行会产生大量冗余 Model States 的空间占用。ZeRO 的本质,是在数据并行的基础上,对冗余空间占用进行深度优化
大规模训练中的显存占用可以分为 Model States 与 Activation 两部分,而 ZeRO 就是为了解决 Model States 而诞生的一项技术。那模型在训练过程中 Model States 是由什么组成的呢?
在传统数据并行下,每个进程都使用同样参数来进行训练。每个进程也会持有对 Optimizer States 的完整拷贝,同样占用了大量显存。在混合精度场景下,以参数量为 Ψ 的模型和 Adam optimzier 为例,Adam 需要保存:
最终需要 2Ψ + 2Ψ + 12Ψ = 16Ψ bytes 的显存。一个 7.5B 参数量的模型,就需要至少 120GB 的显存空间才能装下这些 Model States。当数据并行时,这些重复的 Model States 会在 N 个 GPU 上复制 N 份
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),图中的、
、
分别代表
、
、
论文中的这个图并不好懂,为方便大家更好的理解,我借用李沐画的下图(记为图10)给大家再完整解释一下
和
并不会带来额外的通讯,但
每一次要算
的时候,都得去别的机器拿回来,相当于带来了额外的通讯(增加了50%)
待更..
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
为了更高效地训练,可以将 PP、TP 和 DP 相结合,称为 3D 并行,如下图所示(此图来自博文3D 并行: 扩展到万亿参数模型)
由于每个维度至少需要 2 个 GPU,因此在这里你至少需要 8 个 GPU 才能实现完整的 3D 并行
DeepSpeed 的主要功能之一是 ZeRO,它是 DP 的超级可伸缩增强版,我们在 ZeRO 数据并行 一节中已经讨论过了。通常它是一个独立的功能,不需要 PP 或 TP。但它也可以与 PP、TP 结合使用。
当 ZeRO-DP 与 PP (以及 TP) 结合时,它通常只启用 ZeRO 阶段 1,它只对优化器状态进行分片。ZeRO 阶段 2 还会对梯度进行分片,阶段 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。
用 FP16 训练巨型 LLM 模型是一个禁忌。
我们已经通过花费几个月的时间 训练 104B 模型 自证了这一点,你可以从 Tensorboard 发现,彻头彻尾地失败了。在与不断发散的 lm-loss 作斗争的过程中,我们学到了很多:
我们也从 Megatron-LM 和 DeepSpeed 团队那里得到了相同的建议,在他们训得 530B 模型 后。最近发布的 OPT-175B 也报告说他们在 FP16 上训练得非常艰难。
所以早在一月份,我们就知道我们要在支持 BF16 格式的 A100 上进行训练。Olatunji Ruwase 开发了一个用来训练 BLOOM 的 BF16Optimizer
。
如果您不熟悉这种数据格式,请查看它的 位布局。BF16 格式的关键是它的指数位数与 FP32 相同,因此不会溢出,但 FP16 经常溢出!FP16 的最大数值范围为 64k,您只能进行较小数的乘法。例如你可以做 250*250=62500
,但如果你尝试 255*255=65025
,你就会溢出,这是导致训练出现问题的主要原因。这意味着你的权重必须保持很小。一种称为损失缩放 (loss scaling) 的技术有助于缓解这个问题,但是当模型变得非常大时,FP16 较小的数值范围仍然是一个问题。
BF16 没有这个问题,你可以很容易地做 10_000*10_000=100_000_000
, 完全没问题。
当然,由于 BF16 和 FP16 的大小相同,均为 2 个字节,因此,没有免费的午餐,当使用 BF16 时,代价就是它的精度非常差。然而,你应该还记得我们在训练时采用的随机梯度下降法及其变体,该方法有点像蹒跚而行,如果你这步没有找到完美的方向其实没关系,你会在接下来的步骤中纠正自己。
无论使用 BF16 还是 FP16,都有一个权重副本始终在 FP32 中 —— 这是由优化器更新的内容。因此 16 位格式仅用于计算,优化器以全精度更新 FP32 权重,然后将它们转换为 16 位格式以用于下一次迭代。
所有 PyTorch 组件都已更新,以确保它们在 FP32 中执行任何累加,因此不会发生精度损失。
一个关键问题是梯度累积,它是流水线并行的主要特征之一,因为每个 micro batch 处理的梯度都会累积。在 FP32 中实现梯度累积以保证训练的精确性至关重要,这正是 BF16Optimizer
所做的。
除了其他改进之外,我们认为使用 BF16 混合精度训练将潜在的噩梦变成了一个相对平稳的过程,这可以从以下 lm 损失图中看出:
GPU 主要做两件事。它可以将数据写到显存或从显存读数据,并对这些数据执行计算。当 GPU 忙于读写数据时, GPU 的计算单元就会空闲。如果我们想有效地利用 GPU,我们希望将空闲时间降至最低。
核函数是一组实现特定 PyTorch 操作的指令。例如,当你调用 torch.add
时,它会通过一个 PyTorch 调度器,它会根据输入张量及其他变量的取值来决定它应该运行哪些代码,最后运行它。CUDA 核函数使用 CUDA 来实现这些代码,因此只能在 NVIDIA GPU 上运行。
现在,当使用 GPU 计算 c = torch.add (a, b); e = torch.max ([c,d])
时,一般情况下,PyTorch 将执行的操作是启动两个单独的核函数,一个执行 a
和 b
的加法,另一个执行取 c
和 d
两者的最大值。在这种情况下,GPU 从其显存中获取 a
和 b
,执行加法运算,然后将结果写回显存。然后它获取 c
和 d
并执行 max 操作,然后再次将结果写回显存。
如果我们要融合这两个操作,即将它们放入一个 “融合核函数” 中,然后启动那个内核,我们不会将中间结果 c
写到显存中,而是将其保留在 GPU 寄存器中,并且仅需要获取 d
来完成最后的计算。这节省了大量开销并防止 GPU 空闲,因此整个操作会更加高效。
融合核函数就是这样。它们主要将多个离散的计算和进出显存的数据移动替换为有很少数据移动的融合计算。此外,一些融合核函数会对操作进行数学变换,以便可以更快地执行某些计算组合。
为了快速高效地训练 BLOOM,有必要使用 Megatron-LM 提供的几个自定义 CUDA 融合核函数。特别地,有一个 LayerNorm 的融合核函数以及用于融合缩放、掩码和 softmax 这些操作的各种组合的核函数。Bias Add 也通过 PyTorch 的 JIT 功能与 GeLU 融合。这些操作都是瓶颈在内存的,因此将它们融合在一起以达到最大化每次显存读取后的计算量非常重要。因此,例如,在执行瓶颈在内存的 GeLU 操作时同时执行 Bias Add,运行时间并不会增加。这些核函数都可以在 Megatron-LM repository 代码库 中找到。
Megatron-LM 的另一个重要特性是高效的数据加载器。在首次训练启动前,每个数据集中的每个样本都被分成固定序列长度 (BLOOM 为 2048) 的样本,并创建索引以对每个样本进行编号。基于训练超参,我们会确定每个数据集所需要参与的 epoch 数,并基于此创建一个有序的样本索引列表,然后打乱它
举个例子,如果一个数据集中有 10 个样本并应参与 2 个 epoch 的训练,则系统首先按 [0, ..., 9, 0, ..., 9]
顺序排好样本索引,然后打乱该顺序为数据集创建最终的全局顺序。请注意,这意味着训练不会简单地遍历整个数据集然后重复,你有可能在看到另一个样本之前看到同一个样本两次,但在训练结束时模型将只看到每个样本两次。这有助于确保整个训练过程中的训练曲线平滑。这些索引,包括每个样本在原始数据集中的偏移量,被保存到一个文件中,以避免每次开始训练时都重新计算它们。最后,可以将其中几个数据集以不同的权重混合到训练最终使用的数据中
在我们努力阻止 104B 模型发散的过程中,我们发现在第一个层词嵌入层之后添加一个额外的 LayerNorm 可以使训练更加稳定
该洞察来自对 bitsandbytes 的实验,bitsandbytes 有一个 StableEmbedding
操作,它是一个带有 LayerNorm 的普通嵌入,其使用均匀 xavier 函数来初始化
基于论文 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation,我们还用 AliBi 替换了普通的位置嵌入,它允许外推比训练模型的输入序列更长的输入序列。因此,即使我们训练时使用长度为 2048 的序列,模型也可以在推理过程中处理更长的序列
随着架构、硬件和软件的就位,我们得以在 2022 年 3 月上旬开始训练。然而,从那时起,事情其实并非一帆风顺。在本节中,我们将讨论我们遇到的一些主要障碍
在训练开始之前,有很多问题需要弄清楚。特别是,我们发现了几个问题,这些问题只有在我们开始在 48 个节点上进行训练后才会出现,而不会在小规模时出现
CUDA_LAUNCH_BLOCKING=1
来防止框架挂起,我们需要将优化器组分成更小的组,否则框架会再次挂起。你可以在 训前编年史 中详细了解这些内容我们不可能在本文中详细解释所有内容,因此如果此处介绍的技术激起你的好奇心,使你想了解更多信息,请阅读以下论文:
Megatron-LM:
DeepSpeed:
Megatron-LM 和 Deepspeeed 联合:
ALiBi:
BitsNBytes: