PyTorchモデルをTrainiumへ移すTorchNeuron
AWSは、PyTorchで書かれたAIモデルを自社のAIチップ「Trainium」で実行するソフトウェア「TorchNeuron」を発表しました。既存コードを大きく書き換えず、設定を切り替えて実行できるとしています。PyTorchはAIモデルの開発や学習、実行に使われるオープンソースのソフトウェアです。
AI開発では、NVIDIAのGPU向けソフトウェア基盤CUDAがPyTorchを支えてきました。そのため、PyTorchで作ったアプリをNVIDIA製品で動かす環境は広く整っています。TorchNeuronは、PyTorchの実行先にTrainiumを加え、AWSのチップを使う選択肢を広げます。
PyTorchのコードをAMD製GPUやApple製チップへ移す場合も、設定変更で対応できるケースがあります。AWSは、Trainiumへの移行に加え、自社クラウド上のインフラを利用できる点を利点に挙げています。AWSのコリン・ブレイス氏は、試作から本番運用までの基盤をAWS上で利用できると述べました。
NeuronがPyTorchとチップをつなぐ仕組み
Trainiumは、AWSがクラウド上のAI学習と推論向けに提供するAIチップです。TorchNeuronはPyTorchの標準的な拡張機構を通じてTrainiumをデバイスとして登録し、演算をNeuron側へ振り分けます。開発者はPyTorchの標準APIを使うモデルを「neuron」デバイス上で動かせます。
AWS Neuronは、Trainiumや推論向けチップInferentiaでAIモデルを実行するための開発ソフトウェア群です。コンパイラーや実行時ソフトウェア、学習・推論用ライブラリー、性能分析やデバッグのツールを含みます。TorchNeuronはその中で、PyTorchのモデルをTrainium上で動かすバックエンドとして機能します。
TorchNeuronでは、演算を逐次実行するeagerモードと、モデルをグラフとしてコンパイルする「torch.compile」を利用できます。eagerモードでは対応する演算を個別に実行し、torch.compileではグラフをコンパイルして実行します。さらに、特定の演算を細かく最適化するためのNeuron Kernel Interface（NKI）も利用できます。
導入例とDeepgramの処理効率
AWSは、同じAI学習スクリプトを一般的なチップとTrainiumで実行する比較デモを行いました。AWSによると、コード上の違いは使用するチップを指定する設定だけで、どちらのチップでも同じモデルを同じ速度で学習しました。TorchNeuronはTrainium2とTrainium3に対応し、PyTorchの標準的な機構を通じて既存の開発ツールを利用できます。
音声AIを手がけるDeepgramは、音声認識モデルNova-3をTrainium2に移し、オフラインの一括処理向けに調整しました。AWSとDeepgramは数カ月かけてPyTorchモデルを移植し、最終構成では1台のEC2 trn2.3xlargeインスタンスでモデルを8つ動かしました。Deepgramは、評価時に使用していた計算環境と比べ、1ドル（約160円）あたりに処理した音声の量が11〜70％多かったと報告しています。
この値は、音声をまとめて処理するオフライン推論の測定結果です。Deepgramは、複数の処理を並行して動かすRust製ランタイムも構築しました。処理量が最大となるバッチサイズでは、Pythonのマルチプロセス方式に比べてスループットが7.5％向上したとしています。これらの結果はDeepgramの特定モデルと処理条件で得たものです。
オープンソース化とPyTorchへの変更
AWSはTorchNeuronをApache 2.0ライセンスで公開し、2026年10月20日のPyTorch Conference North Americaに合わせてパブリックベータとして提供する予定です。AmazonはTorchNeuronの提供に加え、特定のチップに結び付いたコードをPyTorch本体から減らす変更をコミュニティに提案しています。この変更は、Trainiumに限らず、複数のハードウェアでPyTorchを使いやすくすることを目指しています。
TorchNeuronは、TorchTitanを使った分散学習にも対応し、複数のTrainiumチップへ学習を広げる標準的な手段を提供します。Hugging Faceの開発ツールにもTrainiumのネイティブ対応が加わりました。開発者は標準的なPyTorchの操作に加え、必要に応じて複数チップへの拡張やNKIによる個別最適化を選べます。
AWSはDeepgramの事例で、既存のPyTorchモデルをTrainiumへ移した後に、カーネルや実行ランタイムを調整した過程も公開しました。TorchNeuronによるコード移行の容易さと、実運用での性能調整は別の段階として示されています。
