トレーニング用の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 モードの選択に役立つよう、以下の表を参照してください。

次のステップ

• fullgraph コンパイルと推論ベンチマークの詳細については、torch.compile for inferenceガイドを参照してください。