BLOOM 176B 训练技术概览

TL;DR

BLOOM 的 176 B 参数模型在 384 × 80 GB NVIDIA A100 GPU 上训练了 3.5 个月，使用了自定义的 Megatron‑DeepSpeed 堆栈，该堆栈结合了 ZeRO 数据并行、张量并行和流水线并行以及 BF16 混合精度，使得在法国的 Jean Zay 超级计算机上实现了 100 万 GPU‑hour 的运行。

项目概述

  • 硬件：384 块 A100 80 GB GPU（48 个节点，每节点 8 块 GPU），配备 AMD EPYC 7543 CPU，每节点 512 GB RAM，Omni‑Path 互连，专用 NCCL 子网，GPFS 存储。
  • 软件：Megatron‑DeepSpeed（一个合并了 Microsoft DeepSpeed 和 NVIDIA Megatron‑LM 的分支）。
  • 架构：GPT‑3 风格的 transformer，加入了改进（ALiBi 位置编码、嵌入层 LayerNorm）。
  • 数据集：来自 46 种语言的 350 B 令牌（约 1.5 TB 清洗文本），词表大小 250 680。
  • 训练时长：约 3.5 个月（≈100 万 计算小时）。

关键贡献者

此工作依赖于六个主要团队：

  1. Hugging Face BigScience 团队（内部工程师和资金）。
  2. Microsoft DeepSpeed 团队（提供 DeepSpeed 库和集成帮助）。
  3. NVIDIA Megatron‑LM 团队（框架和建议）。
  4. IDRIS/GENCI 员工，负责管理 Jean Zay 超级计算机（提供计算资源）。
  5. PyTorch 核心团队（错误修复和可用性改进）。
  6. BigScience 工程工作组的志愿者。

值得注意的个人包括 Olatunji Ruwase、Deepak Narayanan、Jeff Rasley、Jared Casper、Samyam Rajbhandari 和 Rémi Lacroix。

Megatron‑DeepSpeed 堆栈

组件 由 DeepSpeed 提供 由 Megatron‑LM 提供
ZeRO 数据并行 ✅
张量并行 ✅
流水线并行 ✅
BF16 优化器 ✅
融合 CUDA 核 ✅
DataLoader ✅

该堆栈实现了 3‑D 并行：

  • 数据并行 (DP) 将模型复制到多个 GPU 组，每个 GPU 处理一块数据。
  • 张量并行 (TP) 将单个张量切分到多个 GPU 上，降低每个 GPU 的内存占用。
  • 流水线并行 (PP) 将模型在 GPU 之间垂直切分，以流水线方式处理微批次，避免 GPU 空闲。
  • ZeRO 进一步切分优化器状态、梯度，甚至权重，以适配大模型。

并行细节

ZeRO 数据并行

与完整模型复制不同，每个 GPU 只保存参数、梯度和优化器状态的一部分，并在需要时即时重建完整张量。这大幅降低了内存开销。

张量并行

权重矩阵在 GPU 之间按列切分；每个 GPU 计算矩阵乘法的相应切片并在本地进行激活。此方式需要高速互连；在 BLOOM 中，TP 的度数在每个节点上限制为 4（每个张量切片对应一块 GPU）。

流水线并行

模型层被划分为多个阶段；微批次在这些阶段之间流动，使所有 GPU 保持忙碌。“chunks”（或 GAS）超参数控制微批次数量，在流水线空泡与微批大小之间取得平衡。BLOOM 使用了 72 个流水线阶段（包括两个嵌入阶段），以在 GPU 之间平衡内存。

组合 DP + PP + TP（3‑D 并行）

最终的训练配置使用 DP 进行数据分配，TP 进行张量切分，PP 进行层分布，实现了对 384 GPU 集群的高效利用。

BF16 优化器

在 FP16 下训练导致早期实验出现发散（例如 104 B 模型）。BLOOM 切换到 BF16 混合精度，使用自定义的 BF16Optimizer，其：

  • 保持 FP32 的指数范围，避免溢出。
  • 所有累加均在 FP32 中进行。
  • 在流水线微批次之间以 FP32 累积梯度。 该优化器使 176 B 模型的损失曲线保持稳定。

融合 CUDA 核

为最小化 GPU 空闲时间，使用了 Megatron‑LM 的自定义融合核，用于：

  • LayerNorm
  • 组合缩放、掩码和 softmax
  • 加偏置的 GeLU（通过 PyTorch JIT） 这些核通过将中间结果保存在寄存器中，降低了内存流量。

数据集处理

训练数据流水线：

  • 将 1.5 TB 清洗后的多语言文本标记化为 350 B 令牌。
  • 为固定序列长度 2048 创建每个样本的索引。
  • 按 epoch 打乱顺序，以确保均匀曝光。
  • 将索引保存到磁盘，以避免重启时重新计算。
  • 多个数据集以可配置权重进行混合。

架构微调

  • 嵌入层 LayerNorm：在第一个嵌入后添加 LayerNorm 稳定了训练，灵感来源于 bitsandbytes 中的 StableEmbedding 实现。
  • ALiBi 位置编码：用线性偏置注意力（ALiBi）取代绝对位置嵌入，使得能够外推到比训练长度（2048）更长的序列。

工程挑战

  • 硬件故障：每周 1–2 块 GPU 故障；每 3 小时保存一次检查点，将每次故障的损失工作限制在约 1.5 小时。
  • 软件错误：需要 CUDA_LAUNCH_BLOCKING=1、优化器组拆分以及自定义 SLURM 终止开关以进行多用户作业控制。
  • 停机时间：因死锁、磁盘空间耗尽等问题导致 5–10 小时中断；整体训练仍在计划的 3.5 个月窗口内完成。
  • 值班协调：分布在欧洲和加拿大西海岸，实现了 24/7 监控，无需专用呼叫器。

结论

最密集的阶段是两个月的准备工作，包括 BF16 优化器的后期开发以及大规模并行的调试。一旦堆栈稳定，176 B 模型即可顺利训练，证明在提供足够计算资源和协作工具的情况下，开源团队能够训练最先进的多语言模型。

资源

关键论文

  • Megatron‑LM：Efficient Large‑Scale Language Model Training on GPU Clusters（arXiv:2104.04473）
  • DeepSpeed ZeRO：ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（arXiv:1910.02054）
  • ALiBi：Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation（arXiv:2108.12409）

Sources