torch.profiler入門：PyTorchのCPU・GPU待ち時間を読み解く方法

結論から言うと、PyTorchの高速化は、最も長く見える処理を勘で直すのではなく、CPU・GPU・同期の関係をtorch.profilerで可視化することから始まります。GPU計算そのものではなく、処理の依頼、データの受け渡し、完了待ちがボトルネックになっている場合もあるため、トレース全体を読むことが重要です。

torch.profilerでPyTorchの遅さを調べる理由

torch.profilerを使うと、CPU側の処理、GPUで動くカーネル、CUDAランタイム呼び出し、同期処理などを時系列で確認できます。単に「この関数が遅い」と判断するのではなく、誰がGPU処理を起動し、どこで待っているのかまで追跡できる点が重要です。

機械学習のコードでは、PythonからGPUへ計算を依頼した時点で、GPUの仕事が完全に終わるとは限りません。GPUは依頼を受けて非同期に処理を進めるため、CPUは次の処理へ移動できる一方、後の箇所で完了確認が必要になることがあります。

その結果、CPU側で測った処理時間と、GPU側で実際に計算していた時間が一致しないケースがあります。torch.profilerのタイムラインを見れば、計算時間と待ち時間を分けて考えやすくなります。

プロファイリングとは、プログラムの動きを記録し、時間や資源の使われ方を調べる作業です。車の速度だけを見るのではなく、信号待ちや渋滞まで記録するようなものだと考えると、初心者にもイメージしやすいでしょう。

CPUとGPUの非同期処理で時間の見え方がずれる仕組み

CPUとGPUのレーンに表示されるイベントは、同じ処理を別の視点から表しています。CPUがGPUへ仕事を依頼してから次へ進む場合、CPUのイベントが先に終わって見えても、GPU側では計算が続いていることがあります。

小さな行列の例では、この非同期性が特に目立ちます。元記事では64×64の行列を使った処理について、GPUの計算時間そのものは短い一方、CPUからGPUへ依頼する時間や完了を確認する時間が無視できない場合が示されています。

プロファイラ上で「ProfilerStep#2」が長く表示されても、それは単一の計算が長いという意味とは限りません。複数の処理や、途中で発生した待ち時間が一つのステップに含まれている可能性があります。

元記事の例では、CPUとGPUのレーンに約2.5ミリ秒のずれがあり、非同期実行がその理由として説明されています。また、cudaDeviceSynchronizeがCPUの完了待ちとして約1.78ミリ秒を占めるケースも紹介されています。これらの数値は元記事の記録に基づくもので、環境が変われば同じ結果になるとは限りません。

小さな行列と大きな行列で変わるボトルネック

行列が小さい場合は、計算そのものよりも、GPUへ処理を依頼する準備や完了待ちの影響が目立つことがあります。したがって、GPUを使っているという事実だけで、処理全体が必ず速くなるとは言えません。

行列積はmatmulと呼ばれ、複数の行列を掛け合わせる処理です。そこへaddのような足し算を続ける場合、各処理をGPUカーネルとして個別に起動するため、計算以外の起動コストや中間データの受け渡しが発生することがあります。

一方、元記事では4096×4096の大きな行列を使う例も扱われています。このように計算量が増えると、GPU上で実行されるカーネルの処理時間がタイムライン上で目立ちやすくなります。

• 入力データのサイズや形状
• GPU内メモリへのアクセス状況
• 前後の処理との依存関係
• CPUから起動されるタイミング
• 別の処理との競合

同じ種類に見えるカーネルでも、これらの条件によって所要時間が変わる可能性があります。最適化では関数名だけを比較するのではなく、入力、前後関係、起動回数、CPUとのつながりをまとめて確認することが大切です。

torch.compileとカーネル融合をトレースで確認する方法

torch.compileは、PyTorchの処理グラフを解析し、処理のまとめ方や実行方法を改善する仕組みです。うまく働けば、matmulとaddのような処理を個別に実行するのではなく、可能な範囲でまとめて、GPUカーネルの起動回数や中間データの受け渡しを減らせます。

この処理の統合は、カーネル融合と呼ばれます。料理にたとえるなら、材料を切って別の場所へ運ぶ工程を何度も繰り返すのではなく、近い工程を一つの流れにまとめるイメージです。

1. コンパイル前のトレースで、matmulとaddの起動回数を確認する
2. コンパイル後のトレースで、処理がどのように変わったかを見る
3. GPUカーネルが統合されたか、中間処理が減ったかを確認する
4. CPU側のコンパイルやディスパッチ負荷もあわせて比較する

ただし、torch.compileを使えば必ず速くなるという単純な話ではありません。GPU側の処理が効率化されても、コンパイル、処理グラフの管理、ディスパッチ（適切な処理へ振り分ける作業）によってCPU側のオーバーヘッドが増えるケースがあります。

そのため、見るべきなのは最終的な実行時間だけではありません。CUDAカーネルの起動回数、GPU計算の開始と終了、CPUの待機、コンパイルによる追加コストをトレース全体から確認する必要があります。

torch.profilerで最適化の優先順位を決める手順

torch.profilerを使うときは、最初から細かなイベントをすべて理解しようとせず、CPUレーンとGPUレーンの対応関係から読み始めるのがおすすめです。最も長い長方形を見つけたら、すぐにその関数を修正するのではなく、実計算なのか待機なのかを切り分けます。

次の手順で確認すると、初心者でもトレースを整理しやすくなります。

1. 対象の処理を小さな再現例にする
2. CPUとGPUの両方のイベントを表示する
3. GPUカーネルの開始・終了とCPU側の起動処理を対応づける
4. 同期処理やCUDAランタイム呼び出しがないか確認する
5. 起動回数、中間データ、コンパイルの影響を比較する
6. 変更後に同じ観点で再度プロファイルする

特に重要なのは、「どの関数が遅いか」だけでなく、「どこで待たされているか」を見ることです。同期処理が複数の待ち時間を一か所に集めている場合、長いイベントをそのまま計算ボトルネックと判断すると、効果の薄い最適化につながります。

LLMの推論速度を上げたい場合でも、学習時間を短縮したい場合でも、最初にCPU、GPU、メモリ、同期の関係を可視化することが出発点です。「計算を速くする」だけでなく、「待たない」「何度も起動しない」「不要な受け渡しを減らす」という視点で改善策を選びましょう。

プロファイルできないものは、最適化できない。

出典: Profiling in PyTorch (Part 1): A Beginner’s Guide to torch.profiler

よくある質問