米アマゾン、AIチップTrainiumでPyTorchをほぼ無改修で動かす「TorchNeuron」発表

　米アマゾンは10月19日（日本時間）、同社が自社開発するAIチップ「Trainium」上で、AI開発フレームワーク「PyTorch」のコードをほぼ書き換えずに動かせるソフトウェア「TorchNeuron」を発表した。設定を1か所変えるだけで、既存のPyTorchコードをTrainium2とTrainium3で実行できるという。


　TorchNeuronはPyTorchの標準的なデバイス登録の仕組みに接続する。このため、テストやデバッグ、スケーリングに使っているツールを、Trainiumでもそのまま使えるとしている。デモでは同じ学習スクリプトを従来型のチップとTrainiumで実行し、違いは使用するチップを指定する設定だけだった。AIモデルの学習結果と速度は同等だったという。複数チップに広げるためのツールや、性能を引き出すための低レベルのカスタムコードを書く機能も備える。アマゾンはビルドごとに数十万のAIモデルでテストしているとしている。

　ライセンスはApache 2.0で、10月20日に開催されるPyTorch Conference North Americaで公開する予定という。アマゾンによると、現在、AI研究者や開発者の85％がPyTorchでコードを書いている。

　対応や利用の動きも広がっている。Hugging Faceは、人気の開発ツールでTrainiumをネイティブに対応させ、アマゾン固有のコードなしで使えるようにした。PyTorchの分散学習フレームワーク「TorchTitan」とも統合し、大規模学習の標準的な経路を用意した。リアルタイム音声AIのCartesiaは、標準の推論パイプラインをPyTorchコードの最小限の変更でTrainium上で動かした。神経精神疾患の治療法を研究するLieber Institute for Brain Developmentは、創薬に使う分子動力学シミュレーションの高速化をTrainium上で検討している。

　アマゾンは2015年にチップ設計会社のAnnapurna Labsを買収して以降、4つのチップファミリーを自社開発してきた。ビジネス＋ITでは、米シノプシスとの複数年契約も報じている。