TorchNeuronがAWS TrainiumでHFモデル実行を簡素化

AWSとHugging Faceは、AWS Trainium上でPyTorchモデルをより容易に実行するための新バックエンド「TorchNeuron」を発表した。TorchNeuronはNeuronデバイスをPyTorchのネイティブデバイスとして扱えるようにするプラグインで、イーガー実行、torch.compile、DTensor、FSDP2、Tensor Parallelismといった標準的な分散処理機能を直接活用できる。従来はoptimum-neuronとtorch-neuronx、PyTorch XLAを介し、遅延実行、静的形状、グラフ中断時の再コンパイルといった制約を抽象化する設計が必要だったが、今回の統合によりTrainium向けの開発体験が大きく簡素化される。 TorchNeuronの公開ベータ版は2026年10月20日から提供される予定で、torch-neuronxをpip経由でインストールして使える。Hugging Faceのエコシステムでは、Transformers、Diffusers、Accelerate、PEFT、TRLなどのライブラリがTrainiumを他のアクセラレータと同様に扱える。ユーザーは既存のクラスや学習・推論スクリプトをほぼそのまま実行でき、従来推奨されていたOptimum Neuronを介さず直接ライブラリを使う形が推奨される。 推論では、Transformersのテキスト生成対応が大きく改善された。よく利用されるトップ100のテキスト生成モデルアーキテクチャのうち89%がサポートされ、Optimum Neuron v0.4.5の46%から大幅に拡大した。モデル実行時にneuronデバイスを指定するだけで、Neuron上へモデルを配置でき、Tensor Parallelismや静的キャッシュを活用してNeuronインスタンスの性能を引き出せる。 学習でも、TRLのSFTTrainerなどの標準スクリプトがTrainiumで動作する。Tensor ParallelismとFSDPのサイズをDistributedConfigに設定すれば、モデルの重み読み込み時に分散配置が行われ、専用のTrainerやモデルコードのforkは不要になる。8枚のTrainium2チップで32個のNeuronCoreを使い、4x8のTPとFSDP構成で教師あり微調整を完遂した例が示された。保存されたチェックポイントも通常のTransformers形式であり、Hugging Face Hubへそのまま公開し、他のハードウェアで読み込める。 Diffusersでも多くのパイプラインがTorchNeuronに対応し、イーガー実行とコンパイル実行の両方が可能になる。FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1、H3などではTensor Parallelismが使え、Neuronデバイス全体を活かした画像生成が可能になる。加えて、Neuron Kernel Interfaceで記述されたカーネルもHugging Face Hubで共有され、視覚モデルで最大約2.4倍の高速化を期待できる。両社は今回の発表により、Trainiumの使いやすさ向上と、オープンなモデル共有エコシステムの拡大を加速させる構えだ。