Hugging Faceプラットフォームに公開された記事で、Amazonは推論パフォーマンス（学習済みモデルの実行）分析に関するシリーズを開始した。この初回パートでは、AWSの深層学習モデルの学習と推論向けに設計されたチップ「Trainium」に適用されるプロファイリングツールを用いて、時間がどこで消費されているかを特定する方法を説明している。
このガイドは、プロファイリングの事前経験がない読者を対象としている。PyTorchの標準ワークフロー（eagerモード実行やtorch.compileを含む）をコードの変更を最小限にサポートする、Trainium向けのネイティブなオープンソースバックエンド「TorchNeuron」を基盤としている。TorchNeuronはまた、カスタムオペレータの作成を可能にするNeuron Kernel Interface（NKI）とも統合されている。
プロファイラの設定
プロファイリングは2つのオブジェクトに依存する。NeuronConfigは、キャプチャするデータフロー、トレースするNeuronCores、出力フォルダを定義し、その後torch.profiler.profileに渡される。NeuronProfilerは、フレームワークのトレースをtrace.json形式で書き出し、Neuronのシステムトレースと並べて保存するエクスポート関数を提供する。パラメータcorrelate_device_time=Trueを設定すると、デバイス上で費やされた時間がkey_averages()テーブルに表示され、CPU時間と比較できるようになる。
トレースはその後、ブラウザで動作する可視化ツール「Neuron Explorer」で開くことができる。これにより、フレームワークでのディスパッチからランタイムでのスケジューリングを経て、ハードウェア上での実行に至るまでの操作を追跡できる。チップの計算コアであるNeuronCoreの場合、NEFFファイルにはコンパイル済みの実行可能ファイルが含まれ、対応するNTFFファイルにはハードウェア命令のトレースが含まれる。
同一操作の3つのバージョン
この記事では、一連の算術計算とリダクションからなる単純な操作であるRMSNorm関数を検討している。3つの実装が比較されている：PyTorchで操作ごとに記述された素朴なバージョン、組み込み関数torch.rms_norm()の呼び出し、そして操作ごとに記述されたがtorch.compileでコンパイルされたバージョンである。
素朴なバージョンでは、テーブルから、aten::mulへの3回の呼び出しがデバイス上で費やされた時間の大部分を占めていることがわかる。6回のオペレータ呼び出しに対し、Neuronの合計時間は2 355 µsに達する。著者によれば、単一の融合計算カーネルに変換される組み込み関数torch.rms_norm()は、単一の操作aten::_fused_rms_normのみを占め、Neuron時間は403,3 µsである。torch.compileでコンパイルされたバージョンは、単一のコンパイル済み領域として表示され、その持続時間は375,5 µsである。
記事では、これらのテストにおいて、両方のアプローチが素朴なバージョンの計算時間の約6分の1を要したと述べている。融合により、ハードウェア実行間の中断が排除され、x_sqやmean_sqのような中間値が、別々に送信される操作間で具体化されるのを防ぐことができる。
著者は、この方法をより大きなモデルに適用できると指摘している：テーブル内のホットスポットを特定し、トレースで確認し、融合操作やtorch.compileがそのコストを削減できるかどうかを判断する。シリーズの第2部では、NEFFファイルの分析を通じて、デバイスレベルのプロファイルに焦点を当てる予定である。
