用于训练的 torch.compile

torch.compile 将 PyTorch 代码编译为融合内核，以提升运行速度。在训练过程中，它会同时追踪前向和反向传播，并将其编译为优化后的内核，从而降低单个算子启动的开销，并通过融合操作来减少内存带宽的占用。

在 TrainingArguments 中设置 torch_compile=True 即可启用该功能。与推理仅编译前向传播不同，训练会同时编译前向和反向传播。编译发生在第一个训练步骤，因此该步骤的速度会明显慢于后续步骤。

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。

编译模式

使用下表来帮助您选择 torch.compile 模式。

后续步骤

• 有关 fullgraph 编译和推理基准测试的详细信息，请参阅 torch.compile for inference 指南。