Compilación

Resumen

Pytorch 2.0 introdujo torch.compile, una función potente que hace que el código de PyTorch se ejecute más rápido mediante la compilación JIT del código de PyTorch en kernels optimizados. Las funciones clave de torch.compile incluyen:

• Mejora del rendimiento: Acelera significativamente la ejecución del modelo optimizando el grafo de cómputo.
• Facilidad de uso: Requiere cambios mínimos de código para implementarse, lo que lo hace muy accesible.
• Compatibilidad: Funciona de manera fluida con el código y los modelos existentes de PyTorch.

Cuando se usa con Accelerate, torch.compile se integra de forma fluida en los flujos de trabajo de entrenamiento distribuido, lo que permite aprovechar simultáneamente la ejecución distribuida y las optimizaciones de compilación.

La primera ejecución del código compilado suele tardar más, ya que incluye el tiempo de compilación, pero las ejecuciones posteriores son significativamente más rápidas. Para un rendimiento óptimo en diferentes escenarios, torch.compile ofrece varios modos como "default", "reduce-overhead" (que utiliza CUDA graphs para reducir aún más la sobrecarga) y "max-autotune" (que realiza un autotuning exhaustivo para encontrar los mejores kernels para tu modelo).

Usar torch.compile con Accelerate

Accelerate proporciona TorchDynamoPlugin para una integración fácil y sin problemas de torch.compile en tus scripts de entrenamiento.

from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)

Es compatible con todas las demás funciones y complementos de Accelerate, incluida la precisión mixta, el entrenamiento distribuido (DDP, FSDP, Deepspeed), etc.

Compilación regional

En lugar de intentar compilar el modelo completo, que suele tener un espacio de problemas de optimización muy grande, la compilación regional apunta a bloques repetidos de la misma clase y los compila de forma secuencial para aprovechar la caché del compilador. Por ejemplo, en GPT2LMHeadModel, el bloque/clase repetido es GPT2Block, y se puede acceder a él como model.transformer.h[0]. El resto del modelo (p. ej. model.lm_head) se compila por separado.

Esto nos permite acelerar el sobrecoste de compilación / arranque en frío de modelos como LLMs y Transformers en general. Consulte https://pytorch.org/tutorials/recipes/regional_compilation.html para más detalles.

Cómo usar la compilación regional

Se puede habilitar configurando use_regional_compilation=True en la configuración de TorchDynamoPlugin:

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)

También puede utilizar la utilidad accelerate.utils.compile_regions directamente de la misma manera que usaría torch.compile.

Ventajas de la compilación regional

Hemos realizado amplias pruebas de rendimiento que comparan la compilación completa y la compilación regional utilizando la función torch.compile en PyTorch. Los resultados completos están disponibles en el repositorio de accelerate. Las principales conclusiones de nuestras pruebas de rendimiento son:

1. Rendimiento comparable: la compilación regional ofrece mejoras de velocidad de rendimiento similares a las de la compilación completa, especialmente para modelos más grandes.
2. Compilación más rápida: La compilación regional reduce significativamente el tiempo necesario para compilar modelos, lo que la convierte en una opción más eficiente para la implementación.
3. Impacto del tamaño del lote: La diferencia de rendimiento entre las estrategias de compilación disminuye con tamaños de lote más grandes, lo que indica que el sobrecosto de la compilación tiene menos impacto en esos escenarios.
4. Consideraciones sobre el tamaño del modelo: Los beneficios de la compilación regional son más notorios en modelos más grandes, donde el ahorro en el tiempo de compilación puede ser considerable.
5. Aplicación práctica: Para aplicaciones en el mundo real, la compilación regional es una opción práctica para optimizar los tiempos de inicio en frío del entrenamiento, especialmente al trabajar con modelos grandes.

Conclusión

Tanto la compilación completa como la regional pueden acelerar significativamente tus modelos. La compilación regional ofrece un equilibrio práctico entre el tiempo de compilación y el rendimiento en tiempo de ejecución, especialmente para el entrenamiento de modelos grandes con tamaños de lote sustanciales.