BLOOM 训练背后的技术

近年来，训练越来越大的语言模型已成为常态。虽然这些模型未发布以供进一步研究的问题经常被讨论，但关于如何训练这些模型的隐藏知识却很少受到关注。本文旨在改变这一现状，以 1760 亿参数的语言模型 BLOOM 为例，从硬件和软件两方面阐述训练此类模型背后的技术和工程。

但首先，我们要感谢那些公司、关键人物和团队，是他们使得一小群敬业的人能够完成训练一个 1760 亿参数模型的惊人壮举。

然后将讨论硬件设置和主要技术组件。

以下是项目概览

人员

该项目由 Thomas Wolf (Hugging Face 联合创始人兼 CSO) 构思，他敢于与大公司竞争，不仅要训练出最大的多语言模型之一，还要让最终成果惠及所有人，从而将大多数人的梦想变为现实。

本文特别关注模型训练的工程方面。BLOOM 背后最重要的技术是那些分享专业知识、帮助我们进行编码和训练的人员和公司。

主要有 6 组人员需要感谢：

1. Hugging Face 的 BigScience 团队，他们投入了超过六名全职员工，从项目开始到结束负责规划和运行训练，并提供和支付了 Jean Zay 计算资源之外的所有基础设施费用。
2. 微软 DeepSpeed 团队，他们开发了 DeepSpeed 并后来将其与 Megatron-LM 集成。其开发人员花费数周时间满足项目需求，并在训练前后提供了大量宝贵的实践经验建议。
3. NVIDIA Megatron-LM 团队，他们开发了 Megatron-LM，非常热心地回答了我们的大量问题，并提供了一流的经验建议。
4. 管理 Jean Zay 超级计算机的 IDRIS / GENCI 团队，他们为项目捐赠了海量的计算资源和出色的系统管理支持。
5. PyTorch 团队，他们创建了一个功能强大的框架，其他软件都基于此框架。他们在我们准备训练期间给予了大力支持，修复了多个错误，并改善了我们训练中依赖的 PyTorch 组件的可用性。
6. BigScience 工程工作组的志愿者们。

很难一一列举所有为该项目工程方面做出贡献的杰出人士，因此我只提及几位在过去 14 个月中作为该项目工程基石的 Hugging Face 之外的关键人物：

Olatunji Ruwase, Deepak Narayanan, Jeff Rasley, Jared Casper, Samyam Rajbhandari 以及 Rémi Lacroix。

我们也要感谢所有允许其员工为该项目做出贡献的公司。

概述

BLOOM 的架构与 GPT3 非常相似，并增加了一些改进，本文稍后会讨论。

该模型在 Jean Zay 上进行训练，这是一台由法国政府资助的超级计算机，由 GENCI 管理，并安装在法国国家科学研究中心 (CNRS) 的国家计算中心 IDRIS。计算资源由 GENCI 慷慨捐赠给该项目 (拨款号 2021-A0101012475)。

训练期间使用了以下硬件：

• GPU：384 块 NVIDIA A100 80GB GPU (48 个节点) + 32 块备用 GPU
• 每节点 8 块 GPU，使用 NVLink 4 进行 GPU 间互联，4 个 OmniPath 链接
• CPU：AMD EPYC 7543 32 核处理器
• CPU 内存：每节点 512GB
• GPU 内存：每节点 640GB
• 节点间连接：Omni-Path 架构 (OPA)，采用无阻塞胖树结构
• NCCL 通信网络：一个完全专用的子网
• 磁盘 IO 网络：GPFS 与其他节点和用户共享

检查点

• 主要检查点
• 每个包含 fp32 优化器状态和 bf16+fp32 权重的检查点大小为 2.3TB——仅 bf16 权重就为 329GB。

数据集

• 46 种语言，1.5TB 经过去重和大规模清洗的文本，转换为 3500 亿个唯一的词元
• 模型的词汇表大小为 250,680 个词元
• 完整详情请参见 The BigScience Corpus A 1.6TB Composite Multilingual Dataset

1760 亿参数的 BLOOM 模型训练发生在 2022 年 3 月至 7 月，耗时约 3.5 个月 (约 100 万计算小时)。

Megatron-DeepSpeed

1760 亿参数的 BLOOM 模型是使用 Megatron-DeepSpeed 进行训练的，它结合了两种主要技术：

• DeepSpeed 是一个深度学习优化库，使分布式训练变得简单、高效和有效。
• Megatron-LM 是一个由 NVIDIA 应用深度学习研究团队开发的大型、强大的 Transformer 模型框架。

DeepSpeed 团队通过将 DeepSpeed 库中的 ZeRO 分片和流水线并行与 Megatron-LM 中的张量并行相结合，开发了一种基于 3D 并行的实现。下表中可以看到每个组件的更多细节。

请注意，BigScience 的 Megatron-DeepSpeed 是原始 Megatron-DeepSpeed 仓库的一个分支，我们在此基础上添加了多项功能。

下表列出了训练 BLOOM 时各个组件由哪个框架提供：

请注意，Megatron-LM 和 DeepSpeed 都有流水线并行和 BF16 优化器的实现，但我们使用了 DeepSpeed 的版本，因为它们与 ZeRO 集成在一起。

Megatron-DeepSpeed 实现了 3D 并行，以非常高效的方式训练大型模型。让我们简要讨论一下这 3 个维度。

1. 数据并行 (DP) - 相同的设置被复制多次，每个副本处理一部分数据。处理过程并行进行，在每个训练步骤结束时所有设置进行同步。
2. 张量并行 (TP) - 每个张量被分割成多个块，因此整个张量不再驻留在单个 GPU 上，而是每个分片驻留在其指定的 GPU 上。在处理过程中，每个分片在不同的 GPU 上被独立且并行地处理，结果在步骤结束时同步。这可以称为水平并行，因为切分是在水平层面上发生的。
3. 流水线并行 (PP) - 模型在多个 GPU 上垂直 (层级) 拆分，因此只有一个或几个模型层被放置在单个 GPU 上。每个 GPU 并行处理流水线的不同阶段，并处理一小块批次数据。
4. 零冗余优化器 (ZeRO) - 也执行张量的分片，有点类似于 TP，不同之处在于整个张量会在前向或后向计算时被及时重构，因此模型不需要被修改。它还支持各种卸载技术以弥补有限的 GPU 内存。

数据并行

大多数只有几块 GPU 的用户可能都熟悉 `DistributedDataParallel` (DDP) PyTorch 文档。在这种方法中，模型被完全复制到每个 GPU 上，然后在每次迭代后，所有模型相互同步它们的状态。这种方法通过投入更多资源来加快训练速度，但它只有在模型能装入单个 GPU 的情况下才有效。

ZeRO 数据并行

基于 ZeRO 的数据并行 (ZeRO-DP) 在这篇 博文 中的下图有描述

这可能很难理解，但实际上，这个概念相当简单。这只是普通的 DDP，不同之处在于，不是复制完整的模型参数、梯度和优化器状态，每个 GPU 只存储其中的一部分。然后在运行时，当某个特定层需要完整的层参数时，所有 GPU 同步，相互提供它们缺少的部分——就是这样。

该组件由 DeepSpeed 实现。

张量并行

在张量并行 (TP) 中，每个 GPU 只处理张量的一个分片，并且仅在需要整个张量的操作中才聚合完整的张量。

在本节中，我们使用 Megatron-LM 论文中的概念和图表：Efficient Large-Scale Language Model Training on GPU Clusters。

任何 Transformer 的主要构建块是一个全连接层 `nn.Linear`，后跟一个非线性激活函数 `GeLU`。

遵循 Megatron 论文的表示法，我们可以将其点积部分写为 `Y = GeLU(XA)`，其中 `X` 和 `Y` 是输入和输出向量，`A` 是权重矩阵。

如果我们以矩阵形式看待计算，很容易看出矩阵乘法如何在多个 GPU 之间分割：

如果我们将权重矩阵 `A` 按列在 `N` 个 GPU 上分割，并并行执行矩阵乘法 `XA_1` 到 `XA_n`，那么我们将得到 `N` 个输出向量 `Y_1, Y_2, ..., Y_n`，它们可以独立地输入到 `GeLU` 中：。请注意，当 Y 矩阵按列分割时，我们可以将第二个 GEMM 按行分割，使其可以直接接收 GeLU 的输出，而无需任何额外的通信。

利用这个原理，我们可以更新任意深度的 MLP，同时在每个行列序列后同步 GPU。Megatron-LM 论文的作者为此提供了一个有用的插图：

这里 `f` 在前向传播中是恒等操作符，在后向传播中是 all-reduce 操作；而 `g` 在前向传播中是 all-reduce 操作，在后向传播中是恒等操作。

并行化多头注意力层甚至更简单，因为它们由于有多个独立的头，本身就是并行的！

特别注意事项：由于在前向和后向传播中每层都有两次 all-reduce 操作，TP 需要设备之间非常快速的互连。因此，不建议在多个节点之间进行 TP，除非你有非常快的网络。在我们的案例中，节点间通信比 PCIe 慢得多。实际上，如果一个节点有 4 个 GPU，那么最高的 TP 度就是 4。如果你需要 8 的 TP 度，你需要使用至少有 8 个 GPU 的节点。

该组件由 Megatron-LM 实现。Megatron-LM 最近扩展了张量并行，加入了序列并行，该技术将无法像上述那样拆分的操作 (如 LayerNorm) 沿序列维度进行拆分。论文 Reducing Activation Recomputation in Large Transformer Models 提供了该技术的详细信息。序列并行是在 BLOOM 训练后开发的，因此未在 BLOOM 训练中使用。

流水线并行

朴素流水线并行 (naive PP) 是指将模型层的分组分布在多个 GPU 上，然后像在一个大的复合 GPU 上一样简单地将数据从一个 GPU 移动到另一个 GPU。机制相对简单——将所需的层 `.to()` 到所需的设备上，然后每当数据进出这些层时，将数据切换到与层相同的设备上，其余部分保持不变。

这执行的是垂直模型并行，因为如果你还记得大多数模型的绘制方式，我们是垂直切分层的。例如，如果下图显示一个 8 层模型：

===================  ===================
|  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 层时，它需要从 GPU0 传输到 GPU1，这就引入了通信开销。如果参与的 GPU 在同一个计算节点上 (例如同一台物理机器)，这种复制速度相当快，但如果 GPU 位于不同的计算节点上 (例如多台机器)，通信开销可能会大得多。

然后，第 4 层到第 5 层，再到第 6 层，最后到第 7 层，都像正常模型一样。当第 7 层完成后，我们通常需要将数据送回第 0 层，那里有标签 (或者将标签发送到最后一层)。现在可以计算损失，优化器可以开始工作。

问题

• 主要缺陷以及为什么这种方法被称为“朴素”PP，是在任何给定时刻只有一个 GPU 在工作。因此，如果使用 4 个 GPU，这几乎等同于将单个 GPU 的内存量增加四倍，而忽略了其余的硬件。此外，还存在设备之间复制数据的开销。所以，4 张 6GB 的显卡可以容纳与 1 张 24GB 显卡相同大小的模型，但后者会更快地完成训练，因为它没有数据复制的开销。但是，比如说，如果你有 40GB 的显卡，需要容纳一个 45GB 的模型，你可以用 4 张 40GB 的显卡 (但由于梯度和优化器状态的原因，会很勉强)。
• 共享的嵌入层可能需要在 GPU 之间来回复制。

流水线并行 (PP) 与上面描述的朴素 PP 几乎相同，但它通过将输入批次分块为微批次并人为地创建一个流水线来解决 GPU 空闲问题，这使得不同的 GPU 可以同时参与计算过程。

以下来自 GPipe 论文 的插图显示了顶部的朴素 PP 和底部的 PP。

从底部图表中可以很容易地看出 PP 如何减少了 GPU 空闲的死区。这些空闲部分被称为“气泡”。

图表的两个部分都显示了并行度为 4 的情况。也就是说，有 4 个 GPU 参与流水线。因此，有 4 个流水线阶段的前向路径 F0、F1、F2 和 F3，然后是反向路径 B3、B2、B1 和 B0。

PP 引入了一个需要调整的新超参数，称为 `chunks`。它定义了通过同一个流水线阶段连续发送的数据块数量。例如，在底部图表中，你可以看到 `chunks=4`。GPU0 对块 0、1、2 和 3 执行相同的前向路径 (F0,0、F0,1、F0,2、F0,3)，然后等待其他 GPU 完成它们的工作，只有当它们的工作开始完成时，GPU0 才再次开始工作，为块 3、2、1 和 0 执行后向路径 (B0,3、B0,2、B0,1、B0,0)。

注意，从概念上讲，这与梯度累积步数 (GAS) 是相同的概念。PyTorch 使用 `chunks`，而 DeepSpeed 将同一个超参数称为 GAS。

由于 `chunks` 的存在，PP 引入了微批次 (MBS) 的概念。DP 将全局数据批次大小划分为小批次，所以如果你的 DP 度为 4，全局批次大小为 1024 会被划分为 4 个 256 的小批次 (1024/4)。如果 `chunks` (或 GAS) 的数量是 32，我们最终得到的微批次大小为 8 (256/32)。每个流水线阶段一次只处理一个微批次。

要计算 DP + PP 设置的全局批次大小，我们这样做：`mbs * chunks * dp_degree` (`8 * 32 * 4 = 1024`)。

让我们回到图表。

当 `chunks=1` 时，你得到的是朴素 PP，效率非常低。当 `chunks` 值非常大时，你得到的微批次大小会非常小，这可能也不高效。因此，必须通过实验找到能带来最高效 GPU 利用率的值。

虽然图表显示存在一个无法并行的“死区”时间气泡，因为最后一个`前向`阶段必须等待`后向`完成流水线，但找到最佳 `chunks` 值的目的在于实现所有参与 GPU 的高并发利用率，这相当于最小化气泡的大小。

这种调度机制被称为 `all forward all backward`。其他一些替代方案是 one forward one backward 和 interleaved one forward one backward。

虽然 Megatron-LM 和 DeepSpeed 都有自己的 PP 协议实现，但 Megatron-DeepSpeed 使用 DeepSpeed 的实现，因为它与 DeepSpeed 的其他方面集成在一起。

另一个重要问题是词嵌入矩阵的大小。虽然通常词嵌入矩阵消耗的内存比 Transformer 块少，但在我们拥有 25 万词汇表的巨大情况下，嵌入层在 bf16 权重中需要 7.2GB，而 Transformer 块仅为 4.9GB。因此，我们必须指示 Megatron-Deepspeed 将嵌入层视为一个 Transformer 块。所以我们有一个 72 层的流水线，其中 2 层专用于嵌入 (第一层和最后一层)。这使得 GPU 内存消耗得以平衡。如果我们不这样做，第一和最后一个阶段会消耗大部分 GPU 内存，而 95% 的 GPU 将使用少得多的内存，因此训练效率会很低。

DP+PP

以下来自 DeepSpeed 流水线教程 的图表演示了如何将 DP 与 PP 结合起来。

在这里，重要的是要看到 DP rank 0 看不到 GPU2，DP rank 1 看不到 GPU3。对 DP 来说，只有 GPU 0 和 GPU 1，它向它们馈送数据，就好像只有 2 个 GPU 一样。GPU0 “秘密地”使用 PP 将其部分负载卸载到 GPU2。而 GPU1 也通过寻求 GPU3 的帮助来做同样的事情。

由于每个维度至少需要 2 个 GPU，这里你至少需要 4 个 GPU。

DP+PP+TP

为了获得更高效的训练，PP 与 TP 和 DP 结合使用，这被称为 3D 并行。这可以在下图中看到。

这张图来自一篇博文 3D 并行：扩展至万亿参数模型，这也是一篇很好的读物。

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

ZeRO DP+PP+TP

DeepSpeed 的主要特性之一是 ZeRO，它是 DP 的一个超可扩展的扩展。这已经在 ZeRO 数据并行 中讨论过。通常它是一个独立的功能，不需要 PP 或 TP。但它可以与 PP 和 TP 结合使用。

当 ZeRO-DP 与 PP (以及可选的 TP) 结合时，通常只启用 ZeRO stage 1，它只分片优化器状态。ZeRO stage 2 额外分片梯度，而 stage 3 也分片模型权重。

虽然理论上可以将 ZeRO stage 2 与流水线并行一起使用，但这会对性能产生不良影响。每个微批次都需要一个额外的 reduce-scatter 集合操作来在分片前聚合梯度，这会增加潜在的显著通信开销。根据流水线并行的特性，会使用小的微批次，而重点在于平衡计算强度 (微批次大小) 与最小化流水线气泡 (微批次数量)。因此，这些通信成本将会造成损害。

此外，由于 PP 的原因，层数已经比正常情况少，因此内存节省不会很大。PP 已经将梯度大小减少了 `1/PP`，因此在此基础上的梯度分片节省比纯 DP 的情况要小。

ZeRO stage 3 也可以用来训练这种规模的模型，但是，它比 DeepSpeed 3D 并行实现需要更多的通信。在一年前我们对环境进行仔细评估后，我们发现 Megatron-DeepSpeed 3D 并行表现最好。从那时起，ZeRO stage 3 的性能得到了显著提升，如果今天我们再次评估，也许我们会选择 stage 3。

BF16Optimizer

用 FP16 训练大型语言模型是不可行的。

我们通过花费数月时间 训练一个 1040 亿参数模型 证明了这一点，从 tensorboard 可以看出，那次训练完全失败了。在与不断发散的 lm-loss 作斗争的过程中，我们学到了很多东西。

在 Megatron-LM 和 DeepSpeed 团队训练了 5300 亿参数模型 后，我们也从他们那里得到了同样的建议。最近发布的 OPT-175B 也报告说他们在 FP16 训练中遇到了很大的困难。

因此，早在 1 月份，我们就知道我们将在支持 BF16 格式的 A100 上进行训练，Olatunji Ruwase 开发了一个 `BF16Optimizer`，我们用它来训练 BLOOM。

如果你不熟悉这种数据格式，请查看位布局。BF16 格式的关键在于它与 FP32 有相同的指数，因此不会遭受 FP16 经常遇到的溢出问题！对于 FP16，其最大数值范围为 64k，你只能乘以较小的数字。例如，你可以做 `250*250=62500`，但如果你尝试 `255*255=65025`，就会发生溢出，这是训练期间的主要问题。这意味着你的权重必须保持很小。一种称为损失缩放的技术可以帮助解决这个问题，但当模型变得非常大时，FP16 的有限范围仍然是一个问题。

BF16 没有这样的问题，你可以轻松地做 `10_000*10_000=100_000_000`，毫无问题。

当然，由于 BF16 和 FP16 的大小都是 2 字节，天下没有免费的午餐，使用 BF16 会付出精度非常差的代价。然而，如果你还记得使用随机梯度下降及其变体的训练过程是一种蹒跚前行的方式，所以如果你没有立即得到完美的方向也没关系，你会在接下来的步骤中自我修正。

无论使用 BF16 还是 FP16，总会有一个 FP32 的权重副本——这是由优化器更新的。所以 16 位格式只用于计算，优化器用全精度更新 FP32 权重，然后在下一次迭代中将它们转换为 16 位格式。

所有 PyTorch 组件都已更新，以确保它们在 FP32 中执行任何累加，因此不会有损失。

一个关键问题是梯度累积，这也是流水线并行的主要特性之一，因为来自每个微批次处理的梯度会被累积。在 FP32 中实现梯度累积以保持训练的精确性至关重要，而这正是 `BF16Optimizer` 所做的。

除了其他改进之外，我们相信使用 BF16 混合精度训练将一个潜在的噩梦变成了一个相对顺利的过程，这可以从以下 lm loss 图中观察到：

融合 CUDA 核心

GPU 执行两件事。它可以将数据复制到内存或从内存中复制数据，并对这些数据进行计算。当 GPU 忙于复制时，GPU 的计算单元处于空闲状态。如果我们想高效地利用 GPU，我们希望最小化空闲时间。

一个核心 (kernel) 是一组实现特定 PyTorch 操作的指令。例如，当你调用 `torch.add` 时，它会通过一个 PyTorch 调度器，该调度器会查看输入的张量和其他各种东西，然后决定应该运行哪个代码，并执行它。CUDA 核心是一种使用 CUDA API 库的特定实现，只能在 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 操作的各种组合的核心。偏置项的添加也使用 PyTorch 的 JIT 功能与 GeLU 操作融合。这些操作都是内存密集型的，因此将它们融合以最大化从内存中检索到值后完成的计算量非常重要。例如，在已经进行内存密集型的 GeLU 操作时添加偏置项不会增加额外的时间。这些核心都可以在 Megatron-LM 仓库中找到。

数据集

Megatron-LM 的另一个重要特性是高效的数据加载器。在初次训练启动时，每个数据集被分割成请求序列长度 (BLOOM 为 2048) 的样本，并创建一个索引来编号每个样本。根据训练参数计算数据集的 epoch 数，并为那么多 epoch 创建一个顺序，然后进行洗牌。例如，如果一个数据集有 10 个样本，应该遍历两次，系统首先按顺序排列样本索引 `[0, ..., 9, 0, ..., 9]`，然后打乱这个顺序以创建该数据集的最终全局顺序。请注意，这意味着训练不会简单地遍历整个数据集然后重复，有可能在看到另一个样本之前看到同一个样本两次，但在训练结束时，模型将看到每个样本两次。这有助于确保整个训练过程中的平滑训练曲线。这些索引，包括每个样本在基础数据集中的偏移量，都会保存到一个文件中，以避免每次启动训练过程时重新计算。然后可以将这些数据集中的几个以不同的权重混合到训练过程最终看到的数据中。

嵌入层归一化

在我们努力阻止 1040 亿参数模型发散时，我们发现在第一个词嵌入层之后立即添加一个额外的层归一化 (LayerNorm) 会使训练更加稳定。

这一见解来自于对 bitsandbytes 的实验，该库包含一个 `StableEmbedding`，它是一个带有层归一化的普通嵌入层，并使用均匀 xavier 初始化。

位置编码

我们还将通常的位置嵌入替换为 AliBi——基于论文：Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation，这使得模型能够外推到比训练时更长的输入序列。因此，即使我们在长度为 2048 的序列上进行训练，模型在推理时也能处理更长的序列。

训练中的困难

随着架构、硬件和软件的就位，我们于 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 全天候进行的，我们需要有人值班——但由于我们在欧洲和加拿大西海岸都有人，总体上不需要有人携带传呼机，我们的工作时间可以很好地重叠。当然，周末也需要有人照看训练。我们自动化了大部分事情，包括从硬件崩溃中恢复，但有时也需要人工干预。

结论

训练中最困难和最紧张的部分是开始训练前的两个月。我们面临着尽快开始训练的巨大压力，因为资源分配时间有限，而且直到最后一刻我们才接触到 A100。所以那是一段非常艰难的时期，考虑到 `BF16Optimizer` 是在最后一刻编写的，我们需要调试它并修复各种错误。正如上一节所解释的，我们发现了只有在 48 个节点上开始训练时才会出现的新问题，这些问题在小规模下不会出现。

但是一旦我们解决了这些问题，训练本身就出奇地顺利，没有出现重大问题。大多数时候我们只有一个人监控训练，只有几次需要几个人参与故障排除。我们得到了 Jean Zay 管理部门的大力支持，他们迅速解决了训练期间出现的大多数需求。

总的来说，这是一次超级紧张但非常有益的经历。

训练大型语言模型仍然是一项具有挑战性的任务，但我们希望通过开放地构建和分享这项技术，其他人可以在我们的经验基础上继续发展。

资源

重要链接

• 主要训练文档
• tensorboard
• 训练 slurm 脚本
• 训练纪事

论文和文章

我们不可能在本文中详细解释所有内容，所以如果这里介绍的技术激起了您的好奇心，并且您想了解更多，这里有一些值得阅读的论文：

Megatron-LM

• 在 GPU 集群上高效训练大规模语言模型 (Efficient Large-Scale Language Model Training on GPU Clusters).
• 减少大型 Transformer 模型中的激活重计算 (Reducing Activation Recomputation in Large Transformer Models)

DeepSpeed

• ZeRO：内存优化以训练万亿参数模型
• ZeRO-Offload：民主化百亿级模型训练
• ZeRO-Infinity：突破 GPU 内存墙，实现超大规模深度学习
• DeepSpeed：面向所有人的超大规模模型训练 (DeepSpeed: Extreme-scale model training for everyone)

Megatron-LM 和 DeepSpeed 联合 (Joint Megatron-LM and Deepspeeed)

• 使用 DeepSpeed 和 Megatron 训练 Megatron-Turing NLG 530B，一个大规模生成语言模型 (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)
• 如果你有一百万个 GPU 小时，应该训练什么样的语言模型？(What Language Model to Train if You Have One Million GPU Hours?) - 在那里你会找到导致我们选择 ALiBi 的实验。

BitsNBytes

• 通过块级量化实现 8 位优化器 (8-bit Optimizers via Block-wise Quantization) (在嵌入层归一化的背景下，但论文的其余部分和技术都非常出色——我们没有使用 8 位优化器的唯一原因是因为我们已经通过 DeepSpeed-ZeRO 节省了优化器内存)。

博客致谢

非常感谢以下各位朋友，他们提出了很好的问题，并帮助提高了文章的可读性 (按字母顺序列出)：Britney Muller, Douwe Kiela, Jared Casper, Jeff Rasley, Julien Launay, Leandro von Werra, Omar Sanseviero, Stefan Schweter and Thomas Wang。

主图由 Chunte Lee 创作。