DeepSpeedおよびFairScaleを介したHugging Face TransformersのZeRO統合
Hugging Face Transformers v4.2.0では、DeepSpeedおよびFairScaleの実験的なサポートが導入され、Zero Redundancy Optimizer (ZeRO) を統合することで、GPUメモリ使用量を最適化し、より大きなモデルのトレーニングやバッチサイズの拡大を可能にします。この統合により、Trainer APIを通じて、--sharded_ddp (FairScale) および --deepspeed (DeepSpeed) コマンドライン引数を使用して、これらの最適化を直接利用できるようになります。
マルチGPUトレーニングのためのZeROメモリ最適化
DeepSpeedまたはFairScaleを介してZeROを統合することで、標準的なDistributed Data Parallel (DDP) のベースラインと比較して、トレーニングおよび評価時間を大幅に短縮し、同時に可能な最大バッチサイズ (BS) を増加させることができます。2つの24GB Titan RTX GPUを使用してt5-largeモデルを用いたベンチマークでは、以下のパフォーマンス向上が観察されました：
| Method | Max Batch Size | Train Time | Eval Time |
|---|---|---|---|
| Baseline (DDP) | 16 | 30.9458 | 56.3310 |
| fp16 | 20 | 21.4943 | 53.4675 |
| sharded_ddp (FairScale) | 30 | 25.9085 | 47.5589 |
| sharded_ddp + fp16 | 30 | 17.3838 | 45.6593 |
| DeepSpeed (no CPU offload) | 40 | 10.4007 | 34.9289 |
| DeepSpeed (with CPU offload) | 50 | 20.9706 | 32.1409 |
DeepSpeedはバッチサイズとトレーニング速度において最も高い利得を示しましたが、FairScaleは設定ファイルなしで単一のコマンドライン引数のみで済むため、デプロイがより容易であるとされています。
大規模モデルのシングルGPUトレーニング
DeepSpeedは、CPUオフロードを通じて、単一のGPUのメモリ容量を超えるようなモデルのトレーニングを可能にします。単一の24GB RTX-3090カードを使用したt5-3bモデルのテストでは、標準的なシングルGPUセットアップではバッチサイズが1であってもOut of Memory (OOM) エラーが発生しました。DeepSpeedを使用することで、モデルはバッチサイズ20で正常にトレーニングされ、システムはバッチサイズ30で初めてOOMに達しました。
ZeROの技術的メカニズム
ZeRO (Zero Redundancy Optimizer) は、データ並列トレーニングに分散データストレージを追加することでメモリを最適化します。すべてのGPUにモデルの状態全体を複製するのではなく、ZeROはパラメータ、勾配、およびオプティマイザの状態を、利用可能なGPU間で分割（パーティション）します。
分散パーティショニング
各GPUは、パラメータ、勾配、およびオプティマイザの状態の単一のシャードのみを保持します。実行時に、各GPUは特定のレイヤーに必要なデータを他の参加しているGPUからオンザフライで取得するため、データのストレージに重複がなくなります。
ZeRO-Offload
ZeRO-Offloadは、特定の処理とメモリ要件をGPUからホストCPUへ移動させます。これは、大規模なモデル（t5-3bなど）を限られたGPUハードウェアに適合させるために極めて重要です。
メモリ断片化の管理
DeepSpeedは、利用可能な総メモリ量があるにもかかわらず、連続したブロックが十分に大きくないためにOOMエラーが発生する「GPUメモリ断片化」に対処するため、長期的な割り当てと短期的な割り当てを分離してGPUメモリを独立して管理します。
デプロイと統合
ZeROの最適化には、モデルアーキテクチャの変更は必要ありません。変更が必要なのはトレーニングコードのみです。Hugging Face Trainer のユーザーは、以下のフラグを使用してこれらの機能を実装できます：
--sharded_ddp: FairScale統合を有効にします。--deepspeed: DeepSpeed統合を有効にします（JSON設定ファイルが必要です）。
今後の拡張機能として、DeepSpeed Sparse Attention、1-bit Adam、およびFairScaleとDeepSpeedの両方におけるモデルパラメータのシャーディングのサポートが予定されています。