AWS Trainium:あなたのモデルは互換性がありますか?確認できるツール

AWS Trainium:あなたのモデルは互換性がありますか?確認できるツール
1 分の読書

Transformersライブラリのモデルは、オープンソースの機械学習ライブラリであるPyTorchのネイティブバックエンド「TorchNeuron」を介して、AWSのTrainiumチップ上で動作可能です。問題は、実際にどのモデルが機能するかです。Hugging Faceが公開したSpace「torch-neuronx model support」は、Hubで最も人気のあるモデルについてこの疑問に答えます。

モデルを検索する

Spaceを開くと、最初に「Selected」ビューが表示されます。これは、Hubで最もダウンロードされ、トレンドとなっているモデルから選ばれた選択肢です。各行には、テスト結果、モデルが分散されたTrainiumコア数、コンパイルを含む総実行時間、デバイス上の最大使用メモリが表示されます。

特定のモデルを検索するには、「Trending」または「Downloads」ビューに切り替え、モデル名の一部を入力します。これらのビューは、タスクごとにHubのランキングを追跡しており、モデルがテスト済みかどうかに関係なく表示されます。モデルが除外された場合、Spaceはその理由を明記します。ほとんどの場合、優先度が低いためですが、実際に互換性がない場合もあります。ユーザーは、検証対象とするコンパイラNeuron(neuronx-cc)とTransformersのバージョンを選択することもできます。

現時点では、Transformersモデルのみがテスト対象であり、選択範囲はテキスト生成およびテキスト・画像からテキストへのタスクに限定されています。チームは、diffusers、sentence transformers、およびTransformersの他のタスクにもこの仕組みを拡張する計画です。

自分でモデルを実行する

モデルを選択すると、テストに使用されたNeuronソフトウェアスタック、インストールすべきパッケージ、およびテストで使用されたスクリプト「run.py」が表示されます。このスクリプトは、モデルをfrom_pretrainedでロードし、必要に応じて複数コアに分散させ、その後、静的キャッシュとtorch.compile(backend=「neuron」)を使用してテキストを生成します。

例えば「AWS Neuron Deep Learning Container」内など、同じスタックを持つtrn2インスタンス上では、パッケージをインストールした後、python3 run.pyを実行するだけです。torchrunは必要ありません。モデルが分散されている場合、スクリプト自体が分散プロセスを起動します。最初の実行では、モデルのグラフのコンパイルから始まります。これを省略するには、Spaceが提供するコマンド(コンパイラのバージョンを考慮したもの)を使用して、公開NEFFキャッシュバケットからテスト時に生成されたグラフをコピーできます。

「機能する」の意味

各テストでは、モデルをtrn2インスタンスにロードし、その重みを格納できる最小数のTrainiumコアに分散させます。ほとんどのモデルは、計算を複数のコアに分割するテンソル並列処理を使用します。テンソル並列処理計画のない混合エキスパートモデルは、エキスパート並列処理のみを使用します。テストはその後、32トークン(テキストの単位)を生成し、それを検証します。単一コアでは、これらのトークンは同じチェックポイントをCPUで実行した場合のトークンと比較されます。モデルが分散されている場合、小さな判定モデルが、テキストが流暢な英語であることを確認します。

テスト対象は、最大16コア(4つのTrainium2デバイス)に収まるモデルのみです。PASSという結果は、スループットやより長い生成については何も示しません。

すべての結果は、neuronx-ccとライブラリのバージョンの組み合わせごとにParquetファイルを含む公開データセット「aws-neuron/torch-neuronx-model-support」から得られています。各行には、モデルの結果、発生したエラー(もしあれば)、スタックのバージョン、および実行されたrun.pyが記録されています。例えばQwen/Qwen3-0.6Bモデルについて、pandasを使用して直接クエリを実行することが可能です。データセットの説明には、各結果の正確な意味が記載されています。

S’abonner
Notification pour
guest

0 Commentaires
Le plus ancien
Le plus récent Le plus populaire
0
Nous aimerions avoir votre avis, veuillez laisser un commentaire.x