Transformers documentation
トレーニング用のtorch.compile
トレーニング用のtorch.compile
torch.compile は PyTorch のコードを融合カーネルにコンパイルして、実行速度を向上させます。トレーニングでは、フォワードパスとバックワードパスの両方をトレースして最適化されたカーネルにコンパイルし、個々のオペレーション起動のオーバーヘッドを削減するとともに、オペレーションを融合してメモリ帯域幅の使用量を削減します。
torch_compile=True を TrainingArguments に設定して有効にします。推論がフォワードパスのみをコンパイルするのに対し、学習ではフォワードパスとバックワードパスの両方がコンパイルされます。コンパイルは最初の学習ステップで実行されるため、その後のステップよりも大幅に遅くなることを想定してください。
from transformers import TrainingArguments
args = TrainingArguments(
...,
torch_compile=True,
torch_compile_backend="inductor",
torch_compile_mode="reduce-overhead",
)バックエンド
バックエンドが指定されていない場合、TrainingArguments はハードウェアに基づいてバックエンドを選択します。ほとんどの CPU および GPU では、デフォルトは inductor であり、AOTAutograd を使用して Triton カーネルにコンパイルされ、ほとんどのトレーニングワークロードに適しています。Intel Gaudi (HPU) では、デフォルトは hpu_backend です。AWS Trainium および Inferentia (Neuron) では、デフォルトは neuron です。
固定形状のインプットにはcudagraphsを使用します。
Compile モード
torch.compile モードの選択に役立つよう、以下の表を参照してください。
| mode | 説明 |
|---|---|
| default | コンパイル時間と実行時間のバランス |
| reduce-overhead | CUDA グラフを使用して Python/CPU のオーバーヘッドを削減しますが、その代わりにメモリ使用量が増加します |
| max-autotune | コンパイル時に複数のカーネル実装をベンチマークし、最速のものを選択します（コンパイル時間が長くなります）。 |
| max-autotune-no-cudagraphs | max-autotune と同じですが、CUDA graphs を使用しません |
次のステップ
- fullgraph コンパイルと推論ベンチマークの詳細については、torch.compile for inferenceガイドを参照してください。