開発

AWS Trainiumのモデル互換性検索を公開

原題: Does Your Model Run on AWS Trainium? Look It Up

重要提供開始Hugging Face

概要

Hugging Faceは2026年10月23日、Hub上のモデルがAWS Trainiumで動くかを確認できる「torch-neuronx model support」を公開しました。

Spaceでは人気モデルや注目モデルを検索し、動作結果、使用したTrainium core数、コンパイルを含む実行時間、peak device memoryを確認できます。

各モデルについて、Neuron stackと追加パッケージ、実際に検証に使ったrun.pyを表示します。

AWS Neuron Deep Learning ContainerなどのTrainium2環境で、表示されたスクリプトを実行できます。

同じcompilerで生成した公開NEFF cacheを取得すれば、最初のgraph compilationを省略できます。

検証はモデルを収容できる最少数のTrainium coreへshardingし、32tokenを生成して結果を確認します。

単一coreではCPU版とtokenを比較し、sharding時は小型judge modelで英語文の流暢さを確認します。

公開時点ではTransformersのtext-generationとimage-text-to-textが対象で、最大4台・16coreのTrainium2に収まるモデルを試験しています。

結果はaws-neuron/torch-neuronx-model-support datasetでも公開され、compilerとlibraryのversion、エラー、実行スクリプトを直接参照できます。

PASSはthroughputや長い生成の性能保証ではなく、記載条件でモデルを実行できたことを示します。