Hugging Face Accelerate: DeepSpeed と FSDP の精度を調和させる
Hugging Face Accelerate は、PyTorch Fully Sharded Data Parallel (FSDP) と Microsoft DeepSpeed の精度処理を合わせるためのアップデートを導入しました。この変更は Accelerate 0.30.0 リリースに統合されており、ユーザーは 2 つの ZeRO Redundancy Optimizer 実装間を、以前は内部の精度デフォルトの違いに起因する収束の不一致なく切り替えることができます。
精度の不一致と収束
DeepSpeed と FSDP がパラメータの精度を扱う方法の違いは、トレーニング結果の乖離を招く可能性があります。bfloat16 でロードした Mistral-7B ベースモデルを使用したテストでは、DeepSpeed は安定した収束を示したのに対し、FSDP は学習率を GPU の数で手動でスケールするか、1e-5 に下げない限り損失を減少させませんでした。
この挙動は内部のアップキャストによるものです。DeepSpeed の DeepSpeedZeroOptimizer_Stage3 は、_create_fp32_partitions を通じて学習可能なパラメータグループを自動的に float32 にアップキャストし、設計上マスタ重みをフル精度で保持します。これにより、低精度では不安定になる学習率でもオプティマイザが収束できるようになります。
精度ワークフローの比較
DeepSpeed と FSDP は、フラット化されたパラメータとオプティマイザの初期化を管理する方法が根本的に異なります。
| プロセス | FSDP | DeepSpeed |
|---|---|---|
| 準備 | torch_dtype を利用 |
torch_dtype を無視し、float32 で作成 |
| オプティマイザ初期化 | torch_dtype でパラメータを作成 |
float32 でパラメータを作成 |
| オプティマイザ（事前ステップ） | torch_dtype へのアップキャスト（ある場合） |
すべてを float32 にアップキャスト |
| オプティマイザ（実際のステップ） | torch_dtype で実行 |
float32 で実行 |
DeepSpeed のネイティブなアップキャストは収束を保証しますが、メモリ消費が 2 倍になることがあり、GPU が少ない環境でのトレーニングでは重要です。対照的に、torch ネイティブの FSDP 実装はアップキャストを強制せず、オプティマイザが低精度で動作できるため、メモリ制約のあるシナリオで柔軟性が高まります。
Accelerate 0.30.0 の新しい FSDP モード
これらのフレームワークを調和させるため、Hugging Face Accelerate は、収束の安定性とメモリ効率の優先度に合わせて、2 つの異なる FSDP モードをサポートするようになりました。
- Mixed-Precision Mode: 準備段階とオプティマイザステップでパラメータを
fp32にアップキャストし、トレーニングはbf16のままにすることで DeepSpeed と合わせます。 - Memory-Constrained Mode: 準備、トレーニング、オプティマイザステップすべてで低精度 (
bf16) のみで動作します。
フレームワーク比較表
| フレームワーク | モデルロード | 混合精度 | 準備 | トレーニング | オプティマイザ |
|---|---|---|---|---|---|
| FSDP (Memory-Constrained) | bf16 |
None | bf16 |
bf16 |
bf16 |
| FSDP (Mixed Precision) | bf16 |
bf16 |
fp32 |
bf16 |
fp32 |
| DeepSpeed | bf16 |
bf16 |
fp32 |
bf16 |
fp32 |
スループット性能
4 台の A100 GPU 上で IBM Granite 7B モデルをベンチマークした結果、FSDP（整合モード）と DeepSpeed（Zero3）はほぼ同等のスループット性能を示しました。
| フレームワーク | トークン / 秒 / デバイス | ステップ時間 (s) | モデル FLOPs 利用率 (MFU) |
|---|---|---|---|
| FSDP (Aligned) | 3158.7 | 10.4 | 0.41 |
| DeepSpeed | 3094.5 | 10.6 | 0.40 |
移行と設定
Hugging Face は、FSDP と DeepSpeed の間の移行を支援するための concept guide を公開しました。切り替えは主に Accelerate の設定ファイルまたは DeepSpeed と FSDP プラグインクラスを通じて行われます。移行時の主な考慮点は以下の通りです。
- Sharding Strategies: フレームワーク間で同等のシャーディングを実現すること。
- Model Loading: 効率的なロードパターンの実装。
- Weight Prefetching: 重みを GPU 間でどのように転送するかの管理。
- Checkpointing: 各フレームワークがモデル状態を保存・ロードする方法の違いへの対応。