Amazon、AIチップTrainium向けにPyTorchをほぼ無改修で動かす「TorchNeuron」を発表

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

AIによる概要

Amazonは独自AIチップ「Trainium」上でPyTorchのコードをほぼ書き換えずに動かせるソフトウェア「TorchNeuron」を10月19日（日本時間）に発表した。設定を1か所変えるだけで、既存コードをTrainium2とTrainium3で実行できるという。Apache 2.0ライセンスのオープンソースとして、10月20日のPyTorch Conference North Americaで公開する予定である。

解説

TorchNeuronはPyTorchの標準的なデバイス登録の仕組みに接続する。このため、テストやデバッグ、スケーリングに使っているツールをTrainiumでもそのまま使えるとしている。デモでは同じ学習スクリプトを従来型のチップとTrainiumで実行し、違いは使用するチップを指定する設定だけだった。学習結果と速度は同等だったという。

複数チップに広げるためのツールや、性能を引き出すための低レベルのカスタムコードを書く機能も備える。Amazonはビルドごとに数十万のAIモデルでテストしているとしている。また、AI研究者や開発者の85％がPyTorchでコードを書いているという。