--- title: "TorchNeuronとは？TrainiumでPyTorchを動かす導入手順とXLA版からの移行【2026年10月】" url: "https://www.issoh.co.jp/tech/details/19027/" published: 2026-10-28 updated: 2026-10-28 categories: ["AI"] publisher: "株式会社一創" --- # TorchNeuronとは？TrainiumでPyTorchを動かす導入手順とXLA版からの移行【2026年10月】 TorchNeuronは、AWSの機械学習チップTrainiumを、PyTorchの標準デバイス`neuron`として扱えるようにするネイティブバックエンドです。2026年10月にNeuron SDK 2.33.0の公開ベータとして出て、ソースはApache 2.0で公開されました。この記事では、従来のPyTorch/XLA版torch-neuronxとの違い、Ubuntu 26.04での導入コマンド、eagerとtorch.compileの使い分け、FSDP2やHugging Face Transformersでの使い方、初版で動かない機能とTorchTPUとの違いを、公式ドキュメントの記載をもとに整理します。チップの世代や料金は[AWS Trainiumの世代別性能と料金を扱った記事](https://www.issoh.co.jp/tech/details/10448/)にまとめています。 ## まとめ：TorchNeuronを今試すチームと2.9系のXLA版に留まるチーム Trn2かTrn3を使い、PyTorch 2.13か2.14で新しく学習や推論のコードを書くなら、TorchNeuronから始めます。`.to("neuron")`でモデルを載せ、性能が要る部分だけ`torch.compile(model, backend="neuron")`で包む書き方です。PyTorch/XLA特有の`xm.mark_step()`は要りません。 本番の学習ジョブをすでにtorch-neuronx 2.9系で回しているなら、急いで移しません。公開ベータの初版は動的形状に対応せず、`torch.sparse`や`torch.func`も動かないからです。Inf2やTrn1も対象外です。 導入で最初に引っかかるのはOSです。リリースノートでは、単体のwheelはUbuntu 26.04向けだけで、24.04ではtorch-neuronx入りのDLAMIを使うよう案内されています。 ## TorchNeuronの仕組みとPyTorch/XLA版との違い 名前は同じtorch-neuronxでも、中身は別物です。まず構造の差を押さえます。 ### 2026年10月20日の公開ベータとApache 2.0で公開されたリポジトリ [PyTorch公式ブログのTorchNeuron発表](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)によると、TorchNeuronはPyTorch Conference North America（10月20日）に合わせてオープンソース化され、同日から公開ベータとして`pip install`で入れられるようになりました。ライセンスはApache 2.0、ソースは[GitHubのaws-neuron/torch-neuronx](https://github.com/aws-neuron/torch-neuronx)です。 対応チップはTrainium2とTrainium3。[Neuron SDKのリリースノート一覧](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/index.html)では、2.33.0の公開日が10月19日で、TorchNeuron（Beta）の版は2.14.3.0と2.13.3.0です。 ### PrivateUse1登録でneuronデバイスとして見える仕組みと処理の流れ TorchNeuronは、PyTorchが独自ハードウェア向けに用意した拡張口PrivateUse1にTrainiumを登録します。そのため`torch.nn`、`torch.autocast`、`torch.distributed`、`torch.profiler`が`neuron`デバイスの上でそのまま動く仕組みです。デバイスを選ばない書き方の`torch.accelerator`で書いたコードは、変更なしで動くとされています。 `torch.compile`を使う場合は、TorchDynamoがモデルをFXグラフに取り込み、AOTAutogradが逆伝播を先に追跡します。Neuronバックエンドはそれを`Torch-MLIR`経由でStableHLOに落とし、Neuronコンパイラが実行形式のNEFFファイルを作ります。構成図は[Neuron公式のTorchNeuron概要](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html)に掲載済みです。PyTorchのeagerとcompileの基本は[PyTorchの仕組みと2.14の入れ方の記事](https://www.issoh.co.jp/tech/details/1315/)で扱っています。 ### xm.xla\_device()を使う2.9系torch-neuronxとの書き方の差 旧来のtorch-neuronxはPyTorch/XLAを経由していました。違いを表にします。 | 項目 | TorchNeuron | XLA版torch-neuronx | | --------- | ----------- | ----------------- | | デバイス指定 | neuron | xm.xla\_device() | | 対応PyTorch | 2.13・2.14 | 2.9まで | | 最新版 | 2.14.3.0 | 2.9.0.2.15.32035 | | 対応チップ | Trn2・Trn3 | Inf2・Trn1・Trn2など | | 実行の区切り | 不要 | xm.mark\_step() | [TorchNeuronのリリースノート](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)は、ネイティブ版が`torch_xla`と`xm.xla_device()`の代わりに`neuron`デバイスを使うため、XLA版向けのコードは書き換えが要ると明記しています。パッケージ名が同じなので、`pip`の版指定を間違えると意図しない方が入ります。SDK全体の版の固定方法は[Neuron SDKの版ごとの変更点と固定方法の記事](https://www.issoh.co.jp/tech/details/10521/)を参照してください。 ## TorchNeuronの導入手順｜Trn2でpip installから動作確認まで ここではTrn2インスタンスにUbuntu 26.04を入れた前提で進めます。 ### 対応インスタンスTrn2・Trn3とPyTorch 2.13・2.14の版の組み合わせ 2.33.0時点の対応範囲をリリースノートから抜き出します。 | 項目 | 2.33.0時点の対応 | | ------------- | ----------------- | | インスタンス | Trn2・Trn3 | | PyTorch | 2.13.0・2.14.0 | | torch-neuronx | 2.13.3.0・2.14.3.0 | | Python | 3.12・3.13・3.14 | | 単体wheelのOS | Ubuntu 26.04のみ | Inf2とTrn1は対象外です。Trn2・Trn3の提供リージョンや時間単価は[Trainiumの料金と調達方法](https://www.issoh.co.jp/tech/details/10448/)の章で確認してください。 ### Ubuntu 26.04でaptとpipから入れるコマンドと24.04の落とし穴 [Neuron公式の手動インストール手順](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/setup/pytorch/manual.html)から、Ubuntu 26.04の分を1つにまとめました。ドライバ・ランタイム・ツールをaptで入れ、TorchNeuronは仮想環境へpipで入れます。 ``` # Neuron の apt リポジトリを登録（Ubuntu 26.04） . /etc/os-release sudo install -m 0755 -d /etc/apt/keyrings wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB \ | sudo gpg --batch --yes --dearmor -o /etc/apt/keyrings/neuron.gpg echo "deb [signed-by=/etc/apt/keyrings/neuron.gpg] https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main" \ | sudo tee /etc/apt/sources.list.d/neuron.list sudo apt-get update -y # ドライバ・集合通信・ランタイム・ツール sudo apt-get install -y linux-headers-$(uname -r) \ aws-neuronx-dkms=2.* aws-neuronx-collectives=2.* \ aws-neuronx-runtime-lib=2.* aws-neuronx-tools=2.* export PATH=/opt/aws/neuron/bin:$PATH # TorchNeuron（PyTorch 2.14）を仮想環境へ sudo apt-get install -y python3-venv g++ python3 -m venv aws_neuron_venv_pytorch_2_14 source aws_neuron_venv_pytorch_2_14/bin/activate python -m pip install --upgrade pip python -m pip install "torch-neuronx==2.14.*" "neuronx-cc==2.*" nki \ --extra-index-url https://pip.repos.neuron.amazonaws.com \ --extra-index-url https://download.pytorch.org/whl/cpu ``` 複数インスタンスで学習する場合は、別途EFAのインストーラも入れます。 注意点は24.04です。インストール手順のページには24.04向けのコマンドも載っていますが、リリースノートは「このリリースの単体wheelはUbuntu 26.04のみ」と書き、24.04ではtorch-neuronx入りの単一フレームワークDLAMIを使うか、ソースからビルドするよう案内しています。24.04の手元環境でpipが版を解決できないときは、手順の誤りではなくこの制約を疑ってください。22.04向けのDLAMIとコンテナはもう新しく出ません。 ### neuron-lsと.to(“neuron”)で最初のモデルを動かす確認コード 導入後は、デバイスの認識と版を確かめてから小さなモデルを流します。 ``` neuron-ls python3 -c "import torch, torch_neuronx; print(f'PyTorch {torch.__version__}, torch-neuronx {torch_neuronx.__version__}')" ``` 手順書の出力例は`PyTorch 2.14.0+cpu, torch-neuronx 2.14.3.0`です。続いて、[公式クイックスタート](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-build-first-model.html)と同じ大きさの多層パーセプトロンで、eagerとcompileの両方を通します。 ``` import torch import torch.nn as nn class TinyMLP(nn.Module): def __init__(self, d_in=128, d_hidden=256, d_out=10): super().__init__() self.net = nn.Sequential(nn.Linear(d_in, d_hidden), nn.ReLU(), nn.Linear(d_hidden, d_out)) def forward(self, x): return self.net(x) device = "neuron" model = TinyMLP().to(device) x = torch.randn(8, 128, device=device) # eager：演算を1つずつ実行（初回は演算ごとのコンパイルが入る） with torch.no_grad(): y = model(x) print("device:", y.device, "output shape:", tuple(y.shape)) # torch.compile：初回呼び出しでグラフ単位にコンパイルし、NEFF をディスクに保存 compiled = torch.compile(model, backend="neuron", dynamic=False) with torch.no_grad(): y2 = compiled(x) print("max diff:", (y - y2).abs().max().item()) ``` クイックスタートでは、`device: neuron:0`と出力形状`(8, 10)`が表示されれば準備完了としています。初回の実行が遅いのはコンパイルのためで、性能の問題ではないと同ページに書かれています。 ## eager・torch.compile・NKIの使い分けと初回コンパイルの待ち時間 実行方式は3つあります。開発の段階に合わせて切り替えるのが前提の設計です。 ### eagerで未対応演算がCPUへ逃げる挙動とadaptive eagerの位置づけ eagerは演算を呼んだ時点で実行するので、Pythonの分岐や`print`、ブレークポイントがそのまま使えます。研究とデバッグ向けの方式です。リリースノートによると、未対応の演算はeagerではCPUで実行されることがあり、`float64`は警告付きで`float32`として計算されます。CPUへの退避は黙って遅くなる原因になるので、速度が出ないときはまずここを疑います。 発表では、演算をまとめて起動回数を減らすadaptive eager（実験的）も紹介されました。[Amazon Scienceの技術解説](https://www.amazon.science/blog/the-science-behind-torchneuron-amazons-new-pytorch-backend-for-aws-trainium-chips)では、小さな部分グラフをためてコンパイルし、実行と並行させるオプトインの方式と説明されています。性能値は公表されていないため、本番の前提にはしません。 ### torch.compileのdynamic=False指定とNEFFキャッシュの再利用 本番で速度を出すのはtorch.compileです。モデル全体でも一部だけでも包めて、FlexAttentionにも対応します。制約は2つ。`mode="default"`しか使えず、CUDA向けの`max-autotune`は不可です。もう1つは静的形状で、`dynamic=False`が必要になり、入力の形が変わるたびに再コンパイルが走ります。 ``` # 学習ジョブの先頭で一度だけ設定する export TORCH_NEURONX_NEFF_CACHE_DIR=/mnt/shared/neff_cache # 既定は /tmp/neff_cache # Python 側：既定のバックエンドを neuron にしてから compile する torch.compiler.set_default_backend("neuron") model = MyModel().to("neuron") model.compile() # 以降は backend 引数なしで neuron が使われる ``` コンパイル済みのNEFFは`TORCH_NEURONX_NEFF_CACHE_DIR`に保存され、プロセス間や分散ランク間で使い回されます。Amazon Scienceの解説では、64プロセスで動かすLLMで、共有コンパイルの仕組みにより最初の学習ステップまでの時間が約3分の1縮んだという説明です。系列長がばらつくデータは、いくつかの長さに切りそろえてパディングし、形の種類を絞るとコンパイル回数を抑えられます。 ### NKIカーネルをnki\_opでカスタム演算として登録する手順 標準の演算の組み合わせで足りない部分は、[Neuron Kernel Interface（NKI）](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/index.html)でカーネルを書きます。`@nki.jit`で書いたカーネルを`torch_neuronx.nki_op`で包むとPyTorchのカスタム演算になり、`register_autograd`で逆伝播も登録できます。呼び出し元はeagerとcompileのどちらでも対応可能です。 自分で書く前に、[NKI-lib](https://github.com/aws-neuron/nki-library)を確かめます。FlashAttentionやMoEのブロック単位カーネルなど、LLMとマルチモーダル向けの演算がオープンソースで揃う構成です。発表の事例でも、Cartesiaは小さな演算を標準経路に任せ、繰り返し呼ばれる層だけをNKIで書いたと述べています。全層を手書きする必要はありません。 ## 分散学習FSDP2とTransformers v5のdevice\_mapでの使い方 既存の学習コードやモデル定義を、どこまでそのまま持ち込めるかを見ます。 ### neuron c10dバックエンドで動くDDP・FSDP2・テンソル並列の範囲 TorchNeuronは`torch.distributed`向けに`neuron`というc10dバックエンドを持ちます。リリースノートではDDP、FSDP2、DTensorとDeviceMesh、パイプライン並列、コンテキスト並列が対応済みです。発表ではテンソル並列とZeROも挙がっています。 `all_gather`や`all_reduce`などの集合通信は、ノード内ではTrainiumの集合通信専用コア、ノード間ではEFAで流れます。プロセスグループの初期化でバックエンド名をNCCLから`neuron`に替えるのが主な修正点です。NCCL固有の集合通信やCUDA Graphsに頼る処理は対象外なので、移植前にコードを検索して洗い出します。 ### device\_map指定でHugging FaceモデルをTrainiumへ載せる手順 Hugging Face側では、Transformers・Diffusers・TRL・kernelsの各ライブラリにTrainiumが対応デバイスとして入りました。Transformers v5なら`device_map`で読み込み先を指定します。 ``` from transformers import AutoModelForCausalLM, AutoTokenizer model_id = "Qwen/Qwen3-0.6B" tok = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained( model_id, device_map="neuron", # "auto" でも neuron が自動検出される ) print(model.device) # neuron:0 ``` 発表によると`device_map="auto"`でも検出され、以降のスクリプトは変更不要です。ただし文章生成は1トークンごとに系列長が伸びるため、eagerでは形が変わるたびにコンパイルが入ります。高い処理量で配信するなら、TorchNeuronのtorch.compileの上に作られたvLLM Neuron（2.33.0時点で0.29.0.1.2.0、Beta）を使う構成です。Hugging Faceのライブラリ構成は[Hugging Faceのできることと使い方の記事](https://www.issoh.co.jp/tech/details/6108/)で解説しています。 ### torchtitan-neuronでLlama 3 8Bを学習する公式構成 大規模学習のライブラリTorchTitanは、[aws-neuron/torchtitan-neuron](https://github.com/aws-neuron/torchtitan-neuron)で先行して動きます。NKIカーネルやTrainium向けのモデル設定が入り、2.33.0時点の版は0.2.2+neuron.1です。 発表の例は、trn2.48xlarge 1台で`NEURON_RT_NUM_CORES=64`を指定し、`torchrun`で64プロセスを立てる構成です。SimpleFSDPとtorch.compileでLlama 3 8Bを学習し、`--parallelism.tensor-parallel-degree 4`と`--compile.backend neuron`を渡しています。 ## 公開ベータで未対応の機能とTorchTPUとの比較で見る採用判断 「1行で動く」の裏にある制約を確かめてから、移すかどうかを決めます。 ### 動的形状・torch.sparse・TorchAO非対応など初版の制約一覧 リリースノートと発表に書かれた、2.33.0の初版で動かないものを並べます。影響の大きい順です。 - 動的形状とシンボリック形状（開発中。torch.compileは`dynamic=False`が必須） - `torch.sparse`・`torch.func`・`torch.nested`・`torch.distributions` - TorchAOとMXFP8（発表ではFP8を含めプライベートベータで、次のリリースで一般提供予定） - 標準のマルチヘッドアテンションでの任意マスク（マスクなしと因果マスクは動く） - 前向きモードの自動微分、チャネル単位の量子化、双線形・双三次のアップサンプリング - `torch.fft`のうち実数から複素数への変換以外（CPUで実行） 可変長の入力をそのまま流す推論や、FP8で学習コストを下げる計画は、今の版では成り立ちません。前者はパディングで形を固定する設計が、後者は次の版を待つ判断が必要になります。 ### 上流テスト15.4万件と39,000件のNeuron固有スキップの読み方 発表では、PyTorch本家のテスト（ATen・Dynamo・分散・プロファイラ）を2026年8月に流した結果が出ています。約15.4万件の結果のうち約6.9万件がスキップで、約3万件はCUDA専用などPyTorch側の判定によるもの、残る約3.9万件がNeuron側の意図的なスキップです。最大の内訳は、Neuronにネイティブな変換がまだ無い演算の約2.2万件でした。 合格率そのものは示されていません。顧客事例も読み方に注意が要ります。Deepgramの「既存GPU比でドルあたり約60%多くの音声秒数を処理」は、AWSのGenerative AI Innovation Centerと数か月かけて移植した後の数字で、比較したGPUの機種も書かれていません。自社のモデルで同じ結果が出るかは、移植して測るまで分かりません。 ### 同時期にOSS化されたGoogleのTorchTPUとの実行モードの違い 同じPyTorch Conferenceの時期に、GoogleとMetaも[TPU向けのネイティブバックエンドTorchTPU](https://pytorch.org/blog/introducing-torchtpu/)を公開しました。ソースは[google-pytorch/torch\_tpu](https://github.com/google-pytorch/torch%5Ftpu)です。 | 項目 | TorchNeuron | TorchTPU | | --------------- | ------------------- | -------------------- | | デバイス名 | neuron | tpu | | compileのbackend | neuron | tpu | | eagerの既定 | 演算ごとに実行 | Strict（演算ごと） | | 融合するeager | adaptive eager（実験的） | DeferAndFuse | | 動的形状 | 開発中 | bounded dynamismを開発中 | どちらも、まずeagerで動かし、速度はcompileと演算の融合で取りにいく設計です。TorchTPUのDeferAndFuseは環境変数`TORCH_TPU_DEFER_AND_FUSE=1`で有効になり、Strictより50%以上速くなる場合があるとGoogleは書いています。どちらのチップを選ぶかはバックエンドではなく、使っているクラウドと調達できる台数が判断の基準です。TPU側の世代と料金は[TPUとTorchTPUの記事](https://www.issoh.co.jp/tech/details/10447/)で整理しています。 ### TorchNeuronを採用する条件と2.9系XLA版に留める判断の線 TorchNeuronを採用するのは、Trn2・Trn3で新しくコードを書き、入力の形を固定でき、PyTorch 2.13以降の機能を使いたい場合です。GPU向けに`torch.accelerator`で書かれたコードや、Transformers v5に載るモデルは、移植の手間が最も小さく済みます。 見送るのは次の場合です。Inf2やTrn1で動かしている推論、torch-neuronx 2.9系で安定稼働している本番の学習、可変長の入力やFP8に依存する処理。これらは公開ベータの間はXLA版に留め、TorchNeuronは別のインスタンスで評価だけ進めます。XLA版のコードは`xm`系の呼び出しを1つのモジュールに閉じ込めておくと、移行時の書き換え範囲を絞ることが可能です。GPUからTrainiumへの移植評価や、学習・推論基盤を含むモデル開発は[生成AI開発・AI受託開発](https://www.issoh.co.jp/service/ai/development/)でご相談を受けています。 ## よくある質問 TorchNeuronについて検索で多い質問に、公式ドキュメントの記載で答えます。 ### TorchNeuronとtorch-neuronxは同じものですか？ パッケージ名は同じtorch-neuronxですが、版で中身が分かれます。2.13.3.0と2.14.3.0がネイティブバックエンドのTorchNeuronで、使用するのは`neuron`デバイスです。2.9.0.2.15.32035までがPyTorch/XLAを経由する旧来版で、`xm.xla_device()`を使います。両者のコードは互換ではないため、`pip install`では`"torch-neuronx==2.14.*"`のように版を明示してください。 ### TorchNeuronは無料で使えますか？ ソフトウェアはApache 2.0のオープンソースで、ライセンス費用はかかりません。費用が発生するのはTrn2・Trn3インスタンスなど、AWS側の利用料です。コンパイルの待ち時間も課金対象になるため、NEFFキャッシュの共有で再コンパイルを減らすと費用も下がります。 ### TorchNeuronはInferentia（Inf2）で使えますか？ 2.33.0時点では使えません。対応はTrn2とTrn3だけです。Inf2で推論するなら、従来どおりPyTorch/XLA版のtorch-neuronxで`torch_neuronx.trace`を使ってモデルをコンパイルします。Inf2での具体的な手順は[InferentiaとTrainiumの使い分けと推論の始め方の記事](https://www.issoh.co.jp/tech/details/4647/)のとおりです。Trn2・Trn3は学習と推論の両方に使えます。 ### TorchNeuronでvLLMは使えますか？ 使えます。vLLM NeuronはvLLMのプラグインで、TorchNeuronのtorch.compile対応の上に作られています。2.33.0時点の版は0.29.0.1.2.0でBetaです。画像や動画の生成モデルを配信するvLLM Omni Neuron（0.29.0.0.2.0、Beta）もあります。vLLMそのものの仕組みは[vLLMの仕組みと使い方の記事](https://www.issoh.co.jp/tech/details/15786/)を参照してください。 ### GPU向けに書いたPyTorchコードはそのまま動きますか？ `.to("cuda")`を`.to("neuron")`に替えれば動く範囲は広い一方、CUDA Graphs、NCCL固有の集合通信、動的形状は動きません。速度を出すには、torch.compileで包んで入力の形を固定する手直しが前提です。 ## 関連記事 - [AWS Trainiumとは？Trainium3までの世代別性能・料金・Neuron SDK対応と採用判断](https://www.issoh.co.jp/tech/details/10448/)：TorchNeuronが動くTrn2・Trn3の性能と料金です。 - [AWS Neuron SDKとは？最新バージョン2.32の変更点・インストール手順と版の選び方](https://www.issoh.co.jp/tech/details/10521/)：XLA版を含むSDK全体の版の固定方法です。 - [AWS InferentiaとTrainiumの違い｜Inf2とTrn2の使い分けと推論の始め方](https://www.issoh.co.jp/tech/details/4647/)：TorchNeuron対象外のInf2で推論する手順です。 - [TPUとは？GPUとの違い・世代別性能と料金、TorchTPUのOSS公開とPyTorch移行【2026年10月】](https://www.issoh.co.jp/tech/details/10447/)：同時期に公開されたTorchTPUの使い方です。 - [vLLMとは？PagedAttentionの仕組み・使い方とOllama・TensorRT-LLMとの違いを実装者目線で解説](https://www.issoh.co.jp/tech/details/15786/)：vLLM Neuronの土台になる推論エンジンです。 --- 出典: [TorchNeuronとは？TrainiumでPyTorchを動かす導入手順とXLA版からの移行【2026年10月】]()（株式会社一創）