はじめに
近年、大規模言語モデルの学習には多くのGPUリソースが必要とされています。そこで登場したのが、Zero Redundancy Optimizer (Zero)アルゴリズムを実装したDeepSpeedとPyTorch FSDPの2つのフレームワークです。Hugging Face Accelerateはこの両者に対応しており、ユーザーは簡単に切り替えて使用できます。本記事では、AccelerateにおけるFSDPとDeepSpeedの違いや使い分け方について詳しく解説していきます。
FSDPとDeepSpeedは互換性がある？
FSDPとDeepSpeedを比較するため、Mistral-7Bモデルを使った学習を行ってみました。すると、図1のようにDeepSpeed(青線)の損失は収束していますが、FSDP(オレンジ線)は思うように減少しませんでした。
そこで学習率をGPU数の4倍に上げたところ、図2のようにFSDPの損失曲線が改善されました。ただし、別の学習率(1e-5)ではスケーリングなしでも両者の振る舞いが一致しました(図3)。
精度の違いが結果を左右
調査の結果、精度の取り扱いに違いがあることが分かりました。DeepSpeedは内部的にfp32に変換してマスターウェイトを保持する一方、FSDPはtorch_dtypeの設定に従います。表1にまとめたように、FSDPはメモリ制約のあるシナリオでも低精度で動作できる柔軟性を持っています。
AccelerateでFSDPとDeepSpeedを揃える
Accelerate 0.30.0からは、混合精度が有効な場合にFSDPでも自動的にアップキャストを行うようになりました(図4)。これにより、表2のように「混合精度モード」と「低精度モード」の2つの動作が可能になり、DeepSpeedとの互換性が高まりました。
スループット比較
IBM Granite 7Bモデルを使い、4枚のA100 GPUでFSDPとDeepSpeedのスループットを比較しました。表3の通り、トークン/秒/GPUやモデルFLOPS利用率(MFU)はほぼ同等の結果となりました。
概念ガイドとAccelerateの設定
今回の取り組みの成果として、FSDPとDeepSpeed間の移行を助ける概念ガイドを提供しました。シャーディング戦略の同等性やモデルの効率的なロード方法など、ユーザーの疑問に答えています。
Accelerateでは、コマンドラインやプラグインクラスからこれらのフレームワークを簡単に設定できます。設定ファイルを変更するだけでFSDPとDeepSpeedの切り替えがほぼ自動化されているのです。
結論
FSDPとDeepSpeedはどちらもZeroアルゴリズムを実装した強力なフレームワークですが、内部の精度の扱いに違いがあることが分かりました。Hugging Face Accelerateの最新バージョンでは、この違いを吸収してシームレスに切り替えられるようになっています。
参考サイト
コメント