AWS Trainium環境でのPyTorchプロファイリング手法を解説、米Amazon担当者が技術ブログを公開

要点

• 米Amazonの担当者が、AWSのAIアクセラレータ「AWS Trainium」上でPyTorchモデルの推論性能を体系的に分析するプロファイリング手法を解説した。

• 専用バックエンド「TorchNeuron」はPyTorchの拡張機構「PrivateUse1」を採用し、標準ワークフローのサポートや最適化インターフェースとの統合を実現している。

• torch.profiler と専用ツール「Neuron Explorer」を組み合わせることで、CPUとデバイスの実行時間やハードウェア命令トレースの可視化が可能となった。

• 検証事例としてRMSNormを取り上げ、ナイーブな実装、組み込み関数によるカーネル融合、torch.compile の3手法による挙動の違いを比較している。

• 米AmazonのJosh Longenecker氏らは2026年10月21日、Hugging Faceの公式ブログにおいて、AWSの専用AIチップ「AWS Trainium」上で動作するPyTorchモデルの推論性能を体系的に分析するプロファイリング技術の解説記事を公開した。機械学習モデルの処理が遅くなっている原因を特定するための基礎的な手順を、初学者向けに整理して紹介している。本稿はシリーズの第1回にあたり、システムレベルでの性能測定や専用ツールの設定方法が具体的に示されている。

PyTorchとTrainiumをネイティブ統合する「TorchNeuron」

記事ではまず、プロファイリングの前提となるハードウェアとフレームワークの統合方式が説明されている。AWS Trainiumは深層学習の学習および推論ワークロード向けに設計された専用チップであり、そのオープンソースバックエンドとして「TorchNeuron（torch-neuronx）」が提供されている。TorchNeuronは、即時実行モード（eager execution）や torch.compile、標準PyTorch APIをわずかなコード変更のみで利用可能にする。さらに、ハードウェア向けカスタム演算子の作成や最適化を担う「Neuron Kernel Interface（NKI）」とも統合されているという。

この統合の基盤には、PyTorchが備える拡張機構「PrivateUse1」が採用されている。フレームワーク本体に「cuda」のような専用ディスパッチキーを新たに追加することなく外部アクセラレータをネイティブ登録できる仕組みであり、TorchNeuronもこの機構を通じてPyTorchと連携していると解説している。

torch.profiler による測定設定とメトリクス収集

性能測定の設定には、PyTorch標準のプロファイラである torch.profiler と連携する2つの専用オブジェクトが用いられる。

収集対象のデータストリームやNeuronCore、出力先を制御する「NeuronConfig」と、トレース出力のコールバックを提供する「NeuronProfiler」である。NeuronConfigでは、デバイス命令やランタイムトレース（ProfileMode.DEVICE や ProfileMode.RUNTIME など）といった取得対象や、追跡するNeuronCoreの番号を指定する。さらに設定項目 correlate_device_time を有効化することで、デバイス側の処理時間を集計表（key_averages()）へ反映させ、CPUの処理時間とNeuronデバイス側の実行時間を直接比較できるようになる。NeuronProfilerはフレームワークのトレースファイル（trace.json）をシステムトレースと共に出力する役割を果たす。

ブラウザ上でハードウェア実行を追跡する「Neuron Explorer」

数値データの確認にとどまらず、処理の流れを詳細に把握するための専用ビューアとして「Neuron Explorer」が紹介されている。ブラウザ上で動作し、システムトレースやデバイスタイムライン、エンジンごとの内訳を表示できるツールで、利用者はSSHトンネルまたはVisual Studio Codeの拡張機能を通じてアクセスする。

プロファイリング実行後の出力ディレクトリには、システムトレースに加え、コンパイル済みバイナリであるNEFF（Neuron Executable File Format）ファイルと、ハードウェア命令トレースであるNTFFファイルのペアが生成される。両ファイルは共通の数値ハッシュで紐付けられており、ディレクトリ全体または .neff と .ntff のペアをNeuron Explorerにアップロードすることで、ディスパッチからハードウェア実行までの詳細な挙動を追跡できる仕組みとなっている。

RMSNormにおける3つの実装パターンの比較

実践的な分析例として、明確な算術・リダクション演算シーケンスを持つ正規化処理「RMSNorm」を取り上げている。同一の入力条件とプロファイラ設定のもと、以下の3種類の実装が比較された。

第1に、乗算や平均、平方根の逆数などを個別のPyTorch演算で記述するナイーブな実装である。この手法ではコードの各行が基本演算（ATen演算）に対応し、中間テンソルが生成され得る。第2に、単一の融合カーネルへと低水準化されるPyTorch組み込み関数 torch.rms_norm() を用いた実装。第3に、ナイーブな実装に対して torch.compile を適用した構成である。これらを通じて、同一の演算であっても実装形式の違いによってランタイムのスケジューリングやNeuronCore上での実行シグナルがどう変化するかを分析できるとしている。なお、検証用スクリプトは torch-neuronx リポジトリ内で公開されている。