PyTorchにおけるプロファイリング：torch.profiler入門ガイド
PyTorchにおけるプロファイリング：torch.profiler入門ガイド
パフォーマンス最適化のための torch.profiler の理解
プロファイリングは、Large Language Models（LLMs）やディープラーニングパイプラインの最適化における不可欠な第一歩であり、開発者がモデルがオーバーヘッドに依存しているか計算に依存しているかを特定できます。torch.profiler モジュールは、パフォーマンス診断のために 2 つの主要な成果物を提供します：profiler table は、最も時間を要する「何」についての統計的サマリーを提供し、profiler trace は、CPU と GPU のレーン全体で「いつ」および「なぜ」操作が発生するかを時間的に示します。
ボトルネックの特定：オーバーヘッドバウンド vs 計算バウンド
パフォーマンスボトルネックは、profiler table における CPU と GPU の使用時間を比較することでしばしば明らかになります。
- Overhead-Bound:
Self CPU time totalがSelf CUDA time totalより著しく大きい場合（例：CPU 時間がミリ秒単位で、GPU 時間がマイクロ秒単位）アルゴリズムはオーバーヘッドバウンドです。これは、CPU がカーネルの準備と起動に費やす時間が、GPU が実行に費やす時間より多いことを示します。小さな行列乗算でよく見られます。 - Compute-Bound:
Self CPU time totalとSelf CUDA time totalがともにミリ秒単位でほぼ同等の場合、アルゴリズムは計算バウンドです。これは、GPU が主なボトルネックとなる高性能コンピューティングにおいて理想的な状態です。
ワークロードサイズを増やす（例：64×64 行列から 4096×4096 行列へ）と、通常はモデルがオーバーヘッドバウンドの状態から計算バウンドの状態へと移行します。
CPU と GPU のディスパッチチェーンの分析
PyTorch の操作は、Python 呼び出しから CUDA カーネルへと特定のディスパッチチェーンをたどります。典型的な行列乗算と加算のシーケンスは次のようになります：
ProfileStep → record_function (user annotation) → aten::matmul (ATen-level dispatch) → aten::mm (2D matrix-matrix multiply backend) → cudaLaunchKernel.
CPU レーンの主要指標
- Cold-Start Overhead: 最初の
ProfileStepは、ワークスペースの割り当て、cuBLAS のヒューリスティック、遅延モジュールロードなどにより、後続のステップよりも幅が広くなることがよくあります。これはウォームアップイテレーションを使用することで緩和できます。 - CUDA Occupancy Queries:
cudaLaunchKernelの前にcudaOccupancyMaxActiveBlocksPerMultiprocessorが存在することは、GEMM や畳み込みのような「ヘビィ」カーネルであることを示します。CPU はハードウェアの容量に基づいて最適なブロックサイズを決定するためにドライバに問い合わせています。 - Resource-Light Kernels: 要素ごとのカーネルやリダクションカーネルは、リソースフットプリントが固定かつ小さいため、通常は占有率クエリを行いません。
- Synchronization: トレースの最後に長時間の
cudaDeviceSynchronizeがある場合、これはプロファイラがイベントをフラッシュしていることが多く、その期間は CPU が GPU の保留中の作業が完了するのを待った時間を示します。
GPU レーンの主要指標
- Activity Buffer Requests: カーネル間のギャップや CPU と GPU のレーン間の初期オフセットは、プロファイラが独自のイベントバッファを割り当てたり再充填したりすることが原因であることが多いです。
- Runtime Variance: 同一のカーネルでも、GPU のクロック変動、熱状態、電源管理、ドライバのメンテナンスなどにより、ステップごとに実行時間が異なることがあります。
torch.compile の影響
torch.compile を使用すると、イーガーな PyTorch コードが TorchDynamo、AOTAutograd、TorchInductor を介して最適化されたグラフに変換されます。コンパイルされた関数をプロファイルすると、いくつかのアーキテクチャ上の変化が明らかになります：
ディスパッチャーレベルの融合
torch.add(torch.matmul(x, w), b) のような操作に対して、torch.compile はグラフレベルで演算子融合を行い、個別の aten::add と aten::mm 呼び出しを単一の aten::addmm 呼び出しに置き換えます。ただし、これは ディスパッチャーレベル の融合であり、カーネルレベルの融合ではありません。GPU は依然として Memcpy DtoD（バイアスで出力バッファを初期化）を実行し、その後バイアス加算エピローグ付きの GEMM カーネルを実行します。
ランタイムアーキテクチャとオーバーヘッド
コンパイルされた領域は、CPU 側に特定の階層構造を導入します：
- TorchDynamo Cache Lookup: 入力の形状、データ型、デバイスがキャッシュされたコンパイルと一致するかを検証します。これはすべての呼び出しで行われます。
- Torch-Compiled Region: コンパイル済みバージョンに入るラッパーです。
- AOTDispatcher Runtime Wrapper Prologue: テンソルのメタデータとビューの追跡を処理します。
- Call CompiledFxGraph: 生成されたコード（コンテンツハッシュで識別）を実行します。
非常に小さな操作の場合、torch.compile は実際に CPU のオーバーヘッドを増加させることがあります。これは、Dynamo $\rightarrow$ AOTAutograd $\rightarrow$ Inductor スタックをたどるコストが融合による節約を上回るためです。このオーバーヘッドは、モデル内の操作数が増えるにつれて相殺されます。
トレース読み取りサマリーテーブル
| Observation | Likely Meaning |
|---|---|
Self CPU $\gg$ Self CUDA |
オーバーヘッドバウンド；バッチサイズを増やすか、演算を融合する |
cudaOccupancy... before launch |
ヘビィウェイトで、適応的に起動されたカーネル（GEMM/Conv） |
aten::matmul $\rightarrow$ aten::bmm |
3次元以上のテンソルに対するバッチ行列乗算 |
Torch-Compiled Region |
torch.compile によって最適化されたブロック内で実行される |
Memcpy DtoD before GEMM |
コンパイルモードでの addmm エピローグ用バイアスコピー |