世界は立ち止まらない。RSS
BIG CHANGE.

Markdown版

AI-translated from English; not yet reviewed by a fluent editor.

# AWSのTorchNeuron、Trainium向けにPyTorchネイティブバックエンドを導入

> AWSのNeuron 2.33は、Trainium2とTrainium3向けに公開ベータ版のPyTorchバックエンドを導入した。対応バージョン、演算子の対応範囲、コンパイルの制約によって、エンジニアが評価時に確認すべき点が決まる。

By BIG CHANGE Editorial

Published: 2026-10-20T04:20:56.038Z
Updated: 2026-10-20T04:20:56.038Z
Canonical: https://bigchange.ai/blog/aws-torchneuron-native-pytorch-backend-trainium

![Conceptual charcoal illustration of a code-marked monitor beside an open rack drawer containing an unbranded accelerator board.](https://bigchange.ai/api/media/file/torchneuron-native-device-trainium-hero-v1.png)
AI-generated conceptual illustration by BIG CHANGE.

AWSのNeuron 2.33の[リリースノート](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)は、TorchNeuronの公開ベータ版の日付を10月19日としている。同じく10月19日付の[PyTorchの発表](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)では、公開ベータ版とオープンソース版の提供開始は10月20日とされている。`torch-neuronx` パッケージは、Trainium2とTrainium3向けにPyTorchへ `neuron` デバイスを登録し、対応するモデルやテンソルの操作を使い慣れたPyTorch APIで実行できるようにする。これらのチップを検討するエンジニアにとって有用な変化は、既存のPyTorchプログラムからデバイス上での試行までの道のりが短くなることだ。そのプログラムが正しく効率的に動くかどうかは、依然として操作内容、テンソルの形状、ソフトウェア環境に左右される。

従来、TrainiumでPyTorchを使う経路では、XLAベースの `torch-neuronx` ライブラリと `torch_xla` のデバイスコードを利用していた。AWSによると、ネイティブバックエンドはPyTorchの `PrivateUse1` 拡張を使い、`xm.xla_device()` を `neuron` へのデバイス配置に置き換える。パッケージ名はなじみのあるものでも、すでにNeuronを利用するプロジェクトにとっては大きなインターフェース変更だ。このリリースはTrn2とTrn3インスタンス上での学習と推論に適用される。対応表にTrn1は記載されていない。

## PyTorchプログラムで何が変わるのか

AWSの[リリースノート](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)には、モデルを `.to("neuron")` で移動し、入力テンソルを `device="neuron"` 上に作成する例がある。対応する `torch.nn`、autograd、オプティマイザーの操作は、その後バックエンドを通じて実行される。PyTorchのデバイスに依存しない `torch.accelerator` APIで書かれたコードは、デバイス固有の書き直しなしにNeuronを検出できると、[TorchNeuronの概要](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html)で説明されている。

実行経路は2つある。eagerモードでは、Pythonから呼び出されるたびにPyTorchの操作を実行し、通常の制御フロー、出力表示、ブレークポイントを維持する。各操作は初回利用時にコンパイルされ、後続の条件が一致する呼び出しではキャッシュを使える。エンジニアは、モデルまたは選択した領域を `torch.compile(model, backend="neuron", dynamic=False)` でコンパイルすることもできる。ネイティブバックエンドは対応する順伝播・逆伝播のグラフをコンパイルし、取り込んだ領域の外のコードはeagerモードで動作を続ける。[PyTorchの発表](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)には、コンパイル前に `torch.compiler.set_default_backend("neuron")` を一度設定することで、その後のコンパイル呼び出しからバックエンド引数を省略できるという説明もある。

このリリースには `neuron` バックエンドが含まれ、`torch.distributed` 向けの集合通信とポイントツーポイント通信の操作が文書化されている。AWSは対応APIとしてDistributedDataParallel、FSDP2、DTensor、DeviceMesh、パイプライン並列を挙げている。分散学習プロジェクトでは検討する価値がある経路だが、既存の複数ノードのジョブが変更なしに動くことや、特定のスループット目標を満たすことを、リリースノートが証明しているわけではない。

## 評価は整合する環境を用意することから始まる

AWSの[インストールのクイックスタート](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html)では、Trn2またはTrn3インスタンスと、Neuron Base DLAMIまたは対応するUbuntu AMIが必要とされる。このリリースの対応表は、Ubuntu 24.04または26.04上のLinux x86\_64を指定している。Ubuntu 22.04は、このリリースではサポート終了となっている。Python 3.12、3.13、3.14が列挙され、PyTorch 2.13.0には `torch-neuronx` 2.13.3.0、PyTorch 2.14.0には `torch-neuronx` 2.14.3.0を組み合わせる。パッケージ系列はPyTorchのバージョンと一致していなければならない。

クイックスタートでは、NeuronとPyTorchのパッケージインデックスから `torch-neuronx==2.14.*` に加え、`neuronx-cc==2.*` と `nki` をインストールする。その後、デバイスが利用可能かを確認し、Neuron上の小さな行列演算をCPUの結果と比較する。これで検証できるのは、AWSの例におけるインストール済みデバイスへの実行経路だ。読者自身のモデルを検証するものではない。次の評価として妥当なのは、実際のモデルの順伝播・逆伝播を実行し、`torch_neuronx.get_fallback_ops()` で記録されたCPUフォールバックを調べ、現在の環境と出力を比較し、コンパイルと定常状態の実行を別々に測定することだ。これらの後半の手順はBIG CHANGEが提案する評価項目であり、私たちが実施したテストではない。

## ベータ版がなお作業に影響し得る点

eagerモードでは、未対応の操作はCPUにフォールバックし、警告を出すことがある。`torch.compile` では、グラフに取り込まれた操作をNeuronバックエンドが下位の表現へ変換できない場合、デフォルトではコンパイルに失敗する。AWSは、操作ごとの対応範囲について[演算子の互換性表](https://github.com/aws-neuron/torch-neuronx/blob/main/docs/torch_api_compatibility.md)を参照するよう開発者に案内している。したがって、eagerで最後まで実行できたという事実だけでは、すべてがデバイス上で実行されたことや、コンパイル実行との互換性は確認できない。

リリースではさらに、`torch.compile` に `dynamic=False` が必要とされている。入力の形状が変わると、再コンパイルが必要になることがある。AWSは形状の変動を減らすため、パディングやバケット分けを提案している。一致する成果物がキャッシュにない場合、初回実行にはコンパイルが含まれる。ワークロードの時間を測る際には、これらの詳細が重要になる。コンパイルの遅延と入力形状の現れ方によって、評価結果が変わる可能性がある。

ほかの明示的な制約も、互換性の主張を限定している。ベータ版は `torch.sparse`、`torch.func`、`torch.nested`、`torch.distributions` に対応していない。Neuronで実行できるのは `torch.fft` の一部の変換だけで、`float64` は `float32` として実行され、警告が出る。このリリースにはTorchAOとの統合も、複合MXFP8データ型への対応も含まれていない。これらの機能を使うチームは、デバイス配置を移行計画と見なす前に、自分たちのコード経路を確認する必要がある。

### 大きな変化

TorchNeuronは、公開ベータ版でTrainiumをPyTorchのネイティブなデバイス・コンパイルインターフェースの背後に置く。対応するPyTorchプログラムでは最初の試行がより直接的になり、既存のNeuron利用者にはXLAベースのコードとは異なる経路が用意される。採用を判断する決定的な証拠は、ワークロードごとに得る必要がある。演算子が意図したデバイスにとどまるか、数値結果が引き続き許容できるか、そして測定した性能と費用がインスタンスと移行作業に見合うかどうかだ。BIG CHANGEはこの分析のために文書と発表を確認したが、TorchNeuronを実行したり、AWSの性能に関する主張を独立に検証したりはしていない。

## 出典と参考資料

- [AWS Neuron 2.33のTorchNeuronリリースノート](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)：リリース日、ベータの状況、バージョンとOSの対応表、API、既知の制約。
- [AWSによるTorchNeuronの概要](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html)：デバイス統合、eager実行とコンパイル実行、CPUフォールバックの動作、Neuronを経由するソフトウェアの経路。
- [AWSのインストール・検証クイックスタート](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html)：インスタンスとイメージの前提条件、パッケージのインストール、文書化されたデバイス確認。
- [PyTorchブログ：TorchNeuronの紹介](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)：デバイス切り替え、コンパイル、分散APIに関するPyTorch側からの説明。この記事はAWSの寄稿者がPyTorchブログで行った発表であり、独立したベンチマークではない。
- [Amazonの発表](https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai)：目標と初期の採用者についてのAWSの説明。移行、性能、経済性に関する主張は、ワークロードごとに検証する必要がある。

## Sources

- [AWS Neuronリリースノート：TorchNeuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html) — 一次資料。公開ベータ版の日付は2026-10-19で、PyTorchの発表が示す公開ベータ版・オープンソース版の提供開始より1日早い。バージョンの対応表と明示された制約を掲載。
- [TorchNeuronとは何か](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html) — デバイスとバックエンドの利用に関する一次文書。
- [PyTorchブログ：AWS Trainium向けPyTorchネイティブバックエンドTorchNeuron](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/) — AWSの寄稿者によるPyTorch側の発表。10月19日付で、公開ベータ版とオープンソース版の提供開始を10月20日としている。実際のURLはAmazonの記事からたどったもので、独立した性能検証ではない。
- [Amazonの発表](https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai) — ベンダーの発表。移行と経済性に関する主張はAWSの説明であり、独立に検証されたものではない。
- [torch-neuronxのインストールと検証](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html) — 前提条件、対応するパッケージのインストール、文書化されたCPU/Neuronデバイス確認についてのAWSクイックスタート。
- [TorchNeuron API互換性表](https://github.com/aws-neuron/torch-neuronx/blob/main/docs/torch_api_compatibility.md) — 保守担当者による、APIごとのネイティブ実行、CPUフォールバック、未対応の分類。