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, который компилируется в ядра Triton с помощью AOTAutograd и подходит для большинства задач обучения. На Intel Gaudi (HPU) по умолчанию используется hpu_backend. На AWS Trainium и Inferentia (Neuron) по умолчанию используется neuron.
Используйте cudagraphs для входных данных фиксированной формы.
Режим компиляции
Используйте таблицу ниже, чтобы выбрать режим torch.compile.
| mode | Описание |
|---|---|
| default | баланс между временем компиляции и временем выполнения |
| reduce-overhead | снижает накладные расходы Python/CPU с помощью CUDA-графов за счёт дополнительного потребления памяти |
| max-autotune | сравнивает несколько реализаций ядра при компиляции и выбирает самую быструю (компиляция занимает больше времени) |
| max-autotune-no-cudagraphs | аналогично max-autotune, но без CUDA graphs |
Дальнейшие шаги
- См. руководство torch.compile для инференса для получения подробной информации о компиляции fullgraph и бенчмарках инференса.