BLOOM 176B 訓練技術概覽

TL;DR

BLOOM 的 176 B 參數模型在 384 × 80 GB NVIDIA A100 GPU 上訓練了 3.5 個月，使用結合 ZeRO 資料並行、張量並行與流水線並行的自訂 Megatron‑DeepSpeed 堆疊，採用 BF16 混合精度，使得在法國 Jean Zay 超級電腦上完成了 1 M GPU 小時的運算。

專案概述

  • 硬體：384 A100 80 GB GPU（48 個節點，每節點 8 GPU），搭配 AMD EPYC 7543 CPU、每節點 512 GB 記憶體、Omni‑Path 互連、專用 NCCL 子網路、GPFS 儲存。
  • 軟體：Megatron‑DeepSpeed（結合 Microsoft DeepSpeed 與 NVIDIA Megatron‑LM 的分支）。
  • 架構：GPT‑3 風格的 transformer，加入了改進（ALiBi 位置編碼、embedding LayerNorm）。
  • 資料集：來自 46 種語言的 350 B 代幣（約 1.5 TB 清理過的文字），詞彙表大小 250 680。
  • 訓練時長：約 3.5 個月（≈1 M 計算小時）。

主要貢獻者

此工作依賴六個主要團隊：

  1. Hugging Face BigScience 團隊（內部工程師與資金）。
  2. Microsoft DeepSpeed 團隊（提供 DeepSpeed 函式庫與整合協助）。
  3. NVIDIA Megatron‑LM 團隊（框架與建議）。
  4. IDRIS/GENCI 工作人員（管理 Jean Zay 超級電腦，提供捐贈運算資源）。
  5. PyTorch 核心團隊（錯誤修正與可用性改進）。
  6. BigScience Engineering 工作小組的志工。

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

Megatron‑DeepSpeed 堆疊

元件 由 DeepSpeed 提供 由 Megatron‑LM 提供
ZeRO Data Parallelism ✅
Tensor Parallelism ✅
Pipeline Parallelism ✅
BF16 Optimizer ✅
Fused CUDA Kernels ✅
DataLoader ✅

此堆疊實作 3‑D 並行：

  • 資料並行 (DP)：在 GPU 群組間複製模型，每個 GPU 處理一段資料。
  • 張量並行 (TP)：將單一張量切分至多個 GPU，降低每個 GPU 的記憶體需求。
  • 流水線並行 (PP)：將模型垂直切分至多個 GPU，於流水線中處理微批次，以避免 GPU 閒置。
  • ZeRO：進一步切分 optimizer 狀態、梯度，並可選擇性切分權重，以容納大型模型。

並行細節

ZeRO 資料並行

與完整模型複製不同，每個 GPU 僅保有參數、梯度與 optimizer 狀態的一部份，根據需要即時重建完整張量。此方式大幅降低記憶體開銷。

張量並行

權重矩陣以列方式切分至多個 GPU；每個 GPU 計算其對應的矩陣乘法切片，並在本地套用激活函式。此需高速互連；在 BLOOM 中，TP 的度數限制為每節點 4（每個張量切片對應一顆 GPU）。

流水線並行

模型層被劃分為多個階段；微批次在各階段間流動，使所有 GPU 均保持忙碌。「chunks」(或稱 GAS) 超參數控制微批次的數量，以在流水線空洞與微批次大小之間取得平衡。BLOOM 使用了 72 個流水線階段（包含兩個 embedding 階段），以在 GPU 之間平衡記憶體使用。

結合 DP + PP + TP（3‑D 並行）

最終的訓練配置使用 DP 進行資料分配、TP 進行張量切分、PP 進行層分配，達成對 384 GPU 叢集的高效利用。

BF16 Optimizer

訓練時使用 FP16 在早期實驗（例如 104 B 模型）中導致發散。BLOOM 轉而使用 BF16 混合精度，透過自訂 BF16Optimizer，其特性包括：

  • 保持 FP32 的指數範圍，避免溢位。
  • 所有累加皆以 FP32 完成。
  • 在流水線微批次間以 FP32 累積梯度。 此 optimizer 使 176 B 模型的損失曲線保持穩定。

融合 CUDA 核心

為了最小化 GPU 閒置時間，使用 Megatron‑LM 的自訂融合核心處理以下工作：

  • LayerNorm
  • 結合 scaling、masking 與 softmax
  • 加上偏置的 GeLU（透過 PyTorch JIT） 這些核心透過將中間結果保留在暫存器中，減少記憶體流量。

資料集處理

訓練資料管線：

  • 將 1.5 TB 清理過的多語言文字斷詞成 350 B 代幣。
  • 為固定長度 2048 的序列建立每筆樣本的索引。
  • 以 epoch 為單位隨機排序，以確保均勻曝光。
  • 將索引寫入磁碟，避免重新啟動時重新計算。
  • 多個資料集以可設定權重混合。

架構調整

  • Embedding LayerNorm：在第一層 embedding 後加入 LayerNorm，使訓練更穩定，靈感來自 bitsandbytes 中的 StableEmbedding 實作。
  • ALiBi 位置編碼：以線性偏差注意力 (ALiBi) 取代絕對位置嵌入，允許序列長度超過訓練長度（2048）時的外推。

工程挑戰

  • 硬體故障：每週 1–2 顆 GPU 故障；每 3 小時儲存一次 checkpoint，使每次故障損失的工作時間限制在約 1.5 小時。
  • 軟體錯誤：需要設定 CUDA_LAUNCH_BLOCKING=1、optimizer 群組切分，以及自訂 SLURM 終止開關以控制多使用者工作。
  • 停機時間：因死結、磁碟空間耗盡等問題造成 5–10 小時中斷；整體訓練仍在預定的 3.5 個月內完成。
  • 待命協調：分散於歐洲與加拿大西岸，實現 24/7 監控，且無需專屬呼叫器。

結論

最耗時的階段是兩個月的準備工作，包括 BF16 optimizer 的後期開發與大規模並行除錯。堆疊穩定後，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