torch.compile para inferencia

torch.compile compila el código de PyTorch en kernels optimizados que aceleran significativamente la inferencia. Esta función se basa en TorchDynamo para compilar el código en grafos y en TorchInductor para compilar aún más los grafos en kernels optimizados. Es una herramienta de optimización potente y, en muchos casos, solo requiere añadir una línea de código.

Envuelve un modelo con torch.compile para compilarlo y devolver un modelo optimizado.

from transformers import AutoModelForCausalLM

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

La llamada inicial a torch.compile es lenta porque el modelo debe compilarse. Las llamadas posteriores al modelo compilado son mucho más rápidas porque no es necesario volver a compilarlo.

Hay varios parámetros para personalizar el proceso de compilación. A continuación se enumeran dos de los más importantes. Para obtener la lista completa de parámetros, consulte la documentación de torch.compile.

Modos

El parámetro mode ofrece varias opciones de rendimiento para la compilación. Pruebe diferentes modos para ver cuál funciona mejor para su caso de uso.

• default es una opción equilibrada entre velocidad y memoria.
• reduce-overhead reduce la sobrecarga de Python a cambio de un poco más de memoria, pero puede ser más rápido.
• max-autotune ofrece la velocidad más rápida, pero la compilación tarda más.

from transformers import AutoModelForCausalLM

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

Fullgraph

Fullgraph intenta compilar el modelo completo en un único grafo para maximizar el rendimiento. torch.compile genera un error si detecta una interrupción del grafo, lo que significa que no puede compilar el modelo en un ú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 la tabla siguiente para ver los resultados de rendimiento que comparan el tiempo medio de inferencia en milisegundos con torch.compile activado y desactivado en varias GPU y tamaños de lote sobre la misma imagen para diferentes tareas de visión.

Seleccione Subset en la tabla siguiente para cambiar entre diferentes GPUs, así como para ver los benchmarks en PyTorch nightly 2.1.0dev y torch.compile con el modo reduce-overhead habilitado.

Próximos pasos

• Consulte la guía torch.compile for training para aprender a usar torch.compile para el entrenamiento.