推論のためのtorch.compile

torch.compile は、PyTorch のコードを最適化されたカーネルにコンパイルし、推論を大幅に高速化します。この機能は、コードをグラフにコンパイルするために TorchDynamo を、さらにグラフを最適化されたカーネルにコンパイルするために TorchInductor を使用します。強力な最適化ツールであり、多くの場合、1 行のコードを追加するだけで済みます。

モデルを torch.compile でラップして、コンパイルし、最適化されたモデルを返します。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

torch.compileへの最初の呼び出しは、モデルのコンパイルが必要なため遅くなります。コンパイル済みモデルへのその後の呼び出しは、再度コンパイルする必要がないため、はるかに高速になります。

コンパイルプロセスをカスタマイズするためのパラメータはいくつかあります。中でも特に重要な2つを以下に示します。すべてのパラメータのリストについては、torch.compile のドキュメントを参照してください。

モード

mode パラメータには、コンパイル用のパフォーマンスオプションがいくつか用意されています。使用ケースに最適なモードを見つけるために、異なるモードを試してみてください。

• defaultは、速度とメモリ使用量のバランスが取れたオプションです。
• reduce-overhead は、メモリ使用量が増加する代わりに Python のオーバーヘッドを削減し、より高速になる場合があります。
• max-autotune は最速の速度を実現しますが、コンパイルに時間がかかります。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph は、パフォーマンスを最大化するために、モデル全体を単一のグラフにコンパイルしようとします。torch.compile は、グラフのブレイク（グラフの中断）を検出するとエラーを発生させます。これは、モデルを単一のグラフにコンパイルできないことを意味します。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

ベンチマーク

下表は、異なるビジョンタスクにおいて、同じ画像に対して、torch.compileを有効化した場合と無効化した場合の平均推論時間（ミリ秒）を、さまざまなGPUとバッチサイズで比較したパフォーマンスベンチマークを示しています。

下の表でSubsetを選択すると、異なるGPUや、PyTorch nightly 2.1.0devおよびreduce-overheadモードを有効にしたtorch.compileのベンチマークを切り替えることができます。

次のステップ

• トレーニングでtorch.compileを使用する方法を学ぶには、トレーニング用のtorch.compileガイドを参照してください。