DeepSpeed を使用した大規模モデルのトレーニングを加速する

Hugging Face は DeepSpeed の Zero Redundancy Optimizer（ZeRO）を accelerate ライブラリに統合し、GPU 間のメモリ冗長性を削減することで大規模モデルをより効率的にトレーニングできるようにしました。この統合により、開発者は Out of Memory（OOM）エラーを回避し、コードを大幅に変更することなくバッチサイズを増やすことができます。

ZeRO データ並列処理の理解

ZeRO（Zero Redundancy Optimizer）は、利用可能な GPU にトレーニング状態を分割（シャーディング）することでデータ並列処理中のメモリ使用量を最適化します。このフレームワークは、モデル訓練のメモリフットプリントを削減するために複数のステージで動作します：

• Stage 1: オプティマイザ状態をデータ並列ワーカー／GPU 間でシャーディングします。
• Stage 2: オプティマイザ状態と勾配の両方をワーカー／GPU 間でシャーディングします。
• Stage 3: オプティマイザ状態、勾配、モデルパラメータをワーカー／GPU 間でシャーディングします。
• Optimizer Offload: Stage 2 を拡張し、勾配とオプティマイザ状態を CPU またはディスクにオフロードします。
• Param Offload: Stage 3 を拡張し、モデルパラメータを CPU またはディスクにオフロードします。

Accelerate によるコード不要の統合

ユーザーは accelerate config コマンドと Accelerate DeepSpeed プラグインを使用することで、トレーニングコードを変更せずに DeepSpeed ZeRO Stage-2 を活用できます。

24GB の NVIDIA Titan RTX GPU を 2 台使用して DeBERTa-v2-xlarge-mnli モデル（パラメータ数 9億）をファインチューニングしたベンチマークにおいて、DeepSpeed ZeRO Stage-2 は Distributed Data Parallel（DDP）を大幅に上回る性能を示しました：

手法	最大バッチサイズ	エポックあたりの訓練時間 (秒)	エポックあたりの評価時間 (秒)	F1 スコア	精度
DDP	8	103.57	2.04	0.931	0.904
DeepSpeed ZeRO Stage 2	40	28.98	1.79	0.936	0.912

DeepSpeed は、最大バッチサイズを 5 倍 に増加させ、総訓練時間を 約 3.5 倍 高速化しました（性能指標に劣化はありません）。

DeepSpeed 設定ファイルによる高度な構成

より細かい制御が必要な場合、ユーザーは accelerate config を通じて DeepSpeed 設定 JSON ファイルを提供できます。設定ファイルでオプティマイザとスケジューラが定義されている場合、標準の PyTorch オプティマイザとスケジューラを accelerate.utils.DummyOptim と accelerate.utils.DummyScheduler に置き換えるだけで、最小限のコード調整で済みます。

MuDoConv データセットで BlenderBot-400M-distill モデルを使用してこの手法をテストしたところ、DeepSpeed ZeRO Stage-2 は DDP の 100 に対し最大バッチサイズ 200 を実現しました。その結果、訓練が 1.44 倍、評価が 1.23 倍 高速化しました。

CPU/ディスクオフロードによる超大規模モデルの訓練

CPU オフロードを伴う DeepSpeed ZeRO Stage-3 は、バッチサイズ 1 でも GPU メモリに収まらないモデルの訓練を可能にします。

24GB の NVIDIA Titan RTX GPU を 1 台使用して GPT-XL モデル（パラメータ数 15億）を訓練したテストでは、DDP はすぐに OOM エラーとなりました。これに対し、DeepSpeed ZeRO Stage-3 はオプティマイザ状態、勾配、パラメータを CPU にオフロードすることで バッチサイズ 16 の訓練を成功させ、エポックを 6608.35 秒で完了させました。

手法	最大バッチサイズ	エポックあたりの訓練時間 (秒)	備考
DDP	-	-	OOM エラー
DeepSpeed ZeRO Stage 3	16	6608.35	-