AWSのNeuron 2.33のリリースノートは、TorchNeuronの公開ベータ版の日付を10月19日としている。同じく10月19日付のPyTorchの発表では、公開ベータ版とオープンソース版の提供開始は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のリリースノートには、モデルを .to("neuron") で移動し、入力テンソルを device="neuron" 上に作成する例がある。対応する torch.nn、autograd、オプティマイザーの操作は、その後バックエンドを通じて実行される。PyTorchのデバイスに依存しない torch.accelerator APIで書かれたコードは、デバイス固有の書き直しなしにNeuronを検出できると、TorchNeuronの概要で説明されている。

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

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

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

AWSのインストールのクイックスタートでは、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は、操作ごとの対応範囲について演算子の互換性表を参照するよう開発者に案内している。したがって、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の性能に関する主張を独立に検証したりはしていない。

出典と参考資料