最近、推論の速度が思ったより出なくてずっとモヤモヤしていた。
GPUに乗ってるのは確認できてるのに、なぜか遅い。
「とりあえずバッチサイズ変えてみるか」という雑な対処を繰り返して、正直ハマってた。
そんなタイミングでHugging Faceのブログに上がってた `torch.profiler` の解説記事を読んだ。
Part 1は行列積とbias addという、めちゃくちゃシンプルな操作を題材にしたやつだ。
「こんな基礎から？」と最初は思ったけど、読み進めると普通に知らなかったことが詰まってた。
記事の中で一番「えぐい」と思ったのが、64x64の行列演算のトレースを読む部分だ。
CPU laneとGPU laneの間に約2.5msのオフセットがある、という話が出てくる。
GPUに処理を投げてから実際に実行されるまで、けっこうな時間がかかってる。
Pythonのコードを書いてると、この非同期ディスパッチの存在を完全に忘れてしまう。
さらに `cudaDeviceSynchronize` が約1.78msも食ってた、という数値が出てくる。
自分でプロファイル取ったことがない人は、これを見てもピンとこないかもしれない。
でも実際にトレースのUIを開いて色付きの矩形を見ると、「ああここで止まってるのか」と体感として入ってくる。
ブログ記事だけ読んでもフワッとするので、自分も手元でスクリプト走らせてみた。
A100ほどのGPUは手元にないけど、手元のRTX 4070で動かしても同じ構造は確認できた。
トレースを開いたとき、ProfilerStep#2だけ異様に長い、という現象もちゃんと再現した。
これ、warmupが走ってないステップでCUDAの初期化コストが乗るからだ。
はじめて見たとき「バグかな？」と思ったけど、記事読んでたおかげで即わかった。
記事の後半は `torch.compile` を乗せたときの話になる。
matmulとaddが別カーネルとして走るのか、それともfuseされて1つになるのかを確認していく。
ここが個人的に一番気になってた部分だ。
4096x4096の大きい行列になると、同じkernelでも実行時間が変わる。
これ、L2キャッシュに乗り切らなくなるからだ、という説明が記事にある。
サイズが上がるとmemory-boundになってきて、compute-boundとは挙動が変わる。
このあたりを「なぜ？」で掘っていくスタイルが記事全体を通じて一貫していて、読んでて気持ちいい。
`torch.compile` を使ったとき、CUDAのlaunch回数が下がるか、という観点も面白かった。
結果として、期待どおりに減ったもの・減らなかったものが両方あった。
特に「CPU overheadが上がった」という部分は意外だった。
compileをかけたら全部速くなると思ってたけど、初回のコンパイルコストやtrace overheadは普通に乗ってくる。
自分のプロジェクトでも似た状況がある。
推論のhotpathで `torch.compile` を試したことがあって、そのときトータルレイテンシが上がって戻した経験がある。
あのとき、ちゃんとプロファイル取って原因を特定していれば、もっとまともな判断ができたはずだ。
「なんか遅くなった、やめよ」で終わらせてしまったのは完全にもったいなかった。
記事の中で `torch.profiler.record_function` でアノテーションをつけると、トレース上でブロックに名前がつく、という話がある。
自分のコードにはこのアノテーションが一切入っていない。
トレースを見ても、どのブロックが自分が書いた処理なのかが判別しにくい状態だ。
まずここから直す。
各forwardのブロックに `record_function` を入れて、トレースを読める状態にする。
Part 2ではnn.LinearとMLP、Part 3ではLLM全体のprofilingに進む予定だとあったので、このシリーズはちゃんと追いかけていく。
「プロファイルは取るべきとわかってるけど、トレースの読み方がわからない」という状態がずっと続いていた。
その入り口をちゃんと整備してくれている記事なので、PyTorch触ってる人には素直に勧めたい。
GPUに乗ってるのは確認できてるのに、なぜか遅い。
「とりあえずバッチサイズ変えてみるか」という雑な対処を繰り返して、正直ハマってた。
そんなタイミングでHugging Faceのブログに上がってた `torch.profiler` の解説記事を読んだ。
Part 1は行列積とbias addという、めちゃくちゃシンプルな操作を題材にしたやつだ。
「こんな基礎から？」と最初は思ったけど、読み進めると普通に知らなかったことが詰まってた。
CPUとGPUの間に2.5msのズレがある
記事の中で一番「えぐい」と思ったのが、64x64の行列演算のトレースを読む部分だ。
CPU laneとGPU laneの間に約2.5msのオフセットがある、という話が出てくる。
GPUに処理を投げてから実際に実行されるまで、けっこうな時間がかかってる。
Pythonのコードを書いてると、この非同期ディスパッチの存在を完全に忘れてしまう。
さらに `cudaDeviceSynchronize` が約1.78msも食ってた、という数値が出てくる。
自分でプロファイル取ったことがない人は、これを見てもピンとこないかもしれない。
でも実際にトレースのUIを開いて色付きの矩形を見ると、「ああここで止まってるのか」と体感として入ってくる。
ブログ記事だけ読んでもフワッとするので、自分も手元でスクリプト走らせてみた。
# 記事と同じ構成で手元再現した
# NVIDIA A100じゃないのでタイミングはズレるけど
python 01_matmul_add.pyA100ほどのGPUは手元にないけど、手元のRTX 4070で動かしても同じ構造は確認できた。
トレースを開いたとき、ProfilerStep#2だけ異様に長い、という現象もちゃんと再現した。
これ、warmupが走ってないステップでCUDAの初期化コストが乗るからだ。
はじめて見たとき「バグかな？」と思ったけど、記事読んでたおかげで即わかった。
torch.compileでkernelがfuseされたかどうかを確認する
記事の後半は `torch.compile` を乗せたときの話になる。
matmulとaddが別カーネルとして走るのか、それともfuseされて1つになるのかを確認していく。
ここが個人的に一番気になってた部分だ。
4096x4096の大きい行列になると、同じkernelでも実行時間が変わる。
これ、L2キャッシュに乗り切らなくなるからだ、という説明が記事にある。
サイズが上がるとmemory-boundになってきて、compute-boundとは挙動が変わる。
このあたりを「なぜ？」で掘っていくスタイルが記事全体を通じて一貫していて、読んでて気持ちいい。
`torch.compile` を使ったとき、CUDAのlaunch回数が下がるか、という観点も面白かった。
結果として、期待どおりに減ったもの・減らなかったものが両方あった。
特に「CPU overheadが上がった」という部分は意外だった。
compileをかけたら全部速くなると思ってたけど、初回のコンパイルコストやtrace overheadは普通に乗ってくる。
自分のプロジェクトでも似た状況がある。
推論のhotpathで `torch.compile` を試したことがあって、そのときトータルレイテンシが上がって戻した経験がある。
あのとき、ちゃんとプロファイル取って原因を特定していれば、もっとまともな判断ができたはずだ。
「なんか遅くなった、やめよ」で終わらせてしまったのは完全にもったいなかった。
次にやること：自分のコードに record_function を入れる
記事の中で `torch.profiler.record_function` でアノテーションをつけると、トレース上でブロックに名前がつく、という話がある。
自分のコードにはこのアノテーションが一切入っていない。
トレースを見ても、どのブロックが自分が書いた処理なのかが判別しにくい状態だ。
まずここから直す。
各forwardのブロックに `record_function` を入れて、トレースを読める状態にする。
Part 2ではnn.LinearとMLP、Part 3ではLLM全体のprofilingに進む予定だとあったので、このシリーズはちゃんと追いかけていく。
「プロファイルは取るべきとわかってるけど、トレースの読み方がわからない」という状態がずっと続いていた。
その入り口をちゃんと整備してくれている記事なので、PyTorch触ってる人には素直に勧めたい。