Hugging Faceは、PyTorchの新しいバックエンドであるTorchNeuronを通じて、AWSのチップTrainiumのネイティブサポートを発表しました。Transformers、Diffusers、Kernelsがこのチップ上で直接実行できるようになり、Trainium上でのモデル実行とHugging Face Hubでの共有が容易になると見込まれています。情報筋によると、パブリックベータ版は2026年10月20日から利用可能となる予定です。
TorchNeuronは、機械学習モデル開発のデファクトスタンダードであるPyTorchの拡張機能です。これにより、PyTorchにTrainium専用のバックエンドが追加され、 »eager »実行モード、torch.compile（コード最適化ツール）、およびFSDPやTPなどのネイティブ分散APIがサポートされます。研究者や開発者は、pip install torch-neuronxコマンドでインストールできます。
Hugging Faceライブラリのアプローチ変更
これまで、AWSが設計した2つのチップ、Trainium（モデルの経済的なトレーニング用）とInferentia（大規模推論用）の使用は、optimum-neuronライブラリを介して行われていました。これはtorch-neuronxパッケージに依存し、さらにPyTorch XLAに基づいていたため、「コンパイラのように」考える必要がありました：遅延テンソル、静的シェープのみ、グラフの切断ごとの再コンパイルなどです。Hugging Faceは、optimum-neuronがこの複雑さを、特別にメンテナンスされたコードで隠蔽していたと説明しています。
TorchNeuronは、PyTorchに直接統合されることでこのアプローチを変更し、PyTorchがNeuronをネイティブデバイスとして認識するようにします。これにより、TrainiumはTransformers、Diffusers、Accelerate、PEFT、TRLにとって他のアクセラレータと同様になり、同じクラスとスクリプトが使用可能になります。Hugging Faceは現在、NeuronユーザーにOptimum Neuronではなくこれらのライブラリを使用することを推奨しています。
テキスト生成に関しては、サポート範囲は広範です：Hugging Faceによると、最も使用されている100のテキスト生成モデルアーキテクチャのうち89%がすぐに互換性があり、Optimum Neuron v0.4.5の46%と比較して向上しています。モデルをNeuronデバイスに移動させるには、device_map=neuron引数を指定するだけです。
トレーニング、拡散、計算カーネル
トレーニングについては、Neuron固有のいくつかの調整を加えることで、標準的なファインチューニングスクリプト（TRL使用）がTrainiumインスタンス上で動作します。Hugging Faceは、公開された例がQwen/Qwen3-8BとデータセットHuggingFaceH4/ultrachat_200kを使用し、8個のTrainium2チップ、つまり32個のNeuronCores（trn2.48xlargeインスタンスの半分）上で実行されたと説明しています。生成されるチェックポイントは標準的なTransformersチェックポイントであり、他のハードウェアでも使用可能です。
Diffusersの複数のパイプラインが、 »eager »モードとコンパイルモードの両方でTorchNeuronと互換性があります。FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1、H3はテンソル並列処理もサポートしており、モデルの各レイヤーを複数のコアに分散します。Hugging Faceは、他のパイプラインも順次互換性を持つようになり、このサポートを自分で追加するためのガイドが利用可能であると述べています。
最後に、Kernelsライブラリには、Neuron Kernel Interface（NKI）で記述された、Neuron向けに最適化された計算カーネルが含まれる予定です。Hubで公開されているビジョントランスフォーマー向けカーネルは、Trainium上で「最大2.4倍」の高速化を実現すると報告されています。Hugging Faceは、TorchNeuronの開発とコミュニティからのフィードバックへの迅速な対応に対するAWSチームの努力に感謝の意を表しています。
