AWSがPyTorchアプリをTrainiumで動かす「TorchNeuron」を公開

発表されたこと

AWSは米国時間の月曜日、AIスタートアップがPyTorch製のAIアプリをAmazonのAIチップ「Trainium」へ移しやすくするツール「TorchNeuron」を公開した。PyTorchは無料のオープンソースツールキットで、開発者がAIモデルの構築・学習・実行に使う。TorchNeuronはPyTorchから派生させたもので、PyTorchを知る開発者なら使い方はそのまま分かるとしている。

Amazonによれば、移行は設定を1行変えるだけで完了し、コードの書き換えは要らない。同社はAIアプリのおよそ85%がPyTorchで書かれていると説明しており、この母集団を自社チップへ取り込む狙いが読み取れる。PyTorchはFacebookが開発し、2022年からはオープンソースプロジェクトを維持するLinux Foundationの管理下にある。

なぜNvidiaへの挑戦になるのか

Nvidiaは長年、CUDAと呼ばれるソフトウェア基盤を通じて、PyTorchで書かれたAIアプリが同社GPUなどで問題なく動くように整備してきた。その結果、NvidiaのハードウェアはPyTorchの内部に深く組み込まれ、事実上の既定の選択肢になっている。今回の発表は、このソフトウェア面の優位に直接切り込むものだ。

ただし、ハードウェアへの移植を簡単にする取り組みでAmazonが最初ではない。原文によれば、PyTorchの利用者は設定変更だけでAMDのGPUやAppleのチップへアプリを移すこともできる。TorchNeuronは「移植の容易さ」だけで差別化するものではなく、AWSのクラウド全体を使える点を前面に出している。

AWSの主張と利用事例

AWSのカスタムチップ部門であるAnnapurna LabsでTechnology Software担当バイスプレジデントを務めるColin Brace氏は、TechCrunchへのメールで、AWS上では試作から本番までのインフラがすでに揃っており、創業者が素早く動き必要なときに規模を広げられると述べた。Nvidiaのハードウェアに結びつくneocloud（特定ベンダー向けに最適化されたGPUクラウド）としてCoreWeave、Nebius、Firmusの名を挙げ、そうした選択肢より優れると主張している。

Brace氏は、同じチップが学習（AIモデルを教え込む工程）と推論（完成したモデルを実際に動かす工程）の両方を処理し、クラウド基盤と密接に統合されている点がコスト削減につながると説明した。AWS自身、トークンあたりの単価が最安ではないと認めたうえで、価格性能比と学習・推論の一体利用でワークロード全体の経済性が良くなるとしている。

事例として挙げられたのが音声AIアプリを手がけるDeepgramだ。同社はElevenLabsの競合とされ、PyTorchのネイティブサポートを使ってTrainiumへ移行し、実運用に耐える精度を保ちながら1ドルあたり70%多くの音声を処理しているとAWSは説明する。ただしこれはAWS側の主張であり、第三者による検証結果ではない。

PyTorch上流への働きかけと不確実性

Amazonは自社ハード向けのツール提供にとどまらず、PyTorch本体の変更も働きかけている。フレームワークを特定チップに結びつけるコードをPyTorchコミュニティと協力して取り除き、Trainium利用者に限らず開発者コミュニティ全体にとってハードウェア非依存度を高めることが目的だとしている。Brace氏は、開発者から実際の選択肢を奪ってきたロックインを解消することが本質だと語った。

一方で原文は、PyTorchの利点の一つがハードウェア上でAIコードを動かす際の煩雑さを隠すことにあると指摘し、Nvidiaを排除する方向の変更が開発者に選択肢と低価格をもたらすのか、それとも設定の手間を増やすだけなのかは見極めが必要だとしている。上流への提案がどの程度受け入れられるかも現時点では不明だ。