AWS Trainium モデルサポート: Hugging Face が torch-neuronx 互換性ルックアップを公開

TL;DR — Hugging Face と AWS は torch-neuronx モデルサポート Space を公開しました。これは、Hugging Face Hub 上の特定の Transformers モデルが AWS Trainium 上で変更なしで動作するかどうかを示し、テストに使用された正確なスクリプトとスタックを提供する公開ルックアップです。これにより、「このモデルは Trainium で動作するか？」という試行錯誤の質問が、検索可能なデータセットに変わります。

この Space は、以前の Hugging Face の投稿 で紹介された torch-neuronx パッケージに含まれる AWS Trainium 用のネイティブ PyTorch バックエンドである TorchNeuron に基づいています。TorchNeuron を使用すると、Transformers は他のデバイスと同様に Trainium 上で独自のモデリングコードを実行するため、原則として Hub 上の任意の Transformers モデルが実行可能です。新しい Space は、Hub で最も人気のあるモデルについて、実際にどれが動作するかを回答します。

発表内容

aws-neuron/torch-neuronx-model-support Space は、AWS Trainium 用の公開された閲覧可能な互換性テーブルです。リストされた各モデルについて、以下を報告します:

  • モデルが Trainium 上で変更なしで動作するかどうか (PASS 結果)
  • モデルがシャーディングされた Trainium コアの数
  • コンパイルを含む実行の壁時間
  • ピークデバイスメモリ使用量

各モデルについて、Space は実行されたスクリプトも公開しているため、結果は信頼するだけでなく再現可能です。

モデルの検索方法

Space は Selected ビューで開きます。これは、著者自身が Hub 上のトップモデルからダウンロード数とトレンドスコアに基づいて選択したモデルのセットとして説明されています。

Trending および Downloads ビューは、Hub 自身のランキングに従い、タスクごとに表示されます。これらのビューには、テストされたかどうかに関係なくモデルがリストされるため、エントリがない場合の曖昧さはありません。モデルが意図的に除外された場合、Space はその理由を示します。

現時点では、Transformers モデルのみがテストされており、選択は text-generation および image-text-to-text モデルに限定されています。今後、diffusers、sentence transformers、およびその他の transformers タスクに拡大する予定です。

Space では、確認したい Neuron コンパイラ (neuronx-cc) と Transformers のバージョンを選択することもできるため、結果を単一の日付なしスナップショットとして読むのではなく、スタックバージョンでフィルタリングできます。

特定のモデルを見つけるには、Trending または Downloads ビューに切り替えて、名前の一部を入力します。Space はモデルが除外された理由を報告します。最も多いのは優先順位が低かったためで、時には Trainium と本当に互換性がないためです。

自分でモデルを実行する方法

Space でモデルを選択すると、実行方法が表示されます。Space には、モデルがテストされた Neuron スタック、そのスタックに追加でインストールするパッケージ、および実際に実行された run.py スクリプトがリストされます。

スクリプトは TorchNeuron の発表のレシピに従います。from_pretrained でモデルをロードし、必要な場合はシャーディングし、静的キャッシュと torch.compile(backend="neuron") で生成します。

そのスタックを持つ Trainium2 (trn2) インスタンス上で、例えば AWS Neuron Deep Learning Container 内で、リストされたパッケージをインストールしてから実行します:

python3 run.py

torchrun の呼び出しは不要です。モデルがシャーディングされている場合、スクリプトは独自の torch.distributed ランクを起動します。

NEFF キャッシュでコンパイルをスキップ

最初の実行はモデルのグラフのコンパイルから始まり、報告された壁時間の大部分はここで費やされます。このステップをスキップするには、テスト中に同じコンパイラが構築したグラフを公開 NEFF キャッシュ バケットからコピーできます。Space は正しいコンパイラバージョンが入力されたコマンドを提供します:

hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

キャッシュを使用すると、元のテストと同じ neuronx-cc バージョンを使用する限り、再コンパイルをスキップして再現できます。

「動作する」の意味

互換性の主張は狭く具体的であり、Space はそれを正確に文書化しています。

各テストは、モデルの重みを保持できる最小数の Trainium コアにシャーディングして、trn2 インスタンス上でモデルをロードします。ほとんどのモデルはテンソル並列処理でシャーディングされます。テンソル並列プランがない mixture-of-experts モデルは、エキスパート並列処理のみを使用します。テストは 32 トークンを生成し、それらをチェックします。

正確性チェックはシャーディング構成に依存します:

  • 単一コアの場合、生成されたトークンは同じチェックポイントの CPU 実行のものと比較されます。
  • モデルがシャーディングされている場合、小さな判定モデルが生成されたテキストが流暢な英語であるかをチェックします。

PASS の範囲は意図的に制限されています。現時点では、4 つの Trainium2 デバイスに収まるモデルのみがテストされており、最大 16 コアです。PASS はスループットやより長い生成については何も述べていません。

結果はデータセット

Space が表示するすべてのものは、公開 aws-neuron/torch-neuronx-model-support データセットから来ており、neuronx-cc とライブラリバージョンのペアごとに 1 つの Parquet ファイルを保存します。各行には、モデルの結果、失敗時のエラー、スタック全体のバージョン、および実行された run.py が含まれます。

レンダリングされたテーブルではなくデータセットであるため、結果を直接クエリできます:

import pandas as pd

df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]

データセットカード は各結果値の正確な意味を定義しているため、PASS や失敗を推測ではなく正確な方法論に基づいて解釈できます。

カバレッジの制限と今後の予定

現在 Space が証明できることには 4 つの制約があります:

  1. Transformers モデルのみがテストされています。
  2. text-generation と image-text-to-text タスクのみがカバーされています。
  3. 4 つの Trainium2 デバイス (16 コア) に収まるモデルのみがテストされています。
  4. PASS は 32 生成トークンの正確性を報告し、スループットや長いコンテキストの動作は報告しません。

計画されている拡張は、diffusers、sentence transformers、およびその他の Transformers タスクへのカバレッジ拡大です。

なぜ重要か

ハードウェアの移植性に関する質問は、通常、暗黙の知識、アドホックな問題、または高価な試行実行によって答えられます。Trainium の互換性をバージョン管理された検索可能なデータセットとして公開し、再現スクリプトを添付することで、その質問の経済性が変わります。答えはルックアップであり、証拠は検査可能です。

2 つの設計上の選択により、このリソースはリーダーボードよりも有用です。まず、結果は特定の neuronx-cc およびライブラリバージョンにキー付けされているため、実際にデプロイする予定のスタックに対して互換性を評価できます。次に、データセットは成功とともに失敗とそれを生成したエラーを記録するため、以前は文書化されていなかった非互換性が検索可能になります。

実用的な要点: AWS Trainium でモデルを実行する前に、aws-neuron/torch-neuronx-model-support Space を確認するか、その Parquet データセットを直接クエリし、公開されている run.py と NEFF キャッシュエントリを再利用して trn2 インスタンスで結果を再現してください。

Sources