AmazonのTrainiumチップは長らく、書類上はモデルを安く学習させられる場所でありながら、実際には扱いにくい場所でした。PyTorchのコードを手直ししないと動かなかったからです。Trainium向けのネイティブなPyTorchバックエンドTorchNeuronは、この手間をなくそうとしています。2026年10月20日に公開ベータとなり、pip install torch-neuronxで導入でき、Hugging FaceによるとTransformers、Diffusers、Accelerate、PEFT、TRLがその上で直接動きます。
これまでHugging FaceからAWSの学習用チップTrainiumと推論用チップInferentiaを使うには、PyTorch/XLAのパッケージを土台にしたOptimum Neuronを経由していました。そのためコンパイラーのように考える必要がありました。即時実行ではなく遅延評価のテンソル、静的な形状のみ、グラフが変わるたびに再コンパイル、といった具合です。TorchNeuronはチップを通常のneuronデバイスとして扱えるようにし、eager実行、torch.compile、FSDP2やDTensorといったPyTorch自身の分散ツールに対応します。Hugging Faceは今後、Optimum Neuronを通さず、ライブラリを直接使うよう勧めています。
役に立つのは対応範囲の数字です。TorchNeuronでは、最もよく使われるテキスト生成アーキテクチャ上位100のうち89%がそのまま動き、Optimum Neuron v0.4.5の46%を大きく上回ります。微調整は、Neuron固有の設定をいくつか加えた標準的なTRLのスクリプトで行い、できあがるチェックポイントは普通のTransformersのファイルなので、Hubに公開してどのハードウェアでも読み込めます。詳しい例では、1つのインスタンス上の32のプロセスが、テンソル並列で各モデルを4つのコアに分割し 、残りをFSDPで分散させています。Diffusersでは、FLUX.1-dev、FLUX.2、Qwen-Image-2.1などの画像モデルもテンソル並列に対応し、AWSのNeuron Kernel Interfaceで独自のカーネルを書いてHubで共有できます。
同じ週にGoogleのチップ向けのTorchTPUも出ており、Nvidia以外での学習の費用を見積もる人にとって注目すべき変化です。どちらのバックエンドも、PyTorchやTransformersなど私たちがレビューするライブラリから見て、代替アクセラレーターをごく普通のPyTorchデバイスに見せるため、試すコストが下がります。
ベータ版なので、粗さも見えます。再コンパイルを避けるには形状を静的に保つ必要があり、学習スクリプトでは自動の動的形状を無効にしてバッチを埋めています。Transformersの非同期の重み読み込みにはまだ対応しておらず、89%という対応率は人気のあるアーキテクチャの一部がまだ動かないことを意味します。記事にはGPUとの速度やコストの比較がなく、それこそ多くのチームが必要とする数字です。移行を計画する前に、代表的な作業で試すべきです。