2026.10.28 AI

TorchNeuronとは？TrainiumでPyTorchを動かす導入手順とXLA版からの移行【2026年10月】

Q: GPU向けに書いたPyTorchコードはそのまま動きますか？

.to("cuda") を .to("neuron") に替えれば動く範囲は広い一方、CUDA Graphs、NCCL固有の集合通信、動的形状は動きません。速度を出すには、torch.compileで包んで入力の形を固定する手直しが前提です。

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の世代別性能と料金を扱った記事にまとめています。

まとめ：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発表によると、TorchNeuronはPyTorch Conference North America（10月20日）に合わせてオープンソース化され、同日から公開ベータとしてpip installで入れられるようになりました。ライセンスはApache 2.0、ソースはGitHubのaws-neuron/torch-neuronxです。

対応チップはTrainium2とTrainium3。Neuron SDKのリリースノート一覧では、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概要に掲載済みです。PyTorchのeagerとcompileの基本はPyTorchの仕組みと2.14の入れ方の記事で扱っています。

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のリリースノートは、ネイティブ版がtorch_xlaとxm.xla_device()の代わりにneuronデバイスを使うため、XLA版向けのコードは書き換えが要ると明記しています。パッケージ名が同じなので、pipの版指定を間違えると意図しない方が入ります。SDK全体の版の固定方法はNeuron SDKの版ごとの変更点と固定方法の記事を参照してください。

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の料金と調達方法の章で確認してください。

Ubuntu 26.04でaptとpipから入れるコマンドと24.04の落とし穴

Neuron公式の手動インストール手順から、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です。続いて、公式クイックスタートと同じ大きさの多層パーセプトロンで、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の技術解説では、小さな部分グラフをためてコンパイルし、実行と並行させるオプトインの方式と説明されています。性能値は公表されていないため、本番の前提にはしません。

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）でカーネルを書きます。@nki.jitで書いたカーネルをtorch_neuronx.nki_opで包むとPyTorchのカスタム演算になり、register_autogradで逆伝播も登録できます。呼び出し元はeagerとcompileのどちらでも対応可能です。

自分で書く前に、NKI-libを確かめます。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のできることと使い方の記事で解説しています。

torchtitan-neuronでLlama 3 8Bを学習する公式構成

大規模学習のライブラリTorchTitanは、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を公開しました。ソースはgoogle-pytorch/torch_tpuです。

項目	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の記事で整理しています。

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受託開発でご相談を受けています。

よくある質問

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の使い分けと推論の始め方の記事のとおりです。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の仕組みと使い方の記事を参照してください。