Trainium向けTorchNeuronで推論のボトルネックを特定する手順

何が解説されたのか

Hugging Faceのブログで、AWS Trainium上のTorchNeuronを対象としたプロファイリング入門記事のPart 1が公開された。TorchNeuronは、深層学習の学習と推論のために設計されたAWSのチップTrainium上で動く、オープンソースのネイティブPyTorchバックエンドである。記事は、プロファイリングの経験がない読者を想定し、推論性能を体系的に分析して時間がどこで消費されているかを特定する手順を、順を追って説明している。

記事の目的は、単にツールの使い方を示すことではなく、「なぜモデルが遅いのか」を分析できるようになることだと位置づけられている。対象読者は、モデルを速く動かす勘所はあるが、その理由を説明できないML実務者とされている。最終的には、torch.profilerの設定方法、key_averages()テーブルの読み方、Neuron Explorerで1つの演算を追跡する方法、そして同じ演算を3通りに実装したときの指標の変化を理解できる構成になっている。

バックエンドと計測の仕組み

Neuronバックエンドは、PyTorchのPrivateUse1拡張ポイントを利用して統合されている。これは、cudaのような専用のディスパッチキーを追加せずにアクセラレータをフレームワークへ組み込める仕組みで、torch_neuronxパッケージが上流のPyTorchと連携してNeuronをネイティブバックエンドとして登録する。TorchNeuronはeager実行、torch.compile、標準的なPyTorch APIを最小限のコード変更でサポートし、Neuron Kernel Interface（NKI）と連携してカスタム演算子やハードウェア向けの最適化も可能にする。

計測には2つのオブジェクトを使う。NeuronConfigは取得するNeuronのデータストリームやNeuronCore、出力先を制御し、torch.profiler.profileにexperimental_configとして渡す。NeuronProfilerはNeuronConfigを受け取り、フレームワークのトレースであるtrace.jsonをNeuronのシステムトレースと並べて書き出すexport_traceコールバックを提供する。

設定では、modesでデバイス命令やランタイムのトレースを選び、capture_enabled_for_ncで対象のNeuronCoreを指定する。correlate_device_time=Trueにすると、トレースのデバイス時間がkey_averages()の表に反映され、CPU時間とNeuronデバイス時間を並べて比較できる。また、プロファイラのactivitiesにはProfilerActivity.PrivateUse1を含める必要があり、ループ内でtorch.neuron.synchronize()を呼ばないと、デバイスの完了前にエクスポートされてNeuron時間の列が空になる点も注意点として挙げられている。

RMSNormで3実装を比較

検証の題材はRMSNormで、算術演算とリダクションが明確な順序で並ぶ小さな演算として選ばれている。比較されたのは、個別のPyTorch演算（乗算、平均、加算、逆平方根、乗算）を並べたNaive実装、PyTorch組み込みのtorch.rms_norm()、そして同じNaive実装をtorch.compileでコンパイルしたものの3通りである。スクリプトは入力、プロファイラ設定、スケジュールを共通化しており、結果を直接比較できるとしている。

Naive実装では、Self NEURON totalは2.355 msで、3回のaten::mulがデバイス時間の最大の割合を占めた。x * x、x * rms、weightの乗算がそれぞれ対応しており、aten::addはCPU側で始まるスカラーを扱うためホスト・デバイス間のコピーを発生させる。torch.rms_norm()は1つの融合カーネルに下げられ、403.3 μsだった。torch.compile版はkey_averages()上で1つのTorch-Compiled Regionとして現れ、375.5 μsだった。

まとめとして、torch.rms_normとtorch.compileはいずれもRMSNormを1回のカーネルディスパッチに融合し、今回の実行ではどちらもNaive実装のデバイス時間の約6分の1になったと報告されている。融合によってハードウェア実行間の隙間がなくなり、x_sqやmean_sqといった中間値を個別ディスパッチ間で実体化せずに済む、というのが著者の説明である。

Explorerでの可視化

Neuron Explorerは、システムトレース、デバイスのタイムライン、エンジンごとの内訳を表示するブラウザベースのビューアである。SSHトンネルかVisual Studio Code拡張からアクセスでき、記事ではkey_averages()が数値で示した内容を視覚的に確認し、ディスパッチの間に何が起きているかを調べるために使われている。

プロファイル後、出力ディレクトリにはシステムトレースと、デバイスデータを取得した場合は対応する.neffファイルと.ntffファイルが含まれる。NEFF（Neuron Executable File Format）はコンパイル済みのNeuron実行ファイルを、対応するNTFFはそのハードウェア命令トレースを格納しており、両ファイルはペアであることを示す共通の数値ハッシュを持つ。システムまたはcombinedプロファイルでは出力ディレクトリ全体を、deviceのみのプロファイルでは対応する.neffと.ntffを個別にアップロードして開く。

Explorerでは時間が左から右へ、イベントのライフサイクルが上から下へ流れる。各ATen演算は、フレームワークのレーンで生成され、ランタイムのレーンでnrta_execute_scheduleとしてスケジュールされ、ハードウェアのレーンでnc_exec_runningとして実行されるという同じパターンをたどる。eager実行では各演算が独立して走るため、依存関係チェーン・ビューアを使って1つの演算のライフサイクルをフレームワーク、ランタイム、ハードウェアの各レーンにまたがって追跡できる。

現時点での位置づけと注意点

今回の比較はRMSNormという小さな演算に対するマイクロベンチマークであり、記事が示す数値は特定の入力形状と実行環境での結果である。より大きなモデルへ同じワークフローを適用する方法は述べられているが、大規模モデルでの具体的な計測結果はPart 1では示されていない。まず表でホットスポットを見つけ、トレースで確認し、融合演算やtorch.compileでコストを下げられるかを判断する、という進め方が推奨されている。

Part 2ではデバイスレベルのプロファイルに進み、NEFFを開いてTensorEngine、VectorEngine、DMA（直接メモリアクセス）の内訳を調べ、カーネルをルーフライン上に置いて計算律速かメモリ律速かを判定する予定だとされている。Part 1は「どの演算が時間を使っているか」までを扱い、「カーネル内部で何が起きているか」は次回に持ち越す構成である。