Haz que tus ZeroGPU Spaces vayan más rápido con la compilación AoT de PyTorch

TL;DR

• ZeroGPU Spaces ahora aprovecha la compilación ahead-of-time (AoT) de PyTorch para reducir la latencia y los tiempos de arranque en tareas de GPU.
• AoT permite exportar un modelo compilado una vez y recargarlo al instante en nuevos procesos, para demos más ágiles en Hugging Face Spaces con hardware Nvidia H200.
• Se observan mejoras reales de aproximadamente 1.3×–1.8× en modelos como Flux.1-dev, Wan y LTX; la cuantización FP8 puede añadir alrededor de 1.2× de velocidad adicional.
• El flujo de trabajo utiliza spaces.aoti_capture y spaces.aoti_compile, además de spaces.aoti_apply para parchear el pipeline de forma segura sin mantener grandes parámetros del modelo en memoria.
• MIG slices actuales en ZeroGPU: 3g.71gb; se espera un slice completo de 7g.141gb para finales de 2025.
• ZeroGPU es compatible con Flash-Attention 3 (FA3) para acelerar aún más, con la compatibilidad de hardware y controladores CUDA esenciales para AoT. Para más detalles, consulta el post original de Hugging Face. Fuente

Contexto y antecedentes

Spaces es una plataforma de Hugging Face que permite a los profesionales ML publicar demos de apps con GPUs potentes. ZeroGPU evita reservar la GPU entera para cada Space, activando la GPU solo cuando hay actividad real, lo que optimiza el uso de recursos ante tráfico esporádico e irregular. El enfoque tradicional obliga a inicializar el controlador NVIDIA al ejecutar .to(‘cuda’), fijando el proceso CUDA de forma permanente. ZeroGPU adopta una aproximación de arranque bajo demanda: bifurca el proceso, configura CUDA para la tarea, ejecuta el trabajo y libera la GPU al terminar. Este comportamiento se habilita mediante el paquete spaces y el decorador GPU. La AoT ofrece una optimización adicional: exporta un programa compilado y lo reutiliza rápidamente sin necesidad de recompilar en cada ejecución, reduciendo costos y tiempos de arranque y eliminando la sobrecarga típica de la compilación just-in-time. Referencias al post original de Hugging Face sobre AoT. Fuente

Novedades

El objetivo es compilar por adelantado el componente transformer (la parte denoiser en muchos modelos generativos) porque suele ser el cuello de botella computacional. El flujo integra varios elementos del paquete spaces:

• aoti_capture: usado como gestor de contexto para interceptar llamadas del transformer y registrar entradas simuladas para inferencia.
• torch.export.export: exporta el transformer a un PyTorch ExportedProgram, que sirve como base para la compilación AoT.
• spaces.aoti_compile: envoltorio alrededor de torch._inductor.aot_compile que gestiona el guardado y la carga perezosa del modelo compilado.
• spaces.aoti_apply: parche seguro del pipeline para reemplazar forward por el modelo compilado y limpiar parámetros antiguos para evitar consumo de memoria.
• Requiere envoltura en una función @spaces.GPU: la compilación y el ajuste deben ejecutarse en hardware real para obtener resultados óptimos.
• Integración segura: parchear forward solo o cambiar atributos puede causar pérdidas de estado; spaces.aoti_apply maneja esto correctamente. En la práctica, partir del ejemplo base de ZeroGPU y con Flux.1-dev, se observa aproximadamente 1.7× de velocidad con AoT en ZeroGPU. Esto demuestra que un único ciclo de compilación AoT puede entregar beneficios notables para demos en Spaces. La cuantización FP8 puede combinarse con AoT para ganar velocidad adicional. FP8 requiere una capacidad de cómputo CUDA de al menos 9.0, compatible con GPUs H200 en ZeroGPU. TorchAO facilita este flujo, y permite explorar cuantización y formas dinámicas durante la exportación (transformer_dynamic_shapes). En Flux.1-Dev, las entradas de ejemplo y los rangos dinámicos se definen analizando las formas de hidden_states a diferentes resoluciones. En escenarios con dinamismo elevado (como Wan para generación de vídeo), puede ser eficiente compilar un modelo por resolución y despachar la versión adecuada en el momento de la ejecución. FA3 (Flash-Attention 3) es compatible con AoT, lo que añade más aceleración para cálculos de atención; construir FA3 puede llevar tiempo y dependerá del hardware. Para quien quiere practicar, hay demostraciones completas en zerogpu-aoti-multi.py y en la Space Wan 2.2. En conjunto, AoT, cuantización FP8 y formas dinámicas ofrecen una vía práctica para acelerar demos en Spaces con ZeroGPU.

Por qué importa (impacto para desarrolladores/empresas)

AoT con ZeroGPU Spaces permite demos más rápidas sin necesidad de GPUs siempre encendidas. La menor latencia de arranque beneficia especialmente a entornos de demostración donde el tráfico es intermitente. Desde el punto de vista arquitectónico, compilar luego reutilizar el artefacto AoT encaja con el objetivo de ZeroGPU de uso eficiente de recursos GPU a demanda. Las empresas y equipos que utilizan Spaces obtienen varias ventajas:

• Demos más rápidos y ciclos de iteración más cortos.
• Mejor utilización de recursos al evitar inicializaciones CUDA repetidas.
• Mayor rendimiento potencial gracias a la cuantización FP8 y a las formas dinámicas.
• Estrategias por resolución cuando el dynamismo es alto, manteniendo parámetros compartidos. Las referencias se pueden consultar en el post original de Hugging Face para más detalles. Fuente

Detalles técnicos o Implementación

La implementación descrita se apoya en pasos concretos dentro de un flujo Spaces ZeroGPU:

• Preparar entradas de ejemplo con spaces.aoti_capture para representar tensores de inferencia reales.
• Exportar el transformer a un PyTorch ExportedProgram con torch.export.export. El programa exportado sirve como base para la compilación AoT.
• Compilar con spaces.aoti_compile, que envuelve a aot_compile de PyTorch Inductor y gestiona la carga/guardado perezoso del artefacto compilado.
• Aplicar parches con spaces.aoti_apply para reemplazar forward por el modelo compilado y limpiar parámetros antiguos para evitar picos de memoria.
• Envolver todo en una función @spaces.GPU para asegurar el acceso al hardware real y el tuning específico.
• Inferencia con AoT: el transformador compilado puede reutilizarse sin reexportar, reduciendo costos por ejecución. Notas sobre características:
• Cuantización FP8 con TorchAO: acelera la generación de imágenes y vídeos aparte de AoT, siempre que el hardware soporte compute capability 9.0+.
• Formas dinámicas: exportar con transformer_dynamic_shapes para manejar entradas de distintas resoluciones.
• Wan y resoluciones múltiples: cuando el dynamismo es alto, puede ser eficiente compilar por resolución y despachar la versión correcta en tiempo real.
• FA3: Flash-Attention 3 funciona con AoT para acelerar las operaciones de atención; la construcción y ajuste pueden requerir tiempo y depender del hardware. Notas sobre hardware:
• MIG slices: ZeroGPU ofrece actualmente un MIG slice de 3g.71gb; se espera un slice completo de 7g.141gb para finales de 2025.
• Compatibilidad: AoT depende de hardware CUDA real; la emulación de CUDA fuera del bloque @spaces.GPU no es suficiente para optimizaciones completas. Si quieres empezar, los ejemplos Flux.1-Dev, Wan y LTX muestran beneficios reales con AoT. El post original de Hugging Face ofrece una visión completa del proceso. Fuente

Puntos clave

• AoT permite reutilizar de forma inmediata el modelo compilado entre procesos, reduciendo la latencia de arranque.
• La integración con aoti_capture, aoti_compile y aoti_apply facilita mejorar demos existentes en Spaces.
• Cuantización FP8 y formas dinámicas añaden beneficios y amplían escenarios de uso.
• Wan demuestra un enfoque por resolución cuando el dinamismo es alto, manteniendo parámetros compartidos.
• FA3 es compatible y puede aportar aceleraciones adicionales en la atención.

Preguntas frecuentes (FAQ)

• P: ¿Qué es AoT en este contexto?\n R: AoT es la compilación adelantada que exporta un modelo compilado y lo reutiliza instantáneamente en nuevos procesos, reduciendo costos y latencia de inferencia.
• P: ¿Cómo activo AoT en un ZeroGPU Space?\n R: Usa aoti_capture para entradas, exporta con torch.export.export, compila con spaces.aoti_compile y aplica con spaces.aoti_apply dentro de una función @spaces.GPU.
• P: ¿Qué modelos se benefician más?\n R: Modelos generativos donde el transformer es el cuello de botella, como Flux.1-dev, Wan y LTX.
• P: ¿Es compatible la cuantización FP8?\n R: Sí, con TorchAO; FP8 puede aportar ~1.2× de velocidad adicional si el hardware admite compute capability 9.0+.
• P: ¿Qué pasa con la memoria?\n R: Spaces.aoti_apply ayuda a evitar que los parámetros originales se mantengan en memoria, reduciendo el riesgo de OOM.

Referencias

• Post de Hugging Face sobre ZeroGPU AoT: https://huggingface.co/blog/zerogpu-aoti