torch.compile para treinamento

torch.compile compila código do PyTorch em kernels fundidos para acelerar a execução. No treinamento, ele rastreia o passo de ida e o passo de volta juntos e os compila em kernels otimizados, reduzindo o overhead de lançamentos individuais de operações e fundindo operações para diminuir o uso de largura de banda de memória.

Defina torch_compile=True em TrainingArguments para ativá-lo. O treinamento compila tanto o passo de ida (forward) quanto o passo de volta (backward), ao contrário da inferência, que compila apenas o passo de ida. A compilação ocorre no primeiro passo de treinamento, portanto, espere que ele seja significativamente mais lento do que os passos subsequentes.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    torch_compile=True,
    torch_compile_backend="inductor",
    torch_compile_mode="reduce-overhead",
)

Backend

Quando nenhum backend é especificado, o TrainingArguments seleciona um com base no seu hardware. Na maioria dos CPUs e GPUs, o padrão é inductor, que compila para kernels Triton com AOTAutograd e atende à maioria das cargas de trabalho de treinamento. No Intel Gaudi (HPU), o padrão é hpu_backend. No AWS Trainium e Inferentia (Neuron), o padrão é neuron.

Use cudagraphs para entradas com forma fixa.

Modo de compilação

Use a tabela abaixo para ajudar a selecionar um modo de torch.compile.

Próximas etapas

• Consulte o guia torch.compile para inferência para obter detalhes sobre a compilação fullgraph e os benchmarks de inferência.