AWSのAI半導体で動くモデルを検索可能に　Hugging Faceが確認ツール公開

Hugging Faceは、アマゾン・ウェブ・サービス（AWS）の機械学習チップ「AWS Trainium」上で自社ハブのAIモデルが動作するかどうかを確認できるツール「torch-neuronx model support Space」を公開した。モデルをコードの修正なしで実行できるかを判定し、自身で動かすための実行スクリプトを提供する。

同社は以前、AWS Trainium向けのPyTorchバックエンド「TorchNeuron」を発表し、TransformersのコードをTrainium上で動作可能にしていた。今回公開されたツールでは、ダウンロード数やトレンドの上位モデルを対象に動作状況を表示する。一覧では検証結果のほか、モデルの分割に用いたコア数、コンパイルを含む実行時間、ピーク時のメモリ使用量が確認できる。利用者はNeuronコンパイラ「neuronx-cc」やTransformersのバージョンを指定して結果を絞り込むことも可能だ。

検証は「Trainium2（trn2）」インスタンス上で実施されている。重みを保持できる最小のコア数に分割してモデルを読み込み、32トークンを生成して正常性を確認する。単一コアの場合は同一チェックポイントのCPU実行結果と生成トークンを比較し、分割配置の場合は小型の判定モデルで文章の流暢さを評価する。現時点で検証対象となっているのは最大4台のTrainium2デバイス（最大16コア）に収まるモデルに限られる。

テストデータは一般公開されているデータセット「aws-neuron/torch-neuronx-model-support」を基にしており、直接照会することもできる。現状のテスト対象はテキスト生成および画像・テキスト入力モデルに限定されているが、今後は画像生成ライブラリのdiffusersやsentence transformers、その他のタスクへと対応を広げる予定だ。