10月23日、Amazon Scienceが「The science behind TorchNeuron, Amazon's new PyTorch backend for AWS Trainium chips」と題した記事を公開した。AWSのVPかつDistinguished EngineerであるRobert Hundtが執筆したこの記事では、AWSのAI専用チップTrainiumでPyTorchコードをそのまま動かすための新しいバックエンド「TorchNeuron」の技術的な仕組みと評価手法が詳しく解説されている。コンパイル時間を実質ゼロに近づける「Adaptive Eager Mode」と、AIエージェントを活用した大規模テスト戦略の二点が、このプロジェクトの核心だ。
TorchNeuronとは何か
TorchNeuronは、PyTorchのコードをAWS Trainium上で実行可能な命令列に変換するオープンソースのバックエンドだ。既存のPyTorchコードは、他のアクセラレータ向けに書かれたものであっても、コード変更を最小限に抑えてTrainiumで動作する。
PyTorchはAI開発者が最も広く使うフレームワークであり、テンソル（行列の高次元版）に対する演算をPythonで記述できる。TorchNeuronはこのPyTorchとAWS Neuron SDKを橋渡しし、Trainiumが本来持つ大規模テンソルワークロード向けの性能を引き出す。Neuron SDKには、命令をサイクル単位で追跡できるプロファイラや、ほぼアセンブリレベルの制御を可能にするNeuron Kernel Interface（NKI）が含まれる。
なお、Neuron SDK自体はTorchNeuron登場以前からtorch-neuronx等のコンポーネントを提供しており、TorchNeuronはそれらを統合・発展させる形で設計されている。また、GPU向けのCUDAやGoogleのXLAといった既存のアクセラレータバックエンドと同様の役割を担いつつ、Trainiumのハードウェア特性に最適化された独自のコンパイルパスを持つ点が差別化要素となっている。
エンジニアが注目すべき点：「Adaptive Eager Mode」
TorchNeuronの技術的な肝は、「Adaptive Eager Mode（適応的イーガーモード）」という実行モードだ。
PyTorchには元来2つの実行モードがある。イーガーモードはコードを1行ずつ即座に実行するため、デバッグに向く。一方のコンパイルモードは、テンソル演算の制御フローグラフを最適化してから実行するため効率はよいが、コンパイル時間がかかる。
Adaptive Eager Modeはその中間に位置するオプション機能だ。CPU上のPythonインタープリタがPyTorchのオペレータ・ディスパッチから命令を収集し、小さなグラフを構築する。そのグラフをTorchNeuronがコンパイル・キャッシュして実行キューに投入する仕組みで、複数のグラフ収集とコンパイルが並列で走る。
さらにスマートな点は、「どのサブグラフをコンパイルするか」を動的に選択するアルゴリズムだ。グラフのサイズや入力テンソルの形状（shape）はコンパイル時間に直結するため、このアルゴリズムは「すでにコンパイル済みの処理を実行している時間」とコンパイル時間を最大限に重ねるよう制御する。結果として多くのモデルでコンパイル時間を実質ゼロに近づける。
デバッグ上重要な「ソースコードとの対応関係（source correlation）」もこのモードで維持されており、イーガーモードのデバッグ体験を壊さない設計になっている。
複数レベルのキャッシュ機構
コンパイル済みグラフはNEFF（Neuron Executable File Format）という実行ファイルに変換される。TorchNeuronはこのNEFFをインメモリキャッシュとディスク上の永続キャッシュの二段構えで保存し、不要な再コンパイルを避ける。
大規模モデルの場合、EC2 UltraServer（最大144基のTrainiumチップを搭載）でも1台に収まらず、数千チップにまたがってモデルをシャーディングすることがある。この場合、すべてのチップが同一のモデル断片を実行するため、同じNEFFが必要になる。
TorchNeuronはここで共有コンパイルサービスを動かし、1台がコンパイルしたNEFFを他のシステムが再利用できるようにする。この最適化により、大規模言語モデルを複数アクセラレータに分散させたテストでは、最初のトレーニングステップまでの時間を約3分の1削減した。
評価方法：AIエージェントによるモデル収集と段階的テスト戦略
TorchNeuronの開発で特筆すべきは、性能評価の手法だ。
チームはAIエージェントを使ってHugging Faceのリポジトリを広範にクロールし、壊れた設定のモデルを除外した上で、動作可能なコードを持つ大量のモデルを収集した（元記事では収集規模について "hundreds of thousands" 程度の表現が用いられている）。当然そのすべてでテストするのは現実的でない。
そこで別のエージェントが各モデルをアーキテクチャ上の特性（エンコーダの種類、Transformerのアテンション機構の形式、出力の形式など）で分類した。結果として、841件の代表モデルがTorchNeuronが実環境で直面しうる設計上の特性をすべてカバーしていることがわかった。
さらに、各モデルの実行中にTrainiumが処理を求められたテンソルの形状（shape）を記録し続け、実際に遭遇しうる数万件のshapeレジストリを構築した。このレジストリからユニットテストを自動生成し、以下の3段階で回帰テストを構成した：
- ユニットテスト（高速、shapeレジストリベース）
- 841件の代表モデルでのテスト
- 全モデル分布でのテスト
この段階的なテスト戦略により、開発サイクルの早い段階で問題を検出できるようになり、開発速度を大幅に改善したという。Adaptive Eager Modeがランタイム側の課題を解くものだとすれば、このテスト戦略は品質保証側の課題を解く仕組みであり、両輪でTorchNeuronの実用性を支えている。
詳細はThe science behind TorchNeuron, Amazon's new PyTorch backend for AWS Trainium chipsを参照していただきたい。
