Amazon は、機械学習ワークロードにカスタム シリコンを採用する際の最大の実際的な障壁の 1 つを取り除く動きとして、AI チップの Trainium ファミリ用のオープンソース ネイティブ PyTorch バックエンドである TorchNeuron を発表しました。 10月20日のPyTorch Conference North Americaに先立って公開されたPyTorchブログの公式投稿によると、開発者は独自のツールを中心にワークフローを書き直すのではなく、単一行の変更で既存のPyTorchモデルコードをTrainium上で実行できるようになったという。

この発表は、アナリストが長い間Nvidiaの最も永続的な利点として述べてきたもの、つまりチップの生のパフォーマンスではなく、開発者をGPUに結び付けるCUDAソフトウェアエコシステムをターゲットにしているため、すぐに注目を集めました。ロイターは、アマゾンが同社のAIチップがAIで最も人気のある開発フレームワークで動作すると発表したとの報道をし、一方TechCrunchは、トーチニューロンをNVIDIAに対する同社の最大の狙いであると呼んだ。アクセラレータ市場がどのように変化しているかをより広範に把握するために、当社が継続的に行っている AI ハードウェア ニュース報道 では、これらの開発の展開を追跡しています。

トーチニューロンが実際に変えること

TorchNeuron の核となる約束はシンプルさです。 PyTorch ブログで公開されている例では、`.to('neuron')` を使用してモデルを新しいデバイスに移動するだけで、既存のトレーニング ループが Trainium 上で実行されます。投稿によると、PyTorch が推奨するデバイスに依存しない API である「torch.accelerator」に対して書かれたコードは、変更されずに実行されます。コンパイルされたワークロードの場合、開発者は `torch.compiler.set_default_backend("neuron")` でデフォルトのバックエンドを一度設定し、それ以上変更せずに `model.compile()` を呼び出します。

内部では、TorchNeuron は、カスタム ハードウェア バックエンド用の PyTorch の公式拡張ポイントである PrivateUse1 を通じて Trainium を登録します。これは、チップがフレームワークにボルトで固定された付属物ではなく、標準の PyTorch デバイスとして動作することを意味します。 AWS は、統合が複数の実行モードにわたるものであると説明しています。

  • 研究とデバッグのための積極的な実行。実験的な「適応的熱心モード」により、起動オーバーヘッドが削減されます。
  • torch.compile は実稼働パフォーマンスを実現し、Neuron コンパイラーが最適化するグラフにモデルをキャプチャし、FlexAttend をネイティブにサポートします。
  • ニューロン カーネル インターフェイス (NKI) カスタムの低レベル カーネルを必要とする研究者向け

発表によると、このバックエンドはTrainium2チップと新しいTrainium3チップの両方で動作するという。

Nvidia のソフトウェア堀に対する直接攻撃

Nvidia が AI コンピューティングを掌握しているのは、決して高速なハードウェアだけではありません。同社の20年前からあるソフトウェアプラットフォームであるCUDAは、ほとんどのAI研究コードが書かれ、調整され、展開される場所であり、AMDのInstinctシリーズからGoogleのTPUやAmazonのTrainiumに至るまで、あらゆるチャレンジャーチップは、開発者にCUDAから離れるよう求めるコストと闘わなければならなかった。 GeekWire、Benzinga、GuruFocus での TorchNeuron の発売に関する報道は、一律にこの発表をその堀への攻撃として組み立てました。

Amazon は、開発者を説得するために PyTorch から開発者を引き離そうとしているわけではありません。代わりに、Trainium を PyTorch に導入しているため、フレームワークの膨大なユーザー ベースは習慣を変えることなくこのチップを採用できます。 PyTorch のブログ投稿には、フレームワークを維持する Meta からの支持が含まれていました。「AWS Trainium のネイティブサポートにより、その選択肢が広がり、研究者やエンジニアは、すでに知っている熱心な実行と分散 API を使用して実験および反復できるようになります」と、Meta の PyTorch エンジニアリング担当ディレクターである Jana van Greunen 氏が投稿の中で述べています。

その枠組みは AWS のクラウド ビジネスにとって重要です。 Trainium インスタンスは、トレーニングと推論用の Nvidia GPU の低コスト代替品として販売されており、このチップはすでに有力な AI 顧客を魅了しています。チップが PyTorch にネイティブであると感じられるようにすることで、広範な採用を遅らせている問題、つまり移植作業に直接対処できます。

技術的範囲、ノートブックからクラスター規模のトレーニングまで

PyTorch ブログにリストされている機能は、Amazon が趣味の実験をはるかに超えたところを目指していることを示唆しています。 TorchNeuron は、BF16、FP32、FP16 の精度で、テキスト、マルチモーダル、ビジョン言語モデルに加え、画像とビデオの生成、状態空間モデル、ワールド モデルに見られる演算子セットをカバーします。 TorchAO と FP8 のサポートは、次のリリースでの一般提供に先立ってプライベート ベータ版で提供されます。

単一 GPU コードをクラスタースケールの実行に変換する分散トレーニングは、標準の PyTorch API (DDP、FSDP2、SimpleFSDP、テンソルおよびパイプライン並列処理、ZeRO スタイルのシャーディング) と、多次元プロセス グループ用の DTensor および DeviceMesh を通じてサポートされています。同投稿によると、ストリームとイベントはPyTorchのセマンティクスをTrainiumのハードウェアキューにマッピングするという。

推論のために、Amazon の Neuron スタックは同じ基盤上に構築されています。vLLM Neuron は大規模な言語モデルを提供し、vLLM Omni Neuron は画像とビデオ生成の拡散モデルを提供します。 Neuron カーネル インターフェイスを使用すると、研究者は熱心な実行とコンパイルされた実行の両方のためのカスタム オペレーターを作成できます。AWS は、主流の LLM およびマルチモーダル シナリオ向けの FlashAttendant および MoE ブロック単位の計算をカバーする、Trainium に最適化されたカーネルのセットである NKI-lib をオープンソース化しています。

パフォーマンスツールも含まれています。 Neuron Explorer は、標準の「torch.profiler」を拡張して、フレームワークからハードウェアまで実行を追跡し、すべてのイベントをそのソース行にリンクし、Parquet 形式のメトリクスをエクスポートして、自動ボトルネック分析のためにエージェントがポストノートを使用できるようにします。

可用性と今後の予定

ブログ投稿によると、TorchNeuron は 10 月 20 日に開催される PyTorch Conference North America でオープンソースとして公開される予定で、互換性は Hugging Face Transformers v5 まで拡張されます。 AWS は、Trainium をクラウドスケールの AI トレーニングと推論のためにゼロから設計されたチップのポートフォリオであると説明しています。

戦略的な背景は、すべての主要なクラウド プロバイダーが同時に Nvidia アクセラレータの最大の購入者であり、競合他社になる可能性がある市場です。 Microsoft、Google、Meta はいずれもカスタム シリコンを開発しており、Huawei を含む中国企業は CUDA の重力から逃れるために独自のソフトウェア スタックを構築しています。 Amazon は、開発者に移行を依頼するのではなく、PyTorch 内で開発者に会うことで、そのコンテストへの最も混乱の少ない道を選択しました。そして、開発者の経験が維持できれば、実際のトレーニング ワークロードを Nvidia ハードウェアから移行できる可能性が最も高い道を選択しました。

TorchNeuron が本当に CUDA の利点を侵食するかどうかは、大規模なパフォーマンスと現実世界のモデルのロングテールに依存しますが、発売資料では答えられていない質問です。しかし、ほとんどの PyTorch 開発者が Trainium を試行することさえ妨げていた障壁 (書き換え) は、現在では 1 行のコードまで引き下げられています。

---

AI の先を行く

最新の AI ニュース、分析、画期的な情報をすべて 1 か所で入手できます。

AI ニュースをもっと読む→