PyTorch 2.0 正式リリース：torch.compile 導入でモデル実行を大幅高速化

PyTorch 2.0 は、広く使われている深層学習フレームワーク PyTorch の大型アップデートです。Python らしい直感的な動的グラフの操作感をそのまま維持しながら、バックエンドのコンパイラ基盤を根本から刷新。既存のコードを大きく書き換えることなく、モデルの学習や推論を高速化できます。

本バージョンには、1.13.1 以降の 4,541 件のコミットと 428 人のコントリビューターによる成果が凝縮されています。完全な後方互換性を保ちながら設計されています。

できることと仕組み

最大の特徴は、モデルをラップするだけで利用できる torch.compile API の追加です。既存のモデルに対して model = torch.compile(model) と指定するだけで、最適化されたコンパイル済みモデルが生成されます。

内部では 4 つの要素が連動します。TorchDynamo が動的グラフを安全に抽出し、AOTAutograd が逆伝播を事前トレース、PrimTorch が 2,000 以上の演算子を約 250 個の基本プリミティブに集約し、TorchInductor がコード生成を担います。

GPU 向けコード生成には OpenAI Triton を活用しており、熟練者が手書きした CUDA カーネルに匹敵する性能を発揮します。165 種類のオープンソースモデルを用いた検証では、93% で正常に動作し、float32 精度で平均 20%、AMP 混合精度で平均 36% の速度向上が確認されました。

さらに、Transformer 向けに Scaled Dot Product Attention 演算が torch.nn.functional に標準統合されました。FlashAttention や xFormers などの最適化実装が最初から組み込まれており、外部パッケージを手動導入する必要がなくなりました。

Mac 環境向けには MPS バックエンドが強化され、対応演算子が 300 以上に拡大。AWS の Graviton3 プロセッサ上での ResNet-50 や BERT の CPU 推論も高速化されています。

使い方と試し方

導入は公式サイトの案内に従い、通常の pip または conda コマンドでアップデートするだけです。

コード側の変更は最小限で、学習ループや推論コードの前に torch.compile を 1 行挟むだけで動作します。Apple Silicon 上で MPS バックエンドを使う場合は、最新の macOS に更新しておくことが推奨されています。

なぜ重要か

大規模言語モデルや画像生成モデルを扱う開発者なら、今すぐ試す価値があります。これまで FlashAttention などの高速化技術を使うには面倒なビルドや追加パッケージが必要でしたが、それが公式機能として自然に動くメリットは絶大です。

ただし、torch.compile 自体はまだベータ版の位置付けです。大半のモデルで動くとはいえ、特殊な動的分岐を含むコードでは期待通りの恩恵が得られない場合もあります。本番投入の前には、出力精度の検証をしっかり行っておくのが賢明です。