以前の 記事で、AWS TrainiumのネイティブなPyTorchバックエンドであるTorchNeuron（ torch-neuronx パッケージ）を紹介しました。これにより、Transformersは他のデバイスと同様に、Trainium上で独自のモデリングコードを実行できます。つまり、Hub上のあらゆるTransformersモデルが実行できる可能性があります。実際に動くのはどれでしょうか？
この torch-neuronxモデルサポート Spaceは、Hubで最も人気のあるモデルについてその疑問に答えます。モデルがTrainium上で修正なしで動作するかどうかを示し、自分で実行するためのスクリプトを提供します。
モデルを調べる
Spaceは、 Selected ビューで開きます。これは、ダウンロード数とトレンドスコアに基づいてHubの上位モデルから選ばれた、独自のモデルセレクションです。各行には、結果、モデルがシャーディングされているTrainiumコアの数、実行の実時間（コンパイルを含む）、デバイスメモリのピーク値が表示されます。
特定のモデルを調べるには、 Trending ビューまたは Downloads ビューに切り替えて、名前の一部を入力してください。これらのビューは、テスト済みかどうかにかかわらず、Hub自身のランキングにタスクごとに従います。モデルが意図的に除外されていた場合、Spaceはその理由を示します（多くの場合は優先されなかっただけで、時には本当に互換性がないこともあります）。Neuronコンパイラ（neuronx-cc）とTransformersのバージョンを選択して表示することもできます。
執筆時点では、Transformersモデルのみがテストされており、セレクションは
text-generationモデルとimage-text-to-textモデルに限定されています。今後、diffusers、sentence transformers、およびその他のtransformersタスクへ拡大する予定です。
自分で実行する
モデルを選択すると、実行方法が表示されます。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キャッシュ bucket。Spaceには、正しいコンパイラバージョンを指定したコマンドが表示されます:
hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache
「runs」の意味
各テストはモデルを1つの trn2 インスタンスを、その重みを保持できる最小数の Trainium コアにシャーディングします。ほとんどのモデルはテンソル並列でシャーディングされます。テンソル並列のプランがない Mixture-of-experts モデルは、エキスパート並列のみを使用します。その後、テストは 32 トークンを生成して検証します。単一コアの場合、トークンは同じチェックポイントの CPU 実行のものと比較されます。モデルがシャーディングされている場合は、小さなジャッジモデルがテキストが流暢な英語であることを確認します。
執筆時点では、4 台の Trainium2 デバイスに収まるモデルのみがテストされており、最大でも 16 コアを意味します。A 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"]]
この データセットカード には各結果の正確な意味が記載されています。
次にあるモデルが Trainium で動作するかどうか疑問に思ったら、 調べてみてください.