BLOOM 176B トレーニング技術概要

TL;DR

BLOOM の 176 B パラメータモデルは、384 台の 80 GB NVIDIA A100 GPU を使用し、3.5 か月間にわたってカスタム Megatron‑DeepSpeed スタックでトレーニングされました。このスタックは ZeRO データ並列、テンソル並列、パイプライン並列を BF16 混合精度と組み合わせ、フランスの Jean Zay スーパーコンピュータ上で 1 M GPU 時間の実行を可能にしました。

プロジェクト概要

  • Hardware: 384 A100 80 GB GPU（48 ノード、各ノード 8 GPU）に AMD EPYC 7543 CPU、ノードあたり 512 GB RAM、Omni‑Path インターコネクト、専用 NCCL サブネット、GPFS ストレージを使用。
  • Software: Megatron‑DeepSpeed（Microsoft DeepSpeed と NVIDIA Megatron‑LM を統合したフォーク）。
  • Architecture: GPT‑3 スタイルのトランスフォーマーに、ALiBi 位置エンコーディング、埋め込み LayerNorm などの改良を加えたもの。
  • Dataset: 46 言語からの 350 B トークン（約 1.5 TB のクリーンテキスト）、語彙サイズ 250 680。
  • Training duration: 約 3.5 か月（≈1 M 計算時間）。

主な貢献者

The effort relied on six main groups:

  1. Hugging Face BigScience チーム（社内エンジニアと資金提供）。
  2. Microsoft DeepSpeed チーム（DeepSpeed ライブラリと統合支援を提供）。
  3. NVIDIA Megatron‑LM チーム（フレームワークと助言）。
  4. Jean Zay スーパーコンピュータを管理する IDRIS/GENCI スタッフ（計算資源を提供）。
  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 ✅

The stack implements 3‑D parallelism:

  • Data Parallelism (DP) は、モデルを GPU グループに複製し、各 GPU がデータの一部を処理します。
  • Tensor Parallelism (TP) は、個々のテンソルを GPU 間で分割し、GPU あたりのメモリ使用量を削減します。
  • Pipeline Parallelism (PP) は、モデルを垂直に分割し、GPU 間でマイクロバッチをパイプライン処理して GPU のアイドル時間を減らします。
  • ZeRO は、最適化子の状態、勾配、必要に応じて重みもさらに分割し、大規模モデルに対応します。

並列化の詳細

ZeRO データ並列

完全なモデル複製を行う代わりに、各 GPU はパラメータ、勾配、最適化子状態の一部だけを保持し、必要に応じてフルテンソルをオンザフライで再構築します。これによりメモリオーバーヘッドが大幅に削減されます。

テンソル並列

重み行列は列方向に GPU 間で分割され、各 GPU は行列乗算の自分のスライスを計算し、ローカルで活性化を適用します。高速インターコネクトが必要で、BLOOM では TP の度合いをノードあたり 4（テンソルスライスごとに 1 GPU）に制限しました。

パイプライン並列

モデル層はステージに分割され、マイクロバッチがステージ間を流れることで全 GPU が稼働します。"chunks"（または GAS）ハイパーパラメータはマイクロバッチ数を制御し、パイプラインのバブルとマイクロバッチサイズのバランスを取ります。BLOOM は GPU 間のメモリバランスを取るために 72 のパイプラインステージ（埋め込みステージを 2 含む）を使用しました。

組み合わせた DP + PP + TP（3‑D 並列）

最終的なトレーニング構成では、データ分配に DP、テンソル分割に TP、層分配に PP を使用し、384 GPU クラスタの効率的な活用を実現しました。

BF16 最適化子

FP16 でのトレーニングは以前の実験（例：104 B モデル）で発散を引き起こしました。BLOOM はカスタム BF16Optimizer を用いた BF16 混合精度に切り替え、以下を実現しました：

  • FP32 の指数範囲を保持し、オーバーフローを防止。
  • すべての加算を FP32 で実行。
  • パイプラインのマイクロバッチ間で勾配を FP32 で蓄積。 この最適化子により、176 B モデルの損失曲線が安定しました。

結合 CUDA カーネル

GPU のアイドル時間を最小化するため、Megatron‑LM のカスタム結合カーネルが以下に使用されました：

  • LayerNorm
  • スケーリング、マスク、ソフトマックスの結合
  • バイアス付き GeLU（PyTorch JIT 経由） これらのカーネルは中間結果をレジスタに保持することでメモリトラフィックを削減します。

データセット処理

トレーニングデータパイプラインは以下の通りです：

  • クリーンな多言語テキスト 1.5 TB を 350 B トークンにトークナイズ。
  • 固定シーケンス長 2048 のサンプルごとのインデックスを作成。
  • エポック単位で順序をシャッフルし、均一な露出を確保。
  • 再起動時の再計算を防ぐためにインデックスをディスクに保存。 複数のデータセットを構成可能な重みでブレンドしました。

アーキテクチャの調整

  • Embedding LayerNorm: 最初の埋め込みの後に LayerNorm を追加することでトレーニングが安定し、bitsandbytes の StableEmbedding 実装に触発されました。
  • ALiBi Positional Encoding: 絶対位置埋め込みを線形バイアス注意（ALiBi）に置き換え、トレーニング長（2048）より長いシーケンスへの外挿が可能になりました。

エンジニアリング上の課題

  • Hardware failures: 週に 1〜2 台の GPU 故障が発生。チェックポイントを 3 時間ごとに保存することで、故障ごとの作業損失は約 1.5 時間に抑えられました。
  • Software bugs: CUDA_LAUNCH_BLOCKING=1 の設定、最適化子グループ分割、マルチユーザジョブ制御のためのカスタム SLURM キルスイッチが必要でした。
  • Downtime: デッドロック、ディスク容量枯渇、その他の問題により 5〜10 時間の中断が発生しましたが、全体のトレーニングは計画された 3.5 か月の期間内に収まりました。
  • On‑call coordination: ヨーロッパとカナダ西海岸に分散して対応し、専用ページャーなしで 24 時間体制の監視を実現しました。

結論

最も集中的だったフェーズは、BF16 最適化子の後期開発や大規模並列化のデバッグを含む 2 か月間の準備期間です。スタックが安定すると、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