AWSがTorchNeuronをパブリックプレビュー、TrainiumでPyTorchをネイティブ実行可能に　コンパイルは静的形状に限定

AWSは10月20日、TorchNeuronのパブリックプレビューを公開した。Trainium2、Trainium3向けのネイティブPyTorchバックエンドを提供する。開発者は.to('neuron')でモデルと入力を移し、使い慣れたテンソル操作や学習インターフェースをそのまま利用できるほか、Neuron Kernel Interface（NKI）でカスタムカーネルを追加できる。公式発表

技術的な要点は、PyTorchのPrivateUse1機構を使ってneuronデバイスを登録することだ。演算はデフォルトで個別に実行されるが、torch.compileのneuronバックエンドでモデルをコンパイルすることもできる。プロジェクトはDTensor、分散集合通信、シャーディング学習にも対応する。NKIカーネルをコンパイルグラフに組み込めるため、ハードウェア最適化をPyTorchのワークフロー内に保てる。ソースコードと使用方法

ただし、デバイスの切り替えは移行の出発点にすぎない。ドキュメントによると、コンパイルモードがサポートするのは静的形状のみで、新しい入力形状ごとに再コンパイルが発生する。非連続テンソルは呼び出しのたびにコピーが必要だ。float64はfloat32で実行され、TorchAOとの統合もまだ提供されていない。API互換性表では、ネイティブ対応、CPUへのフォールバック、未対応の項目が区別されているため、エンジニアリングチームは精度とデータ転送コストを確認する必要がある。制約の説明

同日に発表されたLaplace 3Bの事例からは、実際の移植作業の一端がうかがえる。チームはGPU専用のTritonアテンションカーネルを標準的なPyTorch演算に書き換え、Muonオプティマイザーの通信方式を再設計し、逆伝播の展開方法も調整した。AWSと共同チームによると、単一のtrn2.48xlargeで学習スループットは毎秒約12万トークンだった。これは特定のアーキテクチャと構成で得られた結果であり、他のモデルにそのまま当てはめることはできない。技術事例

代替アクセラレーターを評価するチームにとって、新しいバックエンドは試験導入のハードルを下げるとともに、実装を確認できるオープンソースの入り口となる。今後は演算子のカバレッジ、初回コンパイル時間、形状バケット化のコスト、エンドツーエンドの数値的一貫性を確認する必要がある。インストール時はバージョンにも注意が必要だ。リポジトリによると、以前から存在する同名のtorch-neuronxパッケージはXLAベースのため、新版で指定されたバージョンの組み合わせに従って導入する必要がある。インストール手順