Faça seus ZeroGPU Spaces trabalharem rápido com compilação ahead-of-time do PyTorch

TL;DR

• O ZeroGPU no Hugging Face Spaces permite acesso sob demanda a GPUs Nvidia H200 para apps de demonstração, evitando reservas de GPU o tempo todo.
• A compilação ahead-of-time (AoT) do PyTorch permite exportar um modelo compilado e recarregá-lo instantaneamente, reduzindo o tempo de cold-start em tarefas do ZeroGPU.
• Espera-se ganhos de ~1,3× a 1,8× em modelos como Flux.1-dev, Wan e LTX; a quantização FP8 agrega ~1,2× a mais em casos adequados.
• A implementação usa spaces.aoti_capture, spaces.aoti_compile e spaces.aoti_apply dentro de um fluxo @spaces.GPU; a compatibilidade com FA3 pode aumentar ainda mais o desempenho.
• Usuários Pro e equipes/Empresas podem criar Spaces ZeroGPU; todos podem usar, com cota 8× maior para Pro/Equipe/Empresa. Fonte.

Contexto e antecedentes

O Hugging Face Spaces oferece uma forma simples de publicar apps de demonstração de ML, mas a execução tradicional com .to(‘cuda’) pode exigir que uma GPU fique reservada durante toda a vida útil do Space, mesmo quando o tráfego está ocioso ou é esparso. O PyTorch inicializa o driver NVIDIA quando CUDA é acionado pela primeira vez, o que tende a manter o processo preso ao CUDA até ser liberado. O ZeroGPU contorna isso forkando o processo, conectando-o ao CUDA apenas quando necessário e liberando-o quando não é mais utilizado, gerando um uso de GPU mais alinhado ao tráfego real. O ZeroGPU utiliza o pacote Spaces e pode reservar um slice de um H200 MIG (3g.71gb) hoje, com a opção de migs maiores (7g.141gb) planejada para o fim de 2025. Esse contexto de hardware é relevante porque frameworks modernos como PyTorch e JAX suportam compilação ahead-of-time (AoT) para otimizar gráficos de computação para velocidade e latência. Em ZeroGPU, AoT difere do torch.compile tradicional, que compila na primeira execução e armazena em cache no sistema de arquivos. Como as tarefas do Space costumam ser criadas rapidamente para cada tarefa de GPU, o AoT oferece uma alternativa mais estável à reutilização de compilação entre processos. O post enfatiza como AoT pode ajudar exportando um programa compilado uma vez e recarregando-o instantaneamente entre processos, reduzindo sobrecarga de framework e eliminando tempos de inicialização a frio. O Flux.1-dev é usado como modelo de exemplo para demonstrar como habilitar AoT no ZeroGPU Spaces, com foco no componente transformer, que costuma ser a parte mais computacionalmente intensiva. A abordagem AoT no ZeroGPU envolve uma sequência de passos baseada em utilitários do pacote spaces e nas capacidades AoT do PyTorch. O objetivo é reduzir a sobrecarga por tarefa de GPU e manter os parâmetros aprendidos estáveis.

O que há de novo

O post apresenta um padrão prático para habilitar AoT em ZeroGPU Spaces, incluindo:

• aoti_capture: usado como gerenciador de contexto para interceptar chamadas do transformer (por exemplo, pipe.transformer) e registrar entradas de exemplo que serão usadas na exportação.
• torch.export.export: exporta o gráfico do transformer e seus parâmetros para um PyTorch ExportedProgram, representando o gráfico de computação e os valores dos parâmetros.
• aoti_compile: wrapper leve ao redor de torch._inductor.aot_compile que compila o programa exportado e cuida da salvamento e do carregamento preguioso do artefato compilado.
• aoti_apply: ajusta o pipeline para que forward do transformer apontem para o modelo compilado e garante que parâmetros antigos sejam liberados da memória para evitar OOM.
• @spaces.GPU: todo o fluxo de captura–export–compile–aplicar é envolvido dentro de uma função spaces.GPU para assegurar a calibração de hardware em CUDA real, não apenas em emulação. O resultado é um transformer compilado que pode ser usado para inferência com muito menor sobrecarga por tarefa. O blog cita um ganho de cerca de 1,7× para o transformador Flux.1-Dev com código adicional modesto. Para ampliar a aplicabilidade, são discutidos caminhos para demonstrar demos com Fluxo e Wan, incluindo exemplos como zerogpu-aoti-multi.py, que mostra como compilar um modelo por resolução em cenários de entrada dinâmica para vídeos. FA3 (Flash-Attention 3) é indicado como compatível com AoT, e a integração com FA3 pode trazer ganhos adicionais. Para ampliar ainda mais, a quantização FP8 pode ser integrada ao fluxo AoT usando APIs de TorchAO, lembrando que FP8 requer uma capacidade de cálculo CUDA ≥ 9.0. Como os ZeroGPU utilizam H200, o FP8 é compatível e pode entregar ganhos adicionais de velocidade. A seção sobre formas dinâmicas descreve como exportar com inputs dinâmicos e ajustar um mapa de dimensões para as entradas, replicando estruturas com PyTorch tree_map. Dependendo do modelo, as entradas que não precisam de dinamicidade podem receber None.

Por que isso importa (impacto para desenvolvedores/empresas)

• Menor desperdício de GPU ociosa: AoT reduz a necessidade de manter o CUDA ativo apenas para tarefas realmente exigidas, tornando o uso de GPU mais eficiente.
• Demos mais rápidas e experiência interativa: AoT pode reduzir significativamente o tempo de inicialização e a latência, aprimorando a experiência de demonstração.
• Alinhamento com hardware: a abordagem é projetada para Nvidia H200 com suporte a quantização FP8 e otimizações CUDA.
• Eficiência operacional para equipes: exportar e recarregar gráficos compilados facilita cenários de várias tarefas e usuários, potencialmente reduzindo custos com GPUs em nuvem.
• Compatibilidade com fluxos PyTorch existentes: AoT aproveita conceitos familiares do PyTorch (ExportedProgram, Inductor) com integração aos recursos do Hugging Face Spaces via o pacote spaces. Para equipes que produzem apps de demonstração, a combinação de AoT com quantização FP8 e compatibilidade FA3 oferece um caminho para maior taxa de transferência e experiência suave do usuário, sem comprometer a integridade dos parâmetros aprendidos. Consulte o post original para orientações e exemplos de código mais detalhados.

Detalhes técnicos ou Implementação

O fluxo AoT no ZeroGPU Spaces depende de passos concretos e utilitários do pacote spaces:

• aoti_capture: captura entradas de exemplos sem realmente executar o transformer, para que possamos compilar com dados representativos.
• torch.export.export: exporta o programa do transformer para PyTorch ExportedProgram.
• aoti_compile: wrapper que gerencia a compilação com torch._inductor.aot_compile, salvando e carregando o modelo compilado quando necessário.
• aoti_apply: aplica o modelo compilado ao pipeline substituindo pipe.transformer.forward e liberando parâmetros antigos para evitar uso excessivo de memória.
• @spaces.GPU: a sequência toda de captura–export–compile–aplicar é executada dentro de uma função spaces.GPU, garantindo que a compilação ocorra com hardware CUDA real. Essa sequência gera um transformer compilado que pode ser carregado e usado para inferência com menor sobrecarga de computação. A métrica mencionada é cerca de 1,7× de ganho para Flux.1-Dev. Dinâmica de formas e quantização:
• Quantização FP8: a quantização FP8 pode ser usada dentro do fluxo AoT via TorchAO, com ganhos adicionais de velocidade. FP8 requer compute capability ≥ 9.0, condição atendida por GPUs H200.
• Formas dinâmicas: durante a exportação, é possível configurar entradas dinâmicas e mapear dimensões dinâmicas para as entradas relevantes, permitindo que o AoT aceite variações de tamanho de entrada.
• Estratégia multi-resolução: para cenários com grande variabilidade (ex.: geração de vídeo Wan), pode-se compilar um modelo por resolução e despachar o modelo adequado em tempo de execução. Integração com FA3:
• FA3 funciona com AoT; compilar FA3 a partir do código pode levar minutos, mas, uma vez disponível, pode trazer ganhos adicionais ao acelerar as operações de atenção durante a inferência. Modelos e demonstrações:
• O Flux.1-dev é usado como modelo de referência para demonstrar o fluxo AoT de compilação de transformer e os ganhos práticos.
• O exemplo zerogpu-aoti-multi.py mostra como compilar modelos dedicados para diferentes resoluções de vídeo, mantendo parâmetros compartilhados.
• O blog aponta para demonstrações prontas no espaço zerogpu-aoti para explorar AoT em Spaces. Notas de implementação e limitações:
• O AoT requer um contexto real de GPU (via @spaces.GPU) para obter otimizações específicas de hardware; a emulação CUDA não fornece tuning adequado.
• AoT reduz a necessidade de compilação repetida entre tarefas e evita depender do cache do filesystem para uso rápido entre processos.
• Como qualquer otimização, os ganhos reais variam com o modelo; no Flux.1-Dev, o ganho reportado é de cerca de 1,7×, e FP8 pode oferecer ganhos adicionais conforme o caso.

Perguntas frequentes (FAQ)

• Q: O que é AoT no contexto de ZeroGPU Spaces? A: Ahead-of-time compilation exporta um gráfico pronto para uso e o recarrega instantaneamente entre processos, reduzindo o overhead de inicialização.
• Q: Qual é a magnitude de ganho com AoT em Flux.1-Dev e outros modelos? A: Ganhos de aproximadamente 1,3×–1,8× são comuns, com FP8 adicionando cerca de 1,2× em cenários apropriados.
• Q: Como a quantização FP8 se encaixa no fluxo AoT? A: FP8 é suportada via TorchAO dentro do fluxo AoT, exigindo compute capability ≥ 9.0, com ganhos adicionais de desempenho.
• Q: E quanto à dinâmica de formas e várias resoluções? A: Entradas dinâmicas podem ser especificadas na exportação; para alta variação (ex.: Wan), pode-se compilar por resolução e despachar a versão apropriada em tempo de execução.
• Q: Quem pode usar ZeroGPU AoT-enabled Spaces? A: Usuários Pro e equipes/Empresas podem criar Spaces; qualquer um pode usar, com cota 8× maior para Pro/Equipe/Empresa.

Referências

• https://huggingface.co/blog/zerogpu-aoti