TorchNeuronのプロファイリング パート1: AWS Trainiumでのシステムレベル推論分析
多くのML実践者はモデルを高速化する方法を理解していますが、モデルが遅い理由を分析するには体系的なプロファイリングが必要です。このガイドでは、TorchNeuronを使用してAWS Trainiumでの推論パフォーマンスを分析する方法を説明し、時間消費のホットスポットを特定し、カーネル融合のパフォーマンス上の利点を示す包括的な方法を提供します。
PyTorchバックエンドとしてのNeuron統合
Neuronバックエンドは、PrivateUse1拡張ポイントを使用してPyTorchと統合します。このメカニズムにより、アクセラレータはcudaなどの専用ディスパッチキーを必要とせずにフレームワークと統合できます。torch_neuronxパッケージは、PyTorchの高度な拡張機能を活用して、NeuronをPyTorchコードのネイティブバックエンドとして登録します。
プロファイラの設定
Neuronワークロードをプロファイリングするには、NeuronConfigとNeuronProfilerの2つのオブジェクトを設定する必要があります。NeuronConfigは、キャプチャするNeuronデータストリームとNeuronCoreを制御し、出力場所を指定します。NeuronProfilerはこの設定を受け取り、export_traceコールバックを提供して、Neuronシステムトレースとともにフレームワークトレース（trace.json）を書き込みます。
主要な設定パラメータは次のとおりです。
modes: デバイス命令やランタイムトレースなど、キャプチャするデータストリームを選択します。capture_enabled_for_nc: トレースする特定のNeuronCoreを選択します。correlate_device_time=True: この設定により、トレースからkey_averages()テーブルにデバイス時間が伝播され、CPUとNeuronデバイス時間を直接比較できるようになります。
ツール: Neuron Explorer
Neuron Explorerは、システムトレース、デバイスタイムライン、エンジン別の内訳を視覚化するように設計されたブラウザベースのビューアです。ユーザーはkey_averages()の数値データを確認し、操作の詳細なライフサイクルを検査できます。
プロファイリング後、出力ディレクトリにはシステムトレースと、デバイスデータがキャプチャされた場合は一致する.neff（Neuron Executable File Format）および.ntff（ハードウェア命令トレース）ファイルが含まれます。Neuron Explorerでプロファイルを表示するには、システムまたは結合プロファイルの場合は出力ディレクトリ全体をアップロードするか、デバイスのみのプロファイルの場合は一致する.neffおよび.ntffファイルを個別にアップロードする必要があります。
ケーススタディ: RMSNormのプロファイリング
この分析にはRMSNormが使用されました。これは、算術演算とリダクション演算の明確なシーケンスを持つ小さな操作だからです。パフォーマンスを比較するために、RMSNormの3つの実装がプロファイリングされました。
- ナイーブなRMSNorm: 個々のPyTorch操作（乗算、平均、加算、逆平方根、乗算）を使用して実装。
torch.rms_norm(): PyTorchの組み込み関数で、単一の融合Neuronカーネルに低減されます。- ナイーブなRMSNorm +
torch.compile: Neuronバックエンドでtorch.compileを使用してコンパイルされた元のナイーブな関数。
すべてのスクリプトは、結果を直接比較できるように、同一の入力、プロファイラ設定、スケジューリングを使用しました。
プロファイル1: ナイーブなPyTorch操作
RMSNormが個々のPyTorch操作を使用して実装されている場合、各行は1つ以上のATen操作に対応し、中間テンソルを作成できます。プロファイリング結果は、3つのaten::mul呼び出しがNeuronデバイス時間の最大のシェアを占めることを示しています。
| 名前 | 呼び出し回数 | 自己CPU | 自己NEURON | % NEURON |
|---|---|---|---|---|
| aten::mul | 3 | 60.9 μs | 1,164.5 μs | 49.5% |
| aten::add | 1 | 26.6 μs | 434.6 μs | 18.5% |
| aten::rsqrt | 1 | 15.6 μs | 415.0 μs | 17.6% |
| aten::mean | 1 | 21.1 μs | 340.8 μs | 14.5% |
自己NEURON合計: 2.355 ms
このナイーブな実装では、各操作がディスパッチ境界を個別に越えます。たとえば、aten::add操作は、スカラーイプシロンがCPU上で開始されるため、ホストからデバイスへのコピーをトリガーします。
プロファイル2: torch.rms_norm()
個々の操作をPyTorchの組み込みrms_norm関数に置き換えることで、Neuronは操作全体を1つの融合カーネルに低減します。キー平均テーブルは、この操作を単一のデバイス操作として示しています。
| 名前 | 呼び出し回数 | 自己CPU | 自己NEURON | % NEURON |
|---|---|---|---|---|
| aten::_fused_rms_norm | 1 | 47.6 μs | 403.3 μs | 100% |
自己NEURON合計: 403.3 μs
この融合実装は、単一のフレームワークからハードウェアへのパスに従い、各算術演算を個別にディスパッチするオーバーヘッドを回避します。
プロファイル3: torch.compileを使用したナイーブなPyTorch
元のナイーブなRMSNorm関数が@torch.compile(backend='neuron', fullgraph=True)を使用してコンパイルされると、TorchDynamoはPythonバイトコードをFXグラフとしてキャプチャします。Neuronバックエンドはこの完全なグラフを受け取り、それを融合して単一のNEFFアーティファクトにコンパイルします。このアーティファクトは単一のハードウェアディスパッチとして実行されます。
| 名前 | 呼び出し回数 | 自己CPU | 自己NEURON | % NEURON |
|---|---|---|---|---|
| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |
自己NEURON合計: 375.5 μs
このコンパイルされたグラフも、ランタイムスケジューリングとハードウェア実行を通る単一のパスに従います。ここで達成されたNeuron時間（375 μs）は、専用のtorch.rms_norm（403 μs）に匹敵します。
並べて比較したパフォーマンスサマリー
| バリアント | NEURON時間 | 自己CPU時間 | オペレーター呼び出し |
|---|---|---|---|
| ナイーブな操作 | 2,355 µs | ~125 µs | 6 |
| torch.rms_norm | 403 µs | ~48 µs | 1 |
| ナイーブ + torch.compile | 376 µs | ~98 µs | 1 |
torch.rms_normとtorch.compileの両方がRMSNormを1つのカーネルディスパッチに正常に融合します。これらの融合方法により、デバイス時間はナイーブな実装のデバイス時間の約6分の1に削減されます。融合により、ハードウェア実行間のギャップがなくなり、個別にディスパッチされた操作間でx_sqやmean_sqなどの中間値のマテリアライズが防止されます。
結論
key_averages()テーブルは、どの操作がデバイス時間を消費するかを識別し、完全な出力はCPU側のセットアップと転送に関するコンテキストを提供します。Neuron Explorerは対応する視覚トレースを提供し、ユーザーはフレームワークディスパッチからランタイムスケジューリング、ハードウェア実行まで操作を追跡し、融合によって除去されるギャップを観察できます。
このワークフローはより大きなモデルに適用できます。テーブル内のパフォーマンスホットスポットを特定し、トレースで視覚的に確認し、融合操作またはtorch.compileがそのコストを削減できるかどうかを判断します。将来の分析では、デバイスレベルのプロファイリングに移行し、TensorEngine、VectorEngine、DMAの内訳を検査して、カーネルが計算バウンドかメモリバウンドかを判断します。