AWSは10月19日から20日にかけて、自社AIチップ「Trainium」向けのネイティブPyTorchバックエンド「TorchNeuron」を公開した。PyTorch公式ブログとAWS Neuronの公式ドキュメントで発表されたもので、米国で開催中の「PyTorch Conference North America」にあわせ、10月20日からApache 2.0ライセンスのオープンソース・パブリックベータとして提供が始まっている。
これまでTrainium上でPyTorchを動かすにはXLAを介した変換が必要だったが、TorchNeuronはPyTorchのデバイス切り替え機構「PrivateUse1」を利用し、既存コードのデバイス指定を「cuda」から「neuron」に変更する程度の少ない書き換えで動作させられるよう設計されている。研究・デバッグ向けにPythonの制御構文やprint文がそのまま使える「eager」実行モードと、本番運用向けに`torch.compile(backend="neuron")`でNeuronコンパイラへ処理を渡す経路の両方に対応する。現時点ではNeuron SDK 2.33のパブリックベータとして、Trn2・Trn3インスタンス向けに提供されている。
Nvidiaの優位性の中核は、CUDAを中心としたソフトウェア環境にあるとされてきた。AmazonはAI開発者の間で広く使われる無料ツールであるPyTorchを足がかりに、その環境への依存を減らす狙いとみられる。AWSはこれまでもTrainium3の投入や、Nvidia製GPUと自社チップを選べる「AI Factories」の展開などを通じてNvidia依存の低減を進めており、今回のTorchNeuronはその流れの延長線上にある一手だ。
出典: PyTorch Blog
