Amazon、PyTorchコードを書き換えずにAWS Trainiumで動かすTorchNeuronを発表

TorchNeuronは、PyTorchが備える新ハードウェア登録の標準的な仕組みに接続する形で実装されている。そのため、テスト、デバッグ、スケーリングに使う既存ツールがそのままTrainiumで使えるという。Amazonのエンジニアは同じ学習スクリプトを従来型チップとTrainiumで実行する比較デモを行い、設定の違いだけでモデルが同じ速度で同じように学習したと説明している。

Amazonは、AI開発のコストはチップ価格だけでなくエンジニアの作業時間にもあると位置づけている。新しいチップへの移行で書き換えや再検証が不要になれば、特にスタートアップは製品改善に時間を回せる。大企業については、オフピーク時に推論用の容量を学習ジョブへ振り向けるといった運用が挙げられている。

初期の利用者として、音声AIのCartesiaと、創薬研究を行うLieber Institute for Brain Developmentが紹介されている。Cartesiaは標準の推論パイプラインを最小限の変更でTrainium上で動かした。Lieber Instituteは分子動力学シミュレーションを加速または代替するAIツールの検討にTorchNeuronを使っている。

エコシステム面では、特定チップに依存するコードを除くため、AmazonがPyTorch本体へ変更を上流提供している。Hugging Faceは自社ツールでTrainiumのネイティブ対応を有効にし、TorchNeuronは分散学習フレームワークTorchTitanとFSDPにも対応する。さらにvLLMへの大きな貢献と、性能カーネル一式のオープンソース化も行われた。ビルドごとに数十万のAIモデルでテストしているとも述べられている。