Hugging Face、AWS Trainiumチップ向けにPyTorchネイティブサポートを追加

Hugging Faceは、PyTorchモデルをAWS Trainiumハードウェア上で直接実行できるようにする新しいネイティブバックエンド「TorchNeuron」を発表した。2026年10月20日にパブリックベータとしてリリースされたこのアップデートにより、専門的なコンパイルワークフローが不要になり、開発者はAmazonのAIアクセラレーターでのトレーニングおよび推論において標準的なPyTorch APIを使用できる。

何が起きたのか

過去数年間、Hugging FaceモデルをAWS TrainiumおよびInferentiaチップ上で実行するには、Optimum Neuronを使用する必要があった。このライブラリはPyTorchとハードウェア間のブリッジとして機能し、PyTorch XLA上に構築されたtorch-neuronxパッケージに依存していた。機能面では問題なかったものの、このアプローチは開発者にコンパイラ中心の思考様式への適応を強いた。彼らは即時実行ではなく遅延テンソルに対応し、静的形状（static shapes）に制限され、計算グラフが変化するたびに頻繁な再コンパイルを管理しなければならなかった。Optimum Neuronは、ユーザーからこの複雑さを隠蔽するために、特定のモデリングクラスやトレーナーを含むカスタムコードを維持してきた。

TorchNeuronは、PyTorch自体に直接統合することでこのアーキテクチャを変更する。これは、開発者がCUDA GPUやCPUデバイスとやり取りするのと同様に、neuronをネイティブデバイスとして公開する。この変化により、操作が呼び出された瞬間に即座に実行されるイージス実行が可能になる。以前のように後でコンパイルのためにキューイングされる必要はない。また、Fully Sharded Data Parallel (FSDP) やTensor Parallelism (TP) などのネイティブ分散APIに加え、DTensorやFSDP2といったPyTorch独自の分散プリミティブもサポートしている。さらに、torch.compileとのシームレスな連携により、標準的なPyTorchワークフローから離れることなくパフォーマンス最適化を行うことが可能だ。

その結果、Transformers、Diffusers、Accelerate、PEFT、TRLなど主要なHugging FaceライブラリがTrainium上で直接動作するようになった。開発者は他のどのアクセラレーター上でも使用するのと同じクラスやスクリプトを利用できる。Hugging Faceは現在、ユーザーに対してOptimum Neuronからの移行を推奨しており、これらのライブラリを新しいバックエンドとともに直接使用することを提案している。このアップデートは、開発者体験を簡素化し、AWSインフラストラクチャーを使用している人々がHugging Face Hub上でモデルを共有しやすくすることを目的としている。

仕組み

TorchNeuronの核心メカニズムは、ネイティブPyTorchバックエンドとしての統合にある。XLAのような外部コンパイラレイヤーを通じてPyTorch操作を変換する代わりに、TorchNeuronはPyTorchがneuronを第一級デバイスとして認識することを可能にする。これにより、各操作がハードウェア上で即座に実行されるイージスモードの実行が可能になる。これは、実行を開始する前に完全な計算グラフを構築する必要があった以前の遅延評価モデルとは対照的である。

ネイティブ分散プリミティブをサポートすることにより、TorchNeuronは複数のチップ間でモデルをスケールアップするためのPyTorch組み込みツールを活用する。FSDP2のような機能により、モデルパラメーターをデバイス間でシャーディングすることで、トレーニング中のメモリ使用効率が高まる。推論においては、バックエンドはテンソル並列処理をサポートしており、モデル層を複数のアクセラレーターに分割して処理速度を向上させる。また、システムは生成タスクにおける静的キャッシュの使用を奨励しており、再コンパイルオーバーヘッドを最小限に抑え、イージス実行の柔軟性を保ちつつ高いパフォーマンスを確保する。

主要な詳細

• TorchNeuronは2026年10月20日からパブリックベータとして利用可能になる。
• 開発者はコマンド pip install torch-neuronx を使用してパッケージをインストールできる。
• 新しいバックエンドは、イージス実行、torch.compile、そしてFSDPやTPなどのネイティブ分散APIをサポートする。
• Transformersにおいて、最も使用されている上位100のテキスト生成モデルアーキテクチャのうち89%がそのままサポートされており、Optimum Neuron v0.4.5時の46%から大幅に増加した。
• ユーザーはコード内で device_map=neuron 引数を渡すことで、モデルをNeuronデバイスへ移動できる。
• Hugging Faceは、Optimum Neuronを経由するのではなく、TransformersやDiffusersなどのライブラリを直接使用することを推奨している。

なぜ重要なのか

このアップデートは、AWS Trainiumを使用したいエンジニアにとって参入障壁を大幅に下げる。従来は、コンパイラのように考え、静的形状を管理する必要があり、学習曲線が急峻だった。Trainiumの体験を標準的なPyTorchワークフローに整合させることで、Hugging Faceはチームが既存のコードと専門知識を再利用できるようにする。これにより、モデルをAWSハードウェアへ移植するために必要なエンジニアリング努力が軽減され、カスタムコンパイルステップに関連するバグのリスクも最小限に抑えられる。

サポートされるモデルアーキテクチャの増加は、本番ワークロードにおけるTrainiumの有用性も拡大する。トップクラスのテキスト生成モデルの89%がネイティブにサポートされているため、組織はカスタム統合を待たずに、より幅広い最先端モデルを展開できる。他のアクセラレーターとの同等性は、ソフトウェア互換性の制約よりもコストと可用性によってハードウェア選択が決定されることを意味する。

できること

• 環境内で pip install torch-neuronx を実行して新しいバックエンドをインストールする。
• Transformersでモデルを読み込む際、推論スクリプトを更新して device_map=neuron を使用する。
• Optimum Neuronへの依存関係を、TransformersやDiffusersなどのHugging Faceライブラリへの直接呼び出しに置き換える。
• テキスト生成タスクで静的キャッシュを実装し、頻繁な再コンパイルを防いでパフォーマンスを向上させる。
• Trainiumインスタンス上でFSDP2やDTensorなどのネイティブPyTorchプリミティブを使用して、既存の分散トレーニングセットアップをテストする。
• TorchNeuronが広範なPyTorchエコシステムとどのように統合されるかに関する追加の技術的詳細については、PyTorchブログを確認する。