乐趣区

关于人工智能:千亿参数开源大模型-BLOOM-背后的技术

jiezi

假如你当初有了数据，也搞到了估算，所有就绪，筹备开始训练一个大模型，一显身手了，“一朝看尽长安花”仿佛近在眼前 …… 且慢！训练可不仅仅像这两个字的发音那么简略，看看 BLOOM 的训练或者对你有帮忙。

近年来，语言模型越训越大已成为常态。大家通常会诟病这些大模型自身的信息未被公开以供钻研，但很少关注大模型训练技术这种背地的常识。本文旨在以 1760 亿参数的语言模型 BLOOM 为例，说明训练此类模型背地的软硬件工程和技术要点，以促成大家对大模型训练技术的探讨。

首先，咱们要感激促成或资助咱们这个小组最终实现了训练 1760 亿参数模型这一惊人壮举的公司、集体和个人。

而后，咱们开始探讨硬件配置和次要技术组件。

以下是对本我的项目的简要总结:
硬件 384 张 80GB A100 GPU
软件 Megatron-DeepSpeed
模型架构 基于 GPT3
数据集 含 59 种语言，共 3500 亿词元
训练时长 3.5 个月

人员组成

该我的项目由 Thomas Wolf (Hugging Face 联结创始人兼 CSO) 发想，他敢于与大公司竞争，提出不仅要训练出立于世界上最大的多语言模型之林的模型，还要让所有人都能够公开拜访训练后果，圆了大多数人的幻想。

本文次要关注模型训练的工程方面。BLOOM 背地的技术中最重要的局部是分享专业知识并帮忙咱们进行编码和训练的人员和公司。

咱们次要须要感激 6 个群体:

  1. HuggingFace 的 BigScience 团队投入了六名以上的全职员工全程参加了训练的钻研和运行，他们还提供或报销了 Jean Zay 计算机之外的所有基础设施。
  2. Microsoft DeepSpeed 团队，开发了 DeepSpeed，起初将其与 Megatron-LM 集成，其开发人员破费数周工夫钻研我的项目需要，并在训练前和训练期间提供了许多很棒的实用教训倡议。
  3. NVIDIA Megatron-LM 团队开发了 Megatron-LM，他们十分乐于答复咱们的大量问题并提供一流的应用倡议。
  4. IDRIS / GENCI 团队治理着 Jean Zay 超级计算机，他们为该我的项目捐献了大量的算力和弱小的系统管理反对。
  5. PyTorch 团队创立了一个超强的框架，其余软件都基于该框架，并且在筹备训练期间十分反对咱们，修复了多个 bug 并进步了咱们所依赖的 PyTorch 组件的训练可用性。
  6. BigScience 工程工作组志愿者

很难说出所有为该项目标工程方面做出奉献的杰出人物的名字，所以我只列举 Hugging Face 之外的几个关键人物，他们在过来 14 个月中为该我的项目奠定了工程根底:

Olatunji Ruwase、Deepak Narayanan、Jeff Rasley、Jared Casper、Samyam Rajbhandari 和 Rémi Lacroix

咱们也感激所有容许其员工为该我的项目做出奉献的公司。

概述

BLOOM 的模型架构与 GPT3 十分类似，只是减少了一些改良，本文稍后将对此进行探讨。

该模型是在 Jean Zay 上训练的，Jean Zay 是由 GENCI 治理的法国政府赞助的超级计算机，装置在法国国家科学研究核心 (CNRS) 的国家计算中心 IDRIS。训练所需的算力由 GENCI 慷慨捐献给本我的项目 (捐献号 2021-A0101012475)。

训练硬件:

  • 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 的存储空间。

数据集:

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

176B BLOOM 模型的训练于 2022 年 3 月至 7 月期间，耗时约 3.5 个月实现 (约 100 万计算时)。

Megatron-DeepSpeed

176B BLOOM 模型应用 Megatron-DeepSpeed 进行训练，它联合了两种次要技术:

  • Megatron-DeepSpeed:
  • DeepSpeed 是一个深度学习优化库，让分布式训练变得简略、高效且无效。
  • Megatron-LM 是由 NVIDIA 的利用深度学习钻研团队开发的大型、弱小的 transformer 模型框架。

DeepSpeed 团队通过将 DeepSpeed 库中的 ZeRO 分片和流水线并行 (Pipeline Parallelism) 与 Megatron-LM 中的张量并行 (Tensor Parallelism) 相结合，开发了一种基于 3D 并行的计划。无关每个组件的更多详细信息，请参见下表。

请留神，BigScience 的 Megatron-DeepSpeed 是基于原始 Megatron-DeepSpeed 代码库，咱们还在其上增加了不少代码。

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

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

请留神，Megatron-LM 和 DeepSpeed 都有流水线并行和 BF16 优化器实现，但咱们应用 DeepSpeed 的实现，因为它们集成进了 ZeRO。

Megatron-DeepSpeed 实现了 3D 并行以容许大模型以十分无效的形式进行训练。咱们简要讨论一下有哪些 3D 组件。

  1. 数据并行 (Data Parallelism，DP) – 雷同的设置和模型被复制多份，每份每次都被馈送不同的一份数据。解决是并行实现的，所有份在每个训练步完结时同步。
  2. 张量并行 (Tensor Parallelism，TP) – 每个张量都被分成多个块，因而张量的每个分片都位于其指定的 GPU 上，而不是让整个张量驻留在单个 GPU 上。在处理过程中，每个分片在不同的 GPU 上别离并行处理，后果在步骤完结时同步。这就是所谓的程度并行，因为是做的程度拆分。
  3. 流水线并行 (Pipeline Parallelism，PP) – 模型在多个 GPU 上垂直 (即按层) 拆分，因而只有一个或多个模型层搁置在单个 GPU 上。每个 GPU 并行处理流水线的不同阶段，并解决 batch 的一部分数据。
  4. 零冗余优化器 (Zero Redundancy Optimizer，ZeRO) – 也执行与 TP 相相似的张量分片，但整个张量会及时重建以进行前向或反向计算，因而不须要批改模型。它还反对各种卸载技术以弥补无限的 GPU 内存。

数据并行

大多数只有几张 GPU 的用户可能比拟相熟 DistributedDataParallel(DDP)，这是相应的 PyTorch 文档。在该办法中，模型被齐全复制到每个 GPU，而后在每次迭代后所有模型互相同步各自的状态。这种办法能够通过投入更多 GPU 资源的形式放慢训练速度，解决问题。但它有个限度，即只有当模型可能放进单个 GPU 时才无效。

ZeRO 数据并行

下图很好地形容了 ZeRO 数据并行 (来自 此博文)。

看上去比拟高大上，可能让你很难分心去了解，但实际上，这个概念非常简单。这只是通常的 DDP，只是没有每个 GPU 都复制残缺的模型参数、梯度和优化器状态，​​而是每个 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。咱们训练 BLOOM 的硬件配置中，节点间的速度比 PCIe 慢很多。实际上，如果节点有 4 个 GPU，则最高 TP 度设为 4 比拟好。如果须要 TP 度为 8，则须要应用至多有 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 层 (或者将标签发送到最初一层)。当初能够计算损失，而后应用优化器来进行更新参数了。

问题:

  • 该办法为什么被称为 奢侈 流水线并行呢，它又有什么缺点呢？次要是因为该计划在任意给定时刻除了一个 GPU 之外的其余所有 GPU 都是闲暇的。因而，如果应用 4 个 GPU，则简直等同于将单个 GPU 的内存量翻两番，而其余资源 (如计算) 相当于没用上。另外还须要加上在设施之间复制数据的开销。所以 4 张 应用奢侈流水线并行的 6GB 卡将可能包容与 1 张 24GB 卡雷同大小的模型，而后者训练得更快，因为它没有数据传输开销。然而，比如说，如果你有 40GB 卡，但须要跑 45GB 模型，你能够应用 4x 40GB 卡 (也就刚刚够用，因为还有梯度和优化器状态须要显存)。
  • 共享嵌入可能须要在 GPU 之间来回复制。
    咱们应用的流水线并行 (PP) 与上述奢侈 PP 简直雷同，但它解决了 GPU 闲置问题，办法是将传入的 batch 分块为 micros batch 并人工创立流水线，从而容许不同的 GPU 同时参加计算过程。

下图来自于 GPipe 论文，其上半局部示意奢侈 PP 计划，下半局部是 PP 办法:

从图的下半局部很容易看出 PP 的死区 (指 GPU 处于闲暇状态) 更少，即“气泡”更少。

图上两种计划的并行度均为 4，即由 4 张 GPU 组成流水线。于是就有了 F0、F1、F2、F3 这 4 个管级的前向门路，而后是 B3、B2、B1、B0 的逆序后向门路。

PP 引入了一个新的超参数来调整，称为 块 (chunks)。它定义了通过同一管级按程序发送多少数据块。例如，在图的下半局部，你能够看到 chunks = 4。GPU0 在 chunk 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) 执行后向门路。

请留神，从概念上讲，这与梯度累积 (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)。

咱们回过头再看一下图。

应用 chunks=1 你最终失去的是奢侈 PP，这是十分低效的。而应用十分大的 块 数，你最终会失去很小的微批量大小，这很可能也不是很无效。因而，必须通过试验来找到能最无效地利用 GPU 的 块数。

该图显示存在无奈并行化的“死”工夫气泡，因为最初一个 forward 阶段必须期待 backward 实现流水。那么，找到最佳的 块 数，从而使所有参加的 GPU 达到高的并发利用率，这一问题其实就转化为最小化气泡数了。

这种调度机制被称为 全前全后。其余一些可选计划有 一前一后 和 交织一前一后。

尽管 Megatron-LM 和 DeepSpeed 都有本人的 PP 协定实现，但 Megatron-DeepSpeed 应用的是 DeepSpeed 实现，因为它与 DeepSpeed 的其余性能集成在一起。

这里的另一个重要问题是词嵌入矩阵的大小。尽管通常词嵌入矩阵比 transfomer 块所需的内存更少，但在 BLOOM 有 250k 词汇表的状况下，嵌入层须要 7.2GB 的 bf16 权重，而变换器块仅为 4.9GB。因而，咱们不得不让 Megatron-Deepspeed 将嵌入层视为一个转换器块。所以咱们有一个 72 级的流水线，其中 2 个是专门用于嵌入的 (第一个和最初一个)。这使得咱们能够均衡 GPU 的内存耗费。如果咱们不这样做，咱们就会让第一级和最初一级耗费很大的 GPU 内存，而 95% 的 GPU 内存应用会很少，因而训练将很不高效。

DP+PP

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

这里重要的是要理解 DP rank 0 是看不见 GPU2 的，DP rank 1 是看不到 GPU3 的。对于 DP 而言，只有 GPU 0 和 1，并向它们馈送数据。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 阶段 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。

BF16 优化器

用 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 损失图中看出:

CUDA 交融核函数

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 代码库 中找到。

数据集

Megatron-LM 的另一个重要个性是高效的数据加载器。在首次训练启动前，每个数据集中的每个样本都被分成固定序列长度 (BLOOM 为 2048) 的样本，并创立索引以对每个样本进行编号。基于训练超参，咱们会确定每个数据集所须要参加的 epoch 数，并基于此创立一个有序的样本索引列表，而后打乱它。举个例子，如果一个数据集中有 10 个样本并应参加 2 个 epoch 的训练，则零碎首先按 [0, ..., 9, 0, ..., 9] 程序排好样本索引，而后打乱该程序为数据集创立最终的全局程序。请留神，这意味着训练不会简略地遍历整个数据集而后反复，你有可能在看到另一个样本之前看到同一个样本两次，但在训练完结时模型将只看到每个样本两次。这有助于确保整个训练过程中的训练曲线平滑。这些索引，包含每个样本在原始数据集中的偏移量，被保留到一个文件中，以防止每次开始训练时都从新计算它们。最初，能够将其中几个数据集以不同的权重混合到训练最终应用的数据中。

嵌入 LayerNorm

在咱们致力阻止 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 来避免框架挂起，咱们须要将优化器组分成更小的组，否则框架会再次挂起。你能够在 训前编年史 中具体理解这些内容。

训练期间遇到的次要问题类型是硬件故障。因为这是一个领有大概 400 个 GPU 的新集群，均匀每周咱们会遇到 1-2 个 GPU 故障。咱们每 3 小时 (100 次迭代) 保留一个检查点。因而，咱们每周因硬件解体均匀损失 1.5 小时的训练成绩。Jean Zay 系统管理员随后将更换有故障的 GPU 并复原节点。与此同时，咱们有备用节点可供使用。

咱们还遇到过屡次导致 5-10 小时停机的各种其余问题，其中一些与 PyTorch 中的死锁谬误无关，另一些则是因为磁盘空间有余。如果您对具体细节有趣味，请参阅 训练编年史。

在对训练这个模型进行可行性剖析时，所有这些停机工夫都被打算在内了，咱们也据此抉择了适合的模型大小和咱们心愿模型耗费的数据量。因而，即便存在这些停机问题，咱们还是胜利地在预计工夫内实现了训练。如前所述，它须要大概 100 万个计算时能力实现。

另一个问题是 SLURM 并非设计为供一组人应用。SLURM 作业由单个用户领有，如果他们不在身边，则该组的其余成员无奈对正在运行的作业执行任何操作。咱们制订了一个终止计划，容许组中的其余用户终止以后过程，而不须要启动该过程的用户在场。这在 90% 的问题上都很无效。如果 SLURM 设计者读到这篇文章，请增加一个 Unix 组的概念，这样一个 SLURM 作业就能够由一个组领有。

因为训练是全天候 24/7 进行的，咱们须要有人随叫随到 – 但因为咱们在欧洲和加拿大西海岸都有人，因而不须要有人携带传呼机，咱们能很好地相互备份。当然，周末的训练也得有人看着。咱们自动化了大部分事件，包含主动从硬件解体中复原，但有时仍须要人工干预。

论断

训练中最艰难和最缓和的局部是训练开始前的 2 个月。咱们接受着尽快开始训练的微小压力，因为资源分配的工夫无限，咱们直到最初一刻才接触到 A100。所以这是一个十分艰难的期间，思考到 BF16Optimizer 是在最初一刻编写进去的，咱们须要调试它并修复各种 bug。正如上一节所述，咱们发现了新问题，这些问题只有在咱们开始在 48 个节点上进行训练后才会呈现，并且不会在小规模时呈现。

然而一旦咱们把这些整顿完，训练自身出奇的顺利，没有呈现大的问题。大多数时候，咱们只有一个人看着，只有少数几个人参加故障排除。咱们失去了 Jean Zay 治理部门的大力支持，他们迅速解决了训练期间呈现的大部分需要。

总的来说，这是一次超级缓和但回报颇丰的经验。

训练大型语言模型依然是一项具备挑战性的工作，但咱们心愿通过公开构建和共享这项技术，其他人能够借鉴咱们的教训。

资源

重要链接

  • 主训练文档
  • Tensorboard
  • 训练用的 slurm 脚本
  • 训练编年史

论文与文章

咱们不可能在本文中具体解释所有内容，因而如果此处介绍的技术激发你的好奇心，使你想理解更多信息，请浏览以下论文:

Megatron-LM:

  • Efficient Large-Scale Language Model Training on GPU Clusters.
  • Reducing Activation Recomputation in Large Transformer Models

DeepSpeed:

  • ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
  • ZeRO-Offload: Democratizing Billion-Scale Model Training
  • ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
  • DeepSpeed: Extreme-scale model training for everyone

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 节俭了优化器内存)。

博文致谢

非常感谢以下这些人，他们提出了很好的问题并帮忙进步了文章的可读性 (按字母序):

  • Britney Muller,
  • Douwe Kiela,
  • Jared Casper,
  • Jeff Rasley,
  • Julien Launay,
  • Leandro von Werra,
  • Omar Sanseviero,
  • Stefan Schweter and
  • Thomas Wang.

本文图表次要由 Chunte Lee 创作。

英文原文: https://hf.co/blog/bloom-megatron-deepspeed

原文作者: Stas Bekman

译者: Matrix Yao (姚伟峰)，英特尔深度学习工程师，工作方向为 transformer-family 模型在各模态数据上的利用及大规模模型的训练推理。

审校、排版: zhongdongy (阿东)

退出移动版