AWS、PyTorchをTrainiumでそのまま動かすTorchNeuronを公開　PrivateUse1の標準バックエンドとしてApache 2.0、10月20日にパブリックベータ

TorchNeuronは、PyTorchコミュニティが用意したアクセラレータ非依存の拡張点（PrivateUse1デバイスバックエンドとtorch.accelerator API）に載る、Trainium向けのネイティブバックエンドだ。AWSは研究・デバッグ向けのeager実行（性能を上げたadaptive eagerを含む）、本番向けのtorch.compile、カスタムカーネル用のNeuron Kernel Interface（NKI）をTrainium2とTrainium3で使えると説明している。プロファイリングはNeuron Explorer経由でtorch.profilerと同じ経路を使い、フレームワークからハードウェアまでのイベントをソース行に紐づけられる。公開リポジトリはgithub.com/aws-neuron/torch-neuronxで、ライセンスはApache 2.0。パブリックベータは10月20日からで、pip install torch-neuronxでの導入が案内されている。Hugging Face側もTransformersやDiffusersでの対応を同日前後に紹介している。従来のoptimum-neuron経由の体験とは別に、標準PyTorchデバイスとしてTrainiumを選ぶ経路が加わった形になる。