Haz ZeroGPU Spaces más rápido con la compilación ahead-of-time (AoT) de PyTorch

Descripción general

ZeroGPU permite utilizar el potente hardware Nvidia H200 en Hugging Face Spaces sin reservar la GPU todo el tiempo. Aunque es eficiente para demos, no siempre se aprovecha al máximo CUDA. La compilación ahead-of-time (AoT) de PyTorch soluciona esto exportando un modelo compilado una vez y recargándolo instantáneamente en nuevos procesos, para demos más rápidas y una experiencia más suave. En la práctica, AoT ofrece mejoras de rendimiento entre 1,3× y 1,8× en modelos como Flux.1-dev, Wan y LTX cuando se usa con ZeroGPU Spaces. AoT es especialmente valiosa para el flujo de ejecución de ZeroGPU, que se basa en procesos que se crean y terminan con la demanda de la GPU. En lugar de compilar en cada ejecución (lo que es costoso cuando los procesos se crean para cada tarea de GPU), AoT precompila y reutiliza el gráfico optimizado. El artículo describe cómo incorporar AoT en ZeroGPU Spaces, con trucos avanzados como cuantización FP8 y formas dinámicas, y ofrece demos disponibles para probar. ZeroGPU utiliza MIG slices en H200 (3g.71gb disponibles; 7g.141gb llegará a finales de 2025). El enfoque mantiene el programa ligero en memoria y asegura que la GPU se asigne solo cuando es necesario.

Características clave

AoT exporta y recarga: compila un modelo una vez y reutiliza el programa compilado a través de múltiples procesos y tareas, reduciendo la sobrecarga de inicio.
Integración con Spaces: usa el paquete spaces con el decorador GPU para activar AoT dentro de ZeroGPU.
Herramientas AoT: spaces.aoti_capture para capturar entradas de ejemplo, spaces.aoti_compile para generar el binario AoT, y spaces.aoti_apply para parchear el pipeline de forma segura.
Núcleo de PyTorch: AoT se apoya en torch._inductor.aot_compile y requiere un PyTorch ExportedProgram compilado.
Cuantización: la cuantización FP8 en pos-entrenamiento puede combinarse con AoT para ganancias adicionales (requisito: compute capability 9.0+).
Formas dinámicas: la exportación puede configurarse para tratar ciertas entradas como dinámicas, permitiendo manejar resoluciones de imagen/video variables.
Modelo por resolución: para generación de video muy dinámica (Wan), puedes compilar un modelo por resolución manteniendo los parámetros compartidos.
Compatibilidad FA3: Flash Attention 3 puede usarse en flujos AoT para mejoras de rendimiento.
Hardware: ZeroGPU funciona con H200s, con opciones MIG como 3g.71gb disponibles; otras tallas están en ruta.
Demostraciones de rendimiento: un ejemplo muestra ~1,7× de mejora en Flux.1-dev; ganancias de 1,3×–1,8× según el modelo.

El artículo enlaza a TorchAO y tutoriales de PyTorch AoTInductor para más contexto.

Casos de uso comunes

Demos en vivo y showcases en Hugging Face Spaces, cuando la actividad del usuario es dispersa y la GPU está ociosa costosa.
Tareas de generación de imágenes y videos (familias Flux y Wan) donde la latencia de inferencia es crítica para demos interactivas.
Flujo que requiere inicialización rápida de modelos compilados entre spaces/tareas, evitando costos de compilación por tarea.
Flujo que combina AoT con cuantización FP8 y formas dinámicas para maximizar el throughput en hardware H200.
Generación de video multi-resolución con despacho del modelo compilado correcto manteniendo parámetros compartidos.

Configuración e instalación

El artículo describe cómo habilitar AoT dentro de ZeroGPU Spaces usando el paquete spaces y las herramientas de Torch Inductor. El flujo implica interceptar entradas, exportar el módulo a un PyTorch ExportedProgram, compilarlo con AoT y parchear el transformer en el pipeline con una utilidad de patching segura. Se señala que se requiere una GPU real durante la compilación debido a que el tuning depende del hardware. Pasos clave (enfoque conceptual):

Envolver el cálculo en una función @spaces.GPU para gestionar el dispositivo y activar el ciclo de vida Just-In-Time de ZeroGPU.
Usar spaces.aoti_capture para interceptar entradas de ejemplo del transformer.
Exportar el transformer con torch.export.export a un PyTorch ExportedProgram.
Compilar el programa exportado con spaces.aoti_compile, que envuelve torch._inductor.aot_compile y gestiona guardado/carga.
Aplicar el parche con spaces.aoti_apply para sustituir forward del transformer por el modelo compilado y liberar parámetros antiguos de la memoria.
Ejecutar la inferencia dentro del contexto Spaces.GPU para asegurar la compatibilidad con el hardware y beneficiarse del gráfico AoT compilado. Nota: se indica que comandos exactos dependen del entorno, y que la compilación requiere una GPU real.

# Comandos exactos no se proporcionan en la fuente. Consulta el artículo para el flujo completo.
# Ejemplos de marcadores (no literales de la fuente):
# 1. from spaces import GPU, aoti_capture, aoti_compile, aoti_apply
# 2. @GPU
# def run(pipe):
# with aoti_capture(pipe.transformer) as call:
# pass
# compiled = aoti_compile(call)
# aoti_apply(pipe, compiled)
# # ejecutar inferencia aquí

Inicio rápido (visión general ejecutable)

Identifica el componente transformer pesado en tu pipeline de difusión/generación.
Enmarca el código de demostración en una función Spaces.GPU para activar el comportamiento JIT de ZeroGPU.
Captura entradas con aoti_capture, exporta con torch.export.export, compila con aoti_compile y aplica con aoti_apply.
Ejecuta una inferencia de prueba para verificar ganancias y asegurar que se está utilizando el camino compilado. El artículo señala que el transformer compilado mostró mejoras de ~1,7× en Flux.1-dev y que la cuantización FP8 puede aportar ~1,2× adicionales cuando corresponde.

Ventajas y desventajas

Ventajas
Reducción significativa de la latencia de inicio/compilación para ZeroGPU Spaces.
Flujo claro entre exportación del modelo, compilación y ejecución en tiempo de ejecución.
Capacidad de combinar AoT con cuantización FP8 y formas dinámicas para ganancias adicionales.
Compatibilidad con FA3 para más aceleración en modelos de generación.
Funciona con el modelo de asignación bajo demanda de ZeroGPU, maximizando la eficiencia de recursos.
Desventajas
Requiere una GPU real durante la compilación debido al acoplamiento al hardware.
Algunos modelos pueden necesitar variantes compiladas por resolución para resultados óptimos.
Configurar AoT implica varias etapas (captura, exportación, compilación, aplicación) que deben mantenerse sincronizadas entre spaces.

Alternativas (comparación breve)

| Enfoque | Ventajas | Desventajas |---|---|---| | AoT con herramientas spaces.aoti_* | Recarga rápida entre procesos; evita compilación por tarea; admite formas dinámicas y FP8 | Requiere GPU para la compilación; integración entre varias APIs; necesita separar componentes |PyTorch torch.compile (JIT) | Ruta más simple en entornos estándar; compila a primera ejecución y cachea en el filesystem | En ZeroGPU, cada tarea lanza un nuevo proceso, haciendo menos efectivo el cache del FS; frío inicial persistente |Inferencia sin AoT | Simple; sin complejidad de export/compile | Mayor latencia de arranque en ZeroGPU |