torch.compile（トーチ コンパイル）
PyTorch 2.x で導入された JIT コンパイル機能。Python のモデルコードを動的にキャプチャし、TorchInductor バックエンドと Triton GPU コンパイラを通じて最適化されたネイティブコードを生成する。1行の追加でモデルの推論・学習速度を 30-200% 向上させる。
torch.compile とは
torch.compile は PyTorch 2.0（2023 年 3 月リリース）で導入された、モデルの JIT コンパイル API である。model = torch.compile(model) の 1 行を追加するだけで、PyTorch の eager execution モードから最適化されたコンパイル実行に切り替わり、推論・学習の速度が大幅に向上する。
動作原理
torch.compile は以下の 3 つのコンポーネントで構成される。
| コンポーネント | 役割 |
|---|---|
| TorchDynamo | Python バイトコードを解析し、PyTorch 演算のグラフ（FX Graph）をキャプチャ |
| AOTAutograd | 勾配計算を事前（Ahead-of-Time）にトレースし、逆伝播のグラフも最適化対象に含める |
| TorchInductor | FX Graph から CPU（C++/OpenMP）および GPU（Triton）のネイティブコードを生成 |
従来の TorchScript（torch.jit.script/trace）とは異なり、Python の動的な制御フロー（if/for/例外処理）を ガード付きで 正しくキャプチャできる。グラフが変化（graph break）した場合は該当部分のみ eager 実行にフォールバックし、残りのグラフは最適化を維持する。
コンパイルモード
torch.compile は 3 つのモード（mode）を提供する。
- default: 安全性と速度のバランス。ほとんどのモデルで推奨
- reduce-overhead: CUDA グラフを活用し、カーネルローンチのオーバーヘッドを削減。小バッチ推論に有効
- max-autotune: 複数のカーネル実装を自動ベンチマークし、最速の組合せを選択。初回コンパイルは遅いが最高性能
model = torch.compile(model, mode="max-autotune")
性能向上の実例
| モデル | モード | 速度向上 | 備考 |
|---|---|---|---|
| ResNet-50 推論 | default | +43% | Eager 比、batch=64 |
| BERT-Large 学習 | default | +41% | A100 GPU |
| GPT-2 推論 | reduce-overhead | +70% | batch=1、CUDA グラフ適用 |
| Stable Diffusion | max-autotune | +49% | UNet の最適化が支配的 |
| LLaMA-7B 推論 | default | +30% | KV キャッシュ付き |
Graph Break の理解
torch.compile が Python コードを完全にグラフ化できない箇所で Graph Break が発生する。Graph Break が多いとコンパイル効果が低下する。
主な Graph Break の原因は以下の通りである。
print()やlogging等の副作用を持つ関数呼び出しtorch.Tensorの.item()や.numpy()による Python スカラーへの変換- サポートされていないサードパーティライブラリの呼び出し
- 動的なテンソル形状に依存する制御フロー
Graph Break の診断には torch._dynamo.explain(model, *inputs) が有用。
FAQ
Q1: torch.compile は全てのモデルで高速化されるか？
ほとんどのモデルで高速化されるが、Graph Break が多いモデル（複雑な制御フロー・カスタム Python ロジックが多い）では効果が限定的になる。また、非常に小さいモデルではコンパイルのオーバーヘッドが速度向上を相殺する場合がある。
Q2: TorchScript から torch.compile に移行すべきか？
PyTorch チームは torch.compile を推奨後継としている。TorchScript は Python の動的機能との互換性が低く、移行が困難なケースが多かった。torch.compile は既存の eager コードをほぼそのまま高速化でき、移行コストが極めて低い。
Q3: コンパイル時間を短縮する方法は？
初回コンパイルは数十秒〜数分かかる場合がある。torch._inductor.config.cache_dir でカーネルキャッシュを永続化すれば、2 回目以降はキャッシュヒットにより瞬時に起動する。max-autotune モードは特にコンパイルが遅いため、本番デプロイ前に 1 回実行してキャッシュを温めておくのが定石。