AI News HubLIVE
站内改写

PyTorch プロファイリング（第1回）：torch.profiler 入門ガイド

本記事は PyTorch プロファイリングシリーズの第1回です。最も単純な行列乗算とバイアス加算から始め、torch.profiler の設定方法、プロファイラテーブルとトレースの読み方、CPU と GPU のアクティビティ間の時間関係、ウォームアップと行列サイズが性能に与える影響について解説します。

記事インテリジェンス

投資家上級

要点

  • torch.profiler は統計テーブルと時間軸トレースを出力し、ホットスポットと実行タイミングを把握できる。
  • 小さい行列演算はオーバーヘッド律速になるが、サイズを大きくすると計算律速に移行する。
  • ウォームアップにより起動オーバーヘッドを排除し、安定したプロファイル結果を得られる。
  • CPU と GPU の間に生じる時間オフセットは、カーネル起動と同期の遅延を示す。

重要な理由

このニュースが重要なのは、torch.profiler は統計テーブルと時間軸トレースを出力し、ホットスポットと実行タイミングを把握できるためです。

技術的影響

モデル選定、推論コスト、プロダクト能力、評価基準に影響する可能性があります。

PyTorch モデルの最適化にはプロファイリングが不可欠です。LLM のトークン毎秒を向上させたい場合、推論レイテンシを短縮したい場合、あるいはトレーニングループが仕様より遅い理由を理解したい場合でも、最終的にはプロファイリングを実施する必要があります。しかし、プロファイリングには学習曲線が急峻であるという問題があります。トレースは色付きの矩形が密集し、イベント名は難解で、ほとんどのチュートリアルはすでに読み方を知っていることを前提としています。そのため、プロファイリングの重要性を理解していても、トレースを開くことは後回しにされがちです。本記事はその敷居を下げることを目的としています。

本シリーズでは、PyTorch プロファイリングの基礎から始め、プロファイラのトレースを読み解くスキルを徐々に身につけ、最適化に活用していきます。第1回である本記事では、行列乗算とバイアス加算という最も単純な操作から始め、プロファイラが返す情報を読み解く方法を学びます。第2回では nn.Linear と小規模 MLP に拡張し、トレースを活用した最適化とカーネルの内部を覗きます。第3回では Transformer を使った大規模言語モデルに適用します。

準備：基本概念

まず、2つの定義を理解しておきましょう。

  • GPU カーネル：GPU 上の多数のスレッドで並列実行されるプログラムです。
  • CPU がこれらのカーネルをスケジュールし、起動します。

PyTorch の操作を使用すると、自動的に1つ以上の GPU カーネルに変換されます。

プロファイリングコードの作成

最も単純な行列乗算と加算から始めます。コードは以下の通りです。

def fn(x, w, b):
    return torch.add(torch.matmul(x, w), b)

torch.profiler を使用したプロファイリングの手順は以下の通りです。

  1. プロファイリングするコードを用意します（上記）。
  2. アルゴリズムを torch.profiler.record_function でアノテーションします（オプションですが推奨）。
  3. コードを torch.profiler.profile コンテキストマネージャでラップします。
  4. プロファイルをエクスポートします。

プロファイラは2つの異なる成果物を出力します。

  • プロファイラテーブル：イベントの統計サマリーを提供し、「何が最も時間を消費しているか」を回答します。
  • トレース：時間軸の実行ビューを提供し、「操作がいつ、なぜ発生したか」を回答します。

プロファイラテーブルの読み方

スクリプトを実行すると、テーブルを含む .txt ファイルとトレースを含む .json ファイルが生成されます。テーブルの最初の列はトリガーされたイベントで、他の列は CPU、GPU などの時間です。イベントの時間を観察することで、直感的にホットスポットを特定できます。「自己 CPU/CUDA 時間」と「合計 CPU/CUDA 時間」の違いに注意してください。前者はイベント自身のみの時間、後者は子イベントも含みます。

例えば、64×64 の行列の場合、自己 CUDA 時間はわずか 23.104 マイクロ秒で、自己 CPU 時間は 2.314 ミリ秒です。これは GPU がほとんどの時間アイドル状態であることを意味し、典型的なオーバーヘッド律速の状態です。行列サイズを 4096×4096 に増やすと、自己 CUDA 時間は 4.495 ミリ秒、自己 CPU 時間は 4.908 ミリ秒となり、GPU 時間が大幅に増加し、計算律速に移行したことがわかります。

トレースの読み方

トレースデータは Perfetto UI で開くことができます。CPU と GPU のレーンがあり、バーの幅がイベントの持続時間、垂直方向のネストが呼び出し階層を示します。64×64 のトレースでは、以下の現象が観察できます。

  • ProfileStep#2 の実行時間が他のステップより明らかに長い：これは GPU がウォームアップされていないために、起動オーバーヘッドが記録されたためです。ウォームアップ（プロファイル前にコードを数回実行する）により、この一時的なオーバーヘッドを排除できます。
  • CPU レーンと GPU レーンの間に約2.5ミリ秒のオフセットがある：これは CPU がカーネルを起動してから、GPU が実際に実行を開始するまでの遅延で、メモリ転送やカーネルのキューイングなどが原因です。

まとめと展望

本記事では、torch.profiler の基本的な使い方と、テーブルとトレースの読み方を紹介しました。行列サイズとウォームアップの調整により、オーバーヘッド律速から計算律速への移行を観察しました。次回は、これらの概念をより複雑なネットワークに拡張し、プロファイル結果を最適化に活用する方法を学びます。

このパートを終えると、以下のことができるようになります。

  • torch.profiler のセットアップとその出力の理解。
  • プロファイラテーブルとトレース（CPU レーン、GPU レーン、その間のギャップ）の読み取り。
  • Python 呼び出しから CUDA カーネルに至るイベントチェーンの追跡。
  • torch.compile が何を変えるか（そして、さらに興味深いことに、何を変えないか）の理解。