インテル® Gaudi® アクセラレーターを活用した FSDP による分散型トレーニング
PyTorch とインテル® Gaudi® アクセラレーターを活用した完全シャード化データ並列処理 (FSDP) のトレーニングによる、スケーラブルなモデル開発方法
どこでも利用できてアクセスしやすい AI コンピューティングの追求は、分野を問わず開発者にとって根強く続く課題です。インテル® Gaudi® アクセラレーターは、この状況を変え、NVIDIA の A100 や H100 プロセッサーに代わるコストとパフォーマンスのバランスの取れた高性能 GPU の選択肢となっています。ハードウェア性能だけでなく、ソフトウェアのサポートが充実している点も特長の 1 つです。
抽象化の面では PyTorch、Hugging Face、DeepSpeed といった主要なフレームワークをサポートし、それぞれに最適化されています。特にモデルの学習処理については、低ランク適応 (LoRA) によるパラメーター効率ファインチューニング (PEFT) や ZeRO 3D 並列処理、そして PyTorch での完全シャード化データ並列処理 (FSDP) といった手法に対応しました。これらの統合と最適化は、インテル® Gaudi® アクセラレーターによるオープンソースの開発者エコシステムに対する積極的な取り組みを示すものです。
最新の FSDP への対応によって、生成 AI モデルを構築する開発者はさらに強力なツールを手に入れたことになります。
＜FSDP とは？＞
完全シャード化データ並列処理 (FSDP) は、現在インテル® Gaudi® AI アクセラレーターでサポートされている PyTorch のデータ並列トレーニング手法の 1 タイプであり、大規模モデルの分散トレーニングを可能にします。
分散データ並列処理 (DDP) と呼ばれる PyTorch の従来の並列トレーニング手法では、モデルの重み付けとオプティマイザーの状態が全ワーカーに複製されるため、学習処理中のメモリー・フットプリントが増大しかねません。FSDP では、モデルのパラメーター、オプティマイザーの状態、勾配がすべてのランクに分散されるため、インテル® Gaudi® アクセラレーターのメモリー・フットプリントが抑えられ、各ワーカーはモデルの完全コピーを保持する必要がなくなります。
図 1. FSDP と DDP の違いを示すダイアグラム。詳しくは、FSDP Theory of Operations を参照してください。- 執筆者提供画像
インテル® Gaudi® アクセラレーターでは FSDP を統合し、eager モードと torch.compile を使用して、大規模モデルの学習処理にシームレスかつ効率的なアプローチを適用できます。これは通常用いられている DeepSpeed 分散トレーニング機能に代わる手法と考えられています。
＜FSDP によるマルチカード・ファインチューニング＞
FSDP を有効にするには、インテル® Gaudi® アクセラレーターのほかのワークフローと同様に、Optimum Habana ライブラリーを使用し FSDP と LoRA で Llama2–70B のファインチューニングを行います。
Optimum Habana は、 Transformers や Diffusers とインテル® Gaudi® AI アクセラレーターの間をつなぐインターフェイスです。このツールセットによって、モデルの読み込み、学習処理、推論がシンプルになり、シングル構成とマルチカード構成のどちらもサポートしているため、多様なダウンストリーム・タスクへの対応が可能になります。
＜環境設定＞
開始する前に、環境を設定する必要があります。
1.まず、インテル® Tiber™ デベロッパー・クラウドからインテル® Gaudi® アクセラレーターのインスタンスをプロビジョニングするところから始めます。詳細については、開発者向けウェブページの Getting Access を参照してください。
2.次に、インテル® Gaudi® アクセラレーターの PyTorch 用 Docker イメージを使用するために、こちらの手順に従う必要があります。コンテナを使用せずにインテル® Gaudi® アクセラレーターを動かすことも可能ですが、最も簡単明瞭な方法として、コンテナの使用を推奨します。
3.コンテナ内から Optimum Habana ライブラリーをインストールし、Optimum Habana GitHub リポジトリーのクローンを作成します。
• pip install optimum-habana==1.11.0
• cd ~
• git clone -b v1.11.0 https://github.com/huggingface/optimum-habana.git
4.optimum-habana/examples/language-modeling フォルダーに移動し、pip install -r requirements.txt を実行して関連する依存関係をインストールします。
インテル® Gaudi® アクセラレーターで FSDP を実装する例
これで環境設定が完了しました。以下のインテル® Gaudi® アクセラレーターでの FSDP 実装の簡略版を使用して、インテル® Gaudi® アクセラレーター関連のコードを詳しく見ていきましょう。
インテル® Gaudi® アクセラレーターに関連するコードの変更点は以下のとおりです。
• 10 行目 — import habana_frameworks.torch.distributed.hccl のインポート
• 12 行目 — eager モードの有効化: os.environment["PT_HPU_LAZY_MODE"] = "0"
• 13 行目 — インテル® Gaudi® アクセラレーター・デバイスを指定 device_hpu = torch.device('hpu')
• 33 行目 — インテル® Gaudi® アクセラレーター・デバイスをモデル実行に指定 model = ToyModel().to(device_hpu)
• 36 行目 — インテル® Gaudi® アクセラレーターでの FSDP 実行を指定 model = FSDP(model,device_id = device+hpu)
ご覧のように、コード変更は最小限に抑えられています。36 行目が、FSDP を有効化している部分です。この設定により、複数のインテル® Gaudi® アクセラレーター・カードを横断してモデルのパラメーターが分散されるため、各インテル® Gaudi® アクセラレーターがパラメーター全体の一部のみを管理するだけで済むようになります。すでに説明したとおり、これはカード当たりのメモリー要件を大幅に軽減して、インテル® Gaudi® アクセラレーターの搭載数を増やして大規模モデルへの拡張を可能にするためには不可欠です。なお、8 枚のインテル® Gaudi® アクセラレーター・カードをすべて使用したい場合は、—ws 8 (world size) フラグを指定する必要がありますのでご注意ください。
この例を実行するには、上記のコードを fsdp_example.py ファイルにコピーし、python3 fsdp_example.py コマンドを実行します。
＜FSDP と LoRA を使用した Llama2–70B のマルチカード・ファインチューニング＞
このシンプルな実装をマスターしたら、次は Optimum Habana リポジトリーから、さらに包括的な例へ進んでみましょう。リポジトリーで紹介されている例の多くに、便利なスクリプト、ユーティリティー、構成ファイルがあらかじめパッケージされています。通常は、いくつかのパラメーターを調整して、コマンドラインの Python 引数を実行するだけで完了です。
Optimum Habana の例の中にある language-modeling フォルダーに移動し、必要な依存関係をインストールして、以下のコマンドライン・スクリプトを実行します。この例では、DeepSpeed ではなく mpi を使用して FSDP と LoRA による Llama2–70B のマルチカード・ファインチューニングを実行していることがお分かりいただけるでしょう。
• cd ~/optimum-habana/examples/language-modeling
• pip install -r requirements.txt
FSDP に関連する引数の変更点は以下のとおりです。
• --fsdp auto_wrap
FSDP ポリシー「auto_wrap」がメモリーと演算効率の高いメソッドを提供し、FSDP のファインチューニング中にすべてのモデル・パラメーターが同時に読み込まれないように防止。これにより、「FSDP ユニット」と呼ばれる複数のレイヤー集合が形成され、互いに独立してパラメーターを読み込みます。
• --fsdp_config fsdp_config.json
FSDP 用の設定パラメーターを提供。全サポート機能のリストについては、FSDP ocumentation を参照してください。ここでは最も重要な機能をいくつか紹介します。
– “BACKWARD_PRE”: 現在のユニットの演算が始まる前に、FSDP ユニットのパラメーターをリクエスト。
– “TRANSFORMER_BASED_WRAP”: ユーザー定義のラッピングポリシーを許可。
– “1”: 通常「FULL_SHARD」と呼ばれるデフォルトの FSDP ポリシー。
– “GaudiLlamaDecoderLayer”: ユーザー定義の FSDP レイヤーをラップ。
＜まとめとディスカッション＞
このページでは、インテル® Gaudi® アクセラレーターに完全シャード化データ並列処理 (FSDP) を実装し、大規模モデルの学習処理の効率と拡張性を向上させる方法について説明しました。環境設定から詳細なコードの例まで実際の手順をまとめ、インテルのハードウェアを活用する方法とあわせて、最先端のモデル・トレーニング手法にも注目しています。パラメーターのシャード化や効率的な分散トレーニングを有効にするといった手法により、メモリー・フットプリントを抑えることで、開発者は大規模な生成 AI モデルの開発に一層効率よく取りかかることができるようになっています。
★次のステップ
•インテル® Gaudi® アクセラレーターを使用した FSDP の用途を拡大:FSDP を比較的複雑な AI モデルやトレーニングのシナリオに統合します。
•ファインチューニングの手法を深堀り: ハイパーパラメーターを変更したり、LoRA など異なるファインチューニングの手法を使用して、パフォーマンスの出力結果を比較します。
Discussion