Hugging Face の Optimum Habana ライブラリーを使用して、さらなる大規模言語モデルの最適化を進めています。

Hugging Face と Habana Labs では最近、GPT-2、GPTNeoX、GPT-J、OPT、LLAMA などのモデルを対象に、Optimum-Habana ライブラリーを使用したインテル® Gaudi® アクセラレーター / インテル® Gaudi® 2 アクセラレーターで実行する複数の LLM 推論リファレンス・モデルを公開しました。

これらのモデルもほかの LLM と同様に、Hugging Face がホストするレポジトリーを使用して開発者が簡単に実行できるようになっています。Habana ハードウェアにアクセスするには、インテル® Tiber™ デベロッパー・クラウドからインテル® Gaudi® 2 アクセラレーター実装インスタンス、または初代のインテル® Gaudi® アクセラレーターを実装した Amazon EC2 DL-1 インスタンスを使用できます。このガイドでは、まず Habana チームがインテル® Gaudi® アクセラレーター上で GPTNeoX の実行を可能にしている手法を簡単に説明してから、単一のインテル® Gaudi® アクセラレーター・デバイスで GPTNeoX モデルを使用した場合と、8 デバイスで DeepSpeed 軽量化フレームワークを使用した場合それぞれの、テキスト生成の実行方法を示しています。

いずれのモデルでも、Python コマンドの –model_name_or_path にモデル名を入れるだけで実行可能です。例えば GPT-J の場合は、–model_name_or_path EleutherAI/gpt-gpt-j-20b と設定します。
ここでは、これらの新しいモデルのパフォーマンスを向上させる最適化の実装方法を 2 タイプ紹介します。既存モデルにはすでに以下の修正が適用され、Optimum Habana ライブラリーの一部として使用する準備ができています。それ以外の LLM の場合は、推論パフォーマンスを向上させるために、最適化の適用に必要な変更を行うことを推奨します。

まず 1 つ目では、新トークン生成中の入力形状を静的に固定することで、不要なグラフ再コンパイルを回避しました。これは、以下のコードに示すとおり、生成関数を呼び出すたびに、自己アテンション・マスクの入力トークンと生成トークンすべてのベクトルを最大トークン長までパディングすることで実現します。また、出力シーケンスで生成中のトークンを追跡するために、新しい変数 ‘token_idx’ を取り入れました。この変数はパディング中に、パディングされたコンテンツ群の最新トークンを追跡するために使用します。静的形状で最適化された現行モデル一覧を参照してください。
https://github.com/huggingface/optimum-habana/blob/11bd194bff283bcd8375a84733367cccc252f52b/optimum/habana/transformers/generation/utils.py%23#L60

2 つ目の最適化には静的キー・バリュー・キャッシュを使用し、新トークンの生成に必要な自己アテンション前方パスの再コンパイルを不要にしました。実装については以下のコマンドを参照してください。
https://github.com/huggingface/optimum-habana/blob/7e8c855434f3ea7a85325facde3751909aa5255b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py%23#L62

これらの最適化を適用したモデル例には、オリジナルの上流モデルコードを継承した、専用のインテル® Gaudi® アクセラレーター・サブクラスが含まれています。同じプロセスを経て、前述の表の各クラスに因果言語モデリング (CLM) のサブクラスを作成しました。唯一の違いは、先ほど説明した静的形状とキー・バリュー・キャッシュの最適化に関連する 2 つの変更だけです。
https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
https://github.com/huggingface/optimum-habana/blob/11bd194bff283bcd8375a84733367cccc252f52b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py%23#L282

Optimum-Habana を使用した HPU_graphs と DeepSpeed 推論はすでにサポートしており、どちらもモデルのパフォーマンスを最適化する付加的な手法と言えます。これらは次のセクションで紹介するコマンドラインのオプションによって有効にすることが可能です。

学習処理または推論用に HPU_graphs を設定できているか確認してください。
https://docs.habana.ai/en/latest/PyTorch/Model_Optimization_PyTorch/HPU_Graphs_Training.html
https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Inference_Using_HPU_Graphs.html

インテル® Gaudi® アクセラレーター搭載プラットフォームでほかの LLM を使用し推論を実行したい開発者は、この 2 つの最適化手法を実装して、推論パフォーマンスの向上につなげることができます。

【単一インテル® Gaudi® アクセラレーター・デバイスでの推論のセットアップと実行】

以下のコマンドを実行して、Optimum-Habana のクローンを作成し、必要な依存関係をインストールすることができます。この例では、Synapse 1.10.0 と Hugging Face v1.6.1 を使用しました。

次に、以下のコマンドを実行して、200億パラメーターの GPTNeoX を使用したテキスト生成の出力を得ます。プロンプトは自由に変更してください。前述の最適化を実装する引数 –use_kv_cache を含める点に注意が必要です。
https://huggingface.co/EleutherAI/gpt-neox-20b

このプロンプトによって、以下のようなインテル® Gaudi® 2 アクセラレーターを使用した出力を確認することができます。

【DeepSpeed を使用した複数インテル® Gaudi® アクセラレーター・デバイス上での推論の実行】

ここからは、インテル® Gaudi® 2 アクセラレーター 8 デバイスで DeepSpeed を実行します。この場合も単一デバイスと同じ引数を使用します。gaudi_spawn.py は、mpirun を呼び出し、multi-card の実行を開始するスクリプトです。

出力は以下のようになります。

＜次のステップ＞
Hugging Face と Habana Labs では引き続きさまざまなリファレンス・モデルを有効化して、誰でもアクセスできるように Optimum-Habana と Model-Referencesで公開していく予定です。ユーザーの皆さんには、インテル® Gaudi® アクセラレーターが持つパフォーマンスの優位性を活用できるように、これらのモデルも独自の新しい LLM どちらも同様に試してみることをお勧めします。インテルの開発者向けサイトでは、導入と運用に役立つ記事やフォーラムの投稿が掲載されています。ぜひご活用ください。

https://github.com/huggingface/optimum-habana
https://github.com/HabanaAI/Model-References/tree/master
https://developer.habana.ai/