PyTorchにNeuronをネイティブ統合、TorchNeuronが公開ベータへ

発表されたこと

Hugging Faceのブログが、AWSのAIチップ「Trainium」向けの新しいPyTorchバックエンド「TorchNeuron」を紹介した。TorchNeuronはPyTorchのプラグインとして提供され、eagerモード、torch.compile、FSDPやテンソル並列（TP）といったPyTorch標準の分散APIをサポートすると説明されている。

提供開始は2026年10月20日からのパブリックベータで、pip install torch-neuronxでインストールできる。TorchNeuron自体の詳細はPyTorchのブログに譲るとして、この記事はHugging Faceライブラリ側の統合に焦点を当てている。

何が新しくなったのか

これまでTrainiumとInferentiaでのHugging Faceの利用は、optimum-neuronというライブラリを通じて提供されてきた。土台となるtorch-neuronxはPyTorch XLAベースで、遅延テンソル、静的シェイプのみ、グラフブレークごとの再コンパイルといった制約があり、利用者は「コンパイラのように考える」必要があった。Optimum Neuronはモデリングクラスやトレーナーなどのカスタムコードを維持することで、この複雑さを隠していた。

TorchNeuronはこの構造を変える。PyTorchに直接統合され、neuronをネイティブなデバイスとして公開するため、eager実行もtorch.compileもPyTorchの作法のまま使える。Hugging Faceは、今後Neuron利用者はOptimum Neuronを経由せず、Transformers、Diffusers、Accelerate、PEFT、TRLを直接使うことを推奨している。

推論と学習の具体像

推論では、上位100のテキスト生成モデルアーキテクチャのうち89%がそのまま対応する。Optimum Neuron v0.4.5では46%だったため、対応範囲は大きく広がったことになる。コード上はdevice_map="neuron"を渡すだけでモデルがNeuronデバイスに配置され、生成時には静的キャッシュの利用が強く推奨されている。頻繁な再コンパイルを避けるためである。

学習は、GPUで使ってきたTRLのSFTTrainerスクリプトがほぼそのまま動く。DistributedConfigでTPとFSDPのサイズを指定すると、重みの読み込み時にモデルがシャーディングされ、Trainer側は既に分散済みのモデルを再ラップしない。--torch_compile trueを渡せばTrainium上ではNeuronバックエンドが自動的に選ばれる。

記事では8個のTrainium2チップ（32 NeuronCores、trn2.48xlargeの半分）で実行した例が示されている。TP_SIZE=4、FSDP_SIZE=8で32プロセスを起動し、学習は完走する。保存されるチェックポイントはNeuronに依存しない標準のTransformers形式で、そのままHubに公開して他のハードウェアで読み込めるとしている。

画像生成とカーネル

Diffusersでも、多くのパイプラインがeagerモードとコンパイルモードの両方でTorchNeuronに対応する。FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1、H3はテンソル並列にも対応しており、Neuronデバイスを活かせるようになった。テンソル並列に対応していないパイプラインについては、ブログで案内されている手順に沿って自分で対応を追加できるという。

さらにNeuronチームによるカーネルもHubで順次公開される。カーネルはNeuron Kernel Interface（NKI）で記述でき、低レベルな制御で演算を最適化できる。記事では、視覚トランスフォーマを高速化するカーネルの例として、Trainium上で最大2.4倍の高速化が挙げられている。

現時点の制約と不明点

現状はパブリックベータであり、いくつかの制約が残る。Neuronバックエンドはtorch.compileの自動dynamic shapesにまだ対応しておらず、学習スクリプトでは無効化している。Transformersの非同期重みロードもNeuronでは未対応のため、HF_DEACTIVATE_ASYNC_LOAD=1で無効化する必要がある。

テンソル並列に対応するDiffusersパイプラインはまだ一部で、今後の拡充が予告されている段階だ。Optimum NeuronからTorchNeuronへの具体的な移行手順や、既存のカスタムコードがどこまで互換なのかは、この記事では明らかにされていない。インスタンスの利用料金についても原文に記載はない。