Hugging Face acelera ZeroGPU con compilación AoT en PyTorch

ZeroGPU es la solución de Hugging Face para ejecutar demos en Spaces sin mantener un GPU ocupado todo el tiempo. ¿El problema? Las demos que generan imágenes o video sufren con arranques fríos porque la compilación en caliente tarda demasiado. Hugging Face propone usar la compilación ahead of time de PyTorch para preparar binarios optimizados una sola vez y recargarlos al vuelo, logrando respuestas más rápidas y demos más fluidas. (huggingface.co)

¿Qué es ZeroGPU y por qué importa?

ZeroGPU permite que una Space use GPUs poderosas como las H200 solo cuando hay tráfico. En lugar de reservar un GPU permanentemente, ZeroGPU crea procesos cortos que inicializan CUDA, ejecutan la tarea y cierran, lo que evita tener hardware ocioso cuando no hay usuarios. Esto es ideal para demos públicas con tráfico intermitente. (huggingface.co)

Nota: la característica está disponible para usuarios Pro y equipos, aunque cualquiera puede usar las ZeroGPU Spaces creadas por esos usuarios. (huggingface.co)

¿Qué aporta la compilación ahead of time (AoT)?

PyTorch ofrece dos vías de compilación: torch.compile para compilación en tiempo de ejecución y un flujo AoT usando torch.export + AOTInductor para compilar antes de servir. En entornos donde los procesos GPU son efímeros como ZeroGPU, torch.compile no puede reaprovechar bien su trabajo, mientras que AoT permite exportar y recargar un modelo ya optimizado casi instantáneamente. Eso reduce el overhead del framework y elimina gran parte del tiempo de arranque. (huggingface.co)

Cómo lo implementan en una Space (resumen práctico)

La guía de Hugging Face muestra un flujo concreto con estos pasos:

• Capturar ejemplos de entrada del componente pesado (por ejemplo pipe.transformer).
• Exportar el componente a un ExportedProgram con torch.export.export usando esos ejemplos.
• Compilar el export con un helper spaces.aoti_compile que envuelve a torch._inductor.aot_compile.
• Aplicar el binario compilado al pipeline con spaces.aoti_apply para no dejar parámetros viejos en memoria.
• Ejecutar todo el proceso de compilación dentro de un @spaces.GPU porque la compilación depende de métricas y pruebas en el hardware real. (huggingface.co)

Ejemplo mínimo de cambios en el código:

import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe.to('cuda')

@spaces.GPU
def generate(prompt):
    return pipe(prompt).images

# Código AoT adicional (resumido):
@spaces.GPU(duration=1500)
def compile_transformer():
    with spaces.aoti_capture(pipe.transformer) as call:
        pipe("ejemplo")
    exported = torch.export.export(pipe.transformer, args=call.args, kwargs=call.kwargs)
    return spaces.aoti_compile(exported)

compiled = compile_transformer()
spaces.aoti_apply(compiled, pipe.transformer)

Este cambio son apenas unas líneas extra pero puede mejorar sensiblemente la latencia. (huggingface.co)

Cosas a tener en cuenta antes de compilar AoT

• Quantización: combinar AoT con quantización (por ejemplo FP8) ofrece ganancias adicionales, pero FP8 exige capacidad CUDA 9.0 o superior. En ZeroGPU, la H200 ya soporta estas técnicas. (huggingface.co)

• Formas dinámicas: imágenes y videos tienen resoluciones variables. Al exportar hay que marcar qué dimensiones serán dinámicas y ajustar rangos mínimos y máximos para torch.export. Eso requiere inspección y algo de prueba y error según el modelo. (huggingface.co)

• Multi-compile: si la variabilidad es muy grande (por ejemplo varios tamaños de video), una estrategia es compilar varias versiones por resolución y compartir pesos entre ellas, despachando la correcta en tiempo de ejecución. (huggingface.co)

• FlashAttention-3: FA3 es compatible y acelera aún más; Hugging Face recomienda usar su hub de kernels precompilados para evitar compilar kernels desde cero y perder tiempo de GPU. (huggingface.co)

Resultados y ejemplos que puedes probar

En los ejemplos que publicaron, las mejoras van desde aproximadamente 1.3× hasta 1.8× según el modelo y las opciones aplicadas. Con FLUX.1-dev reportaron cerca de 1.7× de mejora en latencia y, en combinaciones con FA3, se muestran casos de 1.75×. Si te interesa experimentar, Hugging Face reunió varios ejemplos en una organización llamada ZeroGPU-AOTI para probar directamente. (huggingface.co)

Puedes ver los demos y recursos aquí: ZeroGPU-AOTI en Hugging Face. También es útil revisar el tutorial de AOTInductor en la documentación de PyTorch para entender las opciones avanzadas. (huggingface.co)

¿Qué significa esto para ti?

Si publicas demos en Spaces que generan imágenes, video o cualquier tarea pesada de GPU, AoT en ZeroGPU es una palanca práctica para reducir latencia y costos. ¿Quieres una demo que responda rápido a un usuario o prefieres pagar por GPUs siempre activas? AoT te permite elegir lo primero sin sacrificar rendimiento.

Si eres desarrollador, prueba este flujo en una rama de tu Space: captura ejemplos reales, exporta, compila con spaces.aoti_compile y aplica el binario. Si estás montando una exhibición pública o un prototipo para clientes, esto hace la experiencia más profesional sin aumentar la factura por GPU inactiva. (huggingface.co)

Si quieres profundizar, la publicación original incluye código, demos y enlaces a recursos como kernels precompilados y discusiones de la comunidad. (huggingface.co)