PyTorch コンパイルを極め、AI/ML 性能を最大化する
2023年3月に PyTorch 2.0 と共に導入されて以来、torch.compile は AI および機械学習ワークロードの性能を最適化するための不可欠なツールとして急速に普及しました。PyTorch は当初、その「Pythonic」な設計、使いやすさ、そして eager（逐次）実行モードによって広く人気を博しました。そのため、Just-in-Time (JIT) グラフコンパイルモードの成功的な採用は、自明ではありませんでした。しかし、わずか2年余りで、実行時性能を向上させる上でのその重要性は計り知れません。その強力さにもかかわらず、torch.compile は依然として秘術のように感じられることがあります。その恩恵は機能するときは明らかですが、多数の API 制御ややや分散したドキュメントのため、問題の診断は困難を伴うことがあります。本記事は、torch.compile を解明し、そのメカニズムを説明し、効果的な適用戦略を実演し、モデル性能への影響を評価することを目的としています。
PyTorch のデフォルトの eager 実行モードは、デバッグにはユーザーフレンドリーですが、本質的に最適化の機会を犠牲にします。各 Python コード行が独立して処理されるため、オペレーター融合（複数の GPU 操作を単一のより効率的なカーネルに結合すること）や、メモリレイアウトと実行順序のための先行 (AOT) コンパイルのような効率化が妨げられます。さらに、Python インタプリタと CUDA バックエンド間の constant なハンドオフは、かなりのオーバーヘッドを発生させます。torch.compile は、JIT コンパイラとして機能することでこれらの制限に対処します。コンパイルされた関数が最初に呼び出される際、TorchDynamo を使用して Python コードを中間グラフ表現（しばしば FX Graph と呼ばれる）にトレースします。トレーニングの場合、AOTAutograd は後方パスをキャプチャし、結合された前方および後方グラフを生成します。このグラフはその後、コンパイラバックエンド（通常は TorchInductor）に渡され、カーネル融合や順不同実行などの広範な最適化が実行されます。NVIDIA GPU の場合、TorchInductor は Triton コンパイラを活用して高度に最適化された GPU カーネルを作成し、可能な場合は CUDA Graphs を利用して複数のカーネルを効率的で繰り返し可能なシーケンスに結合します。結果として生成されるマシン固有の計算グラフはキャッシュされ、その後のすべての呼び出しで再利用されるため、Python インタプリタの関与が大幅に減少し、グラフ最適化が最大限に引き出されます。
torch.compile は通常、モデルのスループットを向上させますが、開発者は時として性能が停滞したり、むしろ低下したりするシナリオに遭遇します。低速なデータ入力パイプラインのような外部ボトルネックの他に、2つの主要な「コンパイルキラー」がしばしば原因となります。それはグラフブレーク (graph-breaks) と再コンパイル (recompilations) です。
グラフブレークは、トレーシングライブラリである TorchDynamo または AOTAutograd が、グラフ操作に変換できない Python 操作に遭遇したときに発生します。これにより、コンパイラはコードをセグメント化し、各部分を個別にコンパイルし、セグメント間で Python インタプリタに制御を戻すことを余儀なくされます。この断片化は、カーネル融合のようなグローバル最適化を妨げ、torch.compile の利点を完全に打ち消す可能性があります。一般的な原因としては、print() 文、複雑な条件ロジック、asserts などが挙げられます。残念ながら、グラフブレークはしばしばサイレントに eager 実行にフォールバックします。これに対処するために、開発者はコンパイラを設定してそれらを報告させることができます。例えば、TORCH_LOGS="graph_breaks" を設定するか、fullgraph=True を使用してブレークに遭遇した際にコンパイル失敗を強制します。解決策としては通常、条件ブロックを torch.where や torch.cond のようなグラフフレンドリーな代替手段に置き換えるか、コンパイル中でない場合にのみ print/assert 文を条件付きで実行することです。
二番目の大きな落とし穴はグラフの再コンパイルです。最初のコンパイル中に、torch.compile は入力（例えばテンソルのデータ型や形状）について「ガード」と呼ばれる仮定を立てます。後続のステップでこれらのガードが破られた場合、例えばテンソルの形状が変更された場合、現在のグラフは無効になり、コストの高い再コンパイルがトリガーされます。過度な再コンパイルは、すべての性能向上を帳消しにし、デフォルトの8回の再コンパイル制限に達すると eager モードへのフォールバックを招くことさえあります。再コンパイルの識別は、TORCH_LOGS="recompiles" を設定することで行えます。動的な形状を扱う際には、いくつかの戦略が存在します。デフォルトの動作 (dynamic=None) は動的な特性を自動検出し、外科的に再コンパイルしますが、これは再コンパイル制限に達する可能性があります。動的な形状が事前にわかっている場合、torch._dynamo.mark_dynamic を用いて動的なテンソルと軸を明示的にマークするのが最良のアプローチとなることが多く、これによりコンパイラは再コンパイルなしで動的な特性をサポートするグラフを構築するように指示されます。あるいは、dynamic=True を設定すると、コンパイラは最大限に動的なグラフを作成するように指示されますが、これにより CUDA Graphs のような一部の静的最適化が無効になる可能性があります。より制御されたアプローチとしては、動的なテンソルをいくつかの事前定義された長さにパディングすることで、固定された限られた数の静的グラフをコンパイルし、すべてのグラフバリエーションがウォームアップフェーズ中に作成されるようにする方法があります。
コンパイル問題のデバッグは、しばしば長くて不可解なエラーメッセージを伴うため、困難を極めることがあります。アプローチは、「トップダウン」（コンパイルをモデル全体に適用し、問題が発生するたびに対処する。注意深いログの解読が必要）から、「ボトムアップ」（低レベルのコンポーネントを段階的にコンパイルし、エラーが特定されるまで進める。特定が容易になり、部分的な最適化の恩恵も得られる）まで様々です。これらの戦略を組み合わせることで、多くの場合、最良の結果が得られます。
モデルが正常にコンパイルされた後も、さまざまなチューニングオプションを通じてさらなる性能向上が期待できますが、これらは通常、初期コンパイルと比較して小さな改善にとどまります。「reduce-overhead」や「max-autotune」のような高度なコンパイラモードは、それぞれオーバーヘッドの削減や複数のカーネルオプションのベンチマークに対して積極的に最適化できますが、コンパイル時間は増加します。異なるコンパイラバックエンドを指定することも可能で、NVIDIA GPU のデフォルトは TorchInductor ですが、ipex のような他のバックエンドは Intel CPU により適している場合があります。静的コンポーネントと動的コンポーネントが明確に分かれているモデルの場合、モジュール式コンパイル（個々のサブモジュールに torch.compile を適用する）により、各部分に合わせた最適化設定が可能になります。モデル自体だけでなく、PyTorch 2.2 ではオプティマイザをコンパイルする機能が導入され、トレーニングワークロードの性能がさらに向上しました。例えば、おもちゃの画像キャプションモデルでオプティマイザをコンパイルしたところ、スループットが1秒あたり5.17ステップから5.54ステップに向上しました。
初期コンパイルとウォームアップ時間は長くなることがありますが、通常、モデル全体のトレーニングまたは推論のライフスパンに比べれば無視できる程度です。しかし、コンパイルに数時間かかるような極めて大規模なモデルや、起動時間がユーザーエクスペリエンスに影響する推論サーバーの場合、この時間の短縮が極めて重要になります。2つの主要なテクニックは、コンパイル時キャッシュと地域コンパイルです。コンパイル時キャッシュは、コンパイルされたグラフアーティファクトを永続ストレージ（例：Amazon S3）に保存し、後続の実行で再ロードすることで、ゼロからの再コンパイルを回避します。デモンストレーションでは、これによりコンパイルウォームアップが196秒から56秒に短縮され、3.5倍の高速化が達成されました。地域コンパイルは、大規模モデル全体ではなく、繰り返される計算ブロックに torch.compile を適用します。これにより、そのブロックのすべてのインスタンスで再利用される単一の小さなグラフが作成されます。おもちゃのモデルの場合、これによりウォームアップが196秒から80秒に短縮されました（2.45倍の高速化）。ただし、スループットは1秒あたり7.78ステップから7.61ステップにわずかに減少しました。小さなモデルでのゲインは控えめですが、これらのテクニックは現実世界の大規模展開において不可欠となり得ます。
最終的に、AI/ML モデルが複雑さと規模を増し続けるにつれて、その実行時性能の最適化は最重要課題となります。torch.compile は PyTorch の最も強力な最適化ツールの一つであり、提示された例では、一部の静的グラフシナリオで最大78%高速、動的グラフで最大72%高速という顕著な速度向上を実現できます。グラフブレークや再コンパイルといった一般的な落とし穴を回避することから、設定の微調整、高度な機能の活用に至るまで、そのニュアンスを習得することは、その可能性を最大限に引き出す上で不可欠です。