torch.compile para inferência

torch.compile compila código do PyTorch em kernels otimizados que aceleram significativamente a inferência. Esse recurso depende do TorchDynamo para compilar o código em grafos e do TorchInductor para compilar ainda mais os grafos em kernels otimizados. É uma ferramenta de otimização poderosa e, em muitos casos, exige apenas a adição de uma única linha de código.

Envolver um modelo com torch.compile para compilar e retornar um modelo otimizado.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

A chamada inicial para torch.compile é lenta porque o modelo precisa ser compilado. As chamadas subsequentes ao modelo compilado são muito mais rápidas, pois ele não precisa ser compilado novamente.

Existem vários parâmetros para personalizar o processo de compilação. Dois dos mais importantes estão listados abaixo. Para a lista completa de parâmetros, consulte a documentação do torch.compile.

Modos

O parâmetro mode oferece várias opções de desempenho para compilação. Experimente diferentes modos para ver qual funciona melhor para o seu caso de uso.

• default é uma opção equilibrada entre velocidade e memória.
• reduce-overhead reduz o overhead do Python em troca de um pouco mais de memória, mas pode ser mais rápido.
• max-autotune oferece a velocidade mais rápida, mas a compilação leva mais tempo.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph tenta compilar o modelo inteiro em um único grafo para maximizar o desempenho. torch.compile levanta um erro se encontrar uma quebra de grafo, o que significa que não é possível compilar o modelo em um único grafo.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

Benchmarks

Consulte a tabela abaixo para benchmarks de desempenho comparando o tempo médio de inferência em milissegundos com torch.compile habilitado e desabilitado em vários GPUs e tamanhos de lote na mesma imagem para diferentes tarefas de visão.

Selecione Subset na tabela abaixo para alternar entre diferentes GPUs, bem como benchmarks no PyTorch nightly 2.1.0dev e torch.compile com o modo reduce-overhead ativado.

Próximas etapas

• Consulte o guia torch.compile para treinamento para aprender a usar o torch.compile para treinamento.