PyTorch FSDP を使用した Llama 2 70B のファインチューニング

Hugging Face は、PyTorch Fully Sharded Data Parallelism (FSDP) を使用して Llama 2 70B モデルのファインチューニングの方法論を詳細に説明しており、Hugging Face Transformers、Accelerate、TRL ライブラリを活用しています。このアプローチにより、オプティマイザーの状態、勾配、パラメータをデバイス間でシャードすることで、マルチノード、マルチ GPU セットアップでの巨大モデルのトレーニングが可能になります。

モデルロード中の CPU RAM ボトルネックの克服

Llama 2 70B モデルのロードには通常、相当量の CPU RAM が必要です；ノード上のすべてのプロセスがモデルをロードすると、約 2TB の CPU RAM が必要になる可能性があります（70B パラメータ * 4 バイト * 8 GPU）。Out-of-Memory (OOM) エラーを防ぐため、Hugging Face は transformers と accelerate で実装された特定の初期化戦略を利用します：

  1. Meta Device Initialization: モデルは meta デバイスを使用してすべてのランクで作成され、つまり重みなしで初期化されます。
  2. Rank 0 Loading: ステート ディクショナリはランク 0 でのみロードされます。
  3. Empty Parameter Allocation: 他のすべてのランクは、torch.empty() を使用して meta デバイス上に空のパラメータを作成します。
  4. State Broadcasting: sync_module_states=True を設定することで、FSDP はトレーニング開始前にランク 0 から他のすべてのランクに重みをブロードキャストします。

この方法により、ノードごとに 1 つのプロセスのみが事前学習済みモデルを CPU RAM にロードするため、セットアップ フェーズでのメモリ フットプリントが大幅に削減されます。

シャードされたステート ディクショナリを使用した効率的なチェックポイント

ランク 0 で CPU オフロードを使用して FULL_STATE_DICT でフルの中間チェックポイントを保存すると、NCCL タイムアウト エラーと大きな遅延が発生することがよくあります。これを解決するため、Hugging Face は FSDP の設定で SHARDED_STATE_DICT の使用を推奨しています。

  • Intermediate Checkpoints: SHARDED_STATE_DICT は GPU ごとにシャードを別々に保存し、トレーニングの保存と再開を高速化します。
  • Final Model Export: デプロイ用の標準モデル状態ディクショナリを取得するには、トレーニングの終わりに trainer.save_model() を呼び出す直前にのみ、状態ディクショナリのタイプを FULL_STATE_DICT に切り替えます。

VRAM とトレーニング速度の最適化

計算コストを削減し、トレーニング速度を向上させるため、実装では 2 つの主要な技術が使用されます：Gradient Checkpointing と Flash Attention。

Flash Attention

標準的な注意機構は、要素ごとの操作（マスク、ソフトマックス、ドロップアウト）中の冗長な High Bandwidth Memory (HBM) の読み書きにより、しばしばメモリ束縛になります。Flash Attention は次の方法でこれを最適化します：

  • Kernel Fusion: 中間ステップを SRAM に保持し、最終結果を HBM に一度だけ書き戻します。
  • Tiling: NxN ソフトマックス/スコアの計算は、SRAM の制限内に収まるようにブロックに分割され、オンライン ソフトマックス アルゴリズムが使用されます。
  • Recomputation: バックワードパスでは、フォワードパスからの全 NxN マトリックスを保存する代わりに必要な値を再計算し、メモリ消費を大幅に削減します。

Gradient Checkpointing

Gradient checkpointing を有効にすることで、VRAM 使用量をさらに削減し、70B パラメータ モデルのファインチューニング中にバッチ サイズを大きくしたり、シーケンス長を長くしたりすることが可能になります。

実装とハードウェア仕様

ハードウェア構成

ファインチューニングは以下のハードウェアを使用して実行されました：

  • Nodes: 2 ノード（最小 1 必要）。
  • GPUs: ノードごとに 8 個の A100 (80GB) GPU。
  • インターコネクト: NVLink (intra-node) および Elastic Fabric Adapter (inter-node)。
  • システム RAM: ノードごとに 1TB。
  • CPU: ノードごとに 96 コア。

トレーニング実行

トレーニング プロセスでは、FULL_SHARD ストラテジーと TRANSFORMER_BASED_WRAP 自動ラップ ポリシーを使用した accelerate launch コマンドが利用されました。混合精度トレーニングは bf16 を使用して有効化されました。8 個の A100 80GB GPU を使用したシングルノード セットアップでは、メモリ管理のために bitsandbytes の paged_adamw_32bit オプティマイザーの使用が推奨されます。

ファインチューニングは、meta-llama/Llama-2-70b-chat-hf モデルと smangrul/code-chat-assistant-v1 データセットを使用して約 13.5 時間で完了しました。

Sources