モデルの対応状況を検索できる
Hugging Faceは、AWSのAI向け半導体TrainiumでTransformersのモデルが動作するかを調べられるSpace（ブラウザー上で利用できるデモ環境）を公開しました。背景にあるのは、PyTorchからTrainiumを使うためのバックエンド「TorchNeuron」です。Transformersのモデルは理論上Trainium上で実行できるものの、実際にどのモデルがどの条件で動くかは、モデルごとに試す必要がありました。新しいSpaceでは、人気モデルを中心に動作確認の結果を一覧できます。
画面には、Hugging Face Hubでのダウンロード数やトレンドをもとに選ばれたモデルが表示されます。モデル名の一部を入力して検索できるほか、ランキングをたどることも可能です。モデルごとに、実行結果、使用したTrainiumコア数、処理時間、ピーク時のデバイスメモリー量を確認できます。さらに、NeuronコンパイラーとTransformersのバージョンを切り替えられるため、ソフトウェアの組み合わせによる違いも調べられます。未テストのモデルや、対象から外した理由が表示される場合もあります。
実行スクリプトと検証環境も公開
モデルを選ぶと、検証に使ったNeuron関連ソフトウェアの構成や、追加で必要なパッケージ、実行スクリプトが確認できます。スクリプトはモデルを読み込み、必要に応じて複数のコアに分割して配置し、生成処理を行う構成です。複数コアへの分割には、テンソル並列化などの方式が使われます。利用者は、同じソフトウェア構成のTrainium2環境を用意し、パッケージを入れたうえでpython3 run.pyを実行できます。モデルを分割して動かす場合も、スクリプトが分散実行の準備をするため、別途torchrunを起動する必要はありません。
初回実行では、モデルの計算グラフをTrainium向けにコンパイルします。記事では、検証時と同じコンパイラーが生成したグラフを公開キャッシュから取得し、初回のコンパイルを省く方法も案内しています。これは試行にかかる時間を抑える助けになりますが、キャッシュの利用には対応するコンパイラーのバージョンを合わせる必要があります。
「動く」の意味と確認範囲を理解する
合格の意味を知るには、テスト条件も重要です。記事執筆時点では、Transformersのテキスト生成と画像・テキスト入力モデルが対象で、モデルの重みがTrainium2の4台以内、最大16コアに収まるものを検証しています。モデルは重みを保持できる最小限のコア数に分割され、32トークンを生成して確認します。単一コアの場合はCPUでの生成結果と照合し、複数コアの場合は小型の判定モデルで英文として流暢かを確認します。
したがって、PASSは指定条件で一定の生成確認を通過したという意味であり、処理速度や実運用時のスループット、長文生成の品質を保証するものではありません。性能を判断するには、実際の入力長や生成量、同時利用数に合わせた別途の測定が必要です。対象も今後拡大する予定とされていますが、現時点で未掲載のモデルが必ず非対応とは限りません。
結果はデータセットとしても利用可能
Spaceに表示される情報は、公開データセット「aws-neuron/torch-neuronx-model-support」に基づきます。モデルごとの結果に加えて、エラー、ソフトウェア一式のバージョン、実行スクリプトも含まれ、Pythonのデータ分析ライブラリーpandasから読み込んで絞り込めます。画面での検索だけでなく、複数モデルの結果をまとめて比較したい開発者にも使える形です。今日試すなら、まず候補モデルをSpaceで検索し、結果と使われたバージョンを確認してから、手元のTrainium2環境が同じ構成か照合しましょう。
参考リンク
Hugging Face：Does Your Model Run on AWS Trainium? Look It Up
本記事はAIが生成し、自動公開しています。内容の正確性には努めていますが、重要な判断の際は必ず一次情報をご確認ください。