Faça seus ZeroGPU Spaces ficarem mais rápidos com compilação AoT do PyTorch

TL;DR

• ZeroGPU Spaces agora usa compilação ahead-of-time (AoT) do PyTorch para reduzir latência e tempos de cold-start em tarefas de GPU.
• AoT permite exportar um modelo compilado e recarregá-lo instantaneamente em novos processos, proporcionando demos mais ágeis no Hugging Face Spaces com hardware Nvidia H200.
• Ganhos reais observados ficam entre 1,3× e 1,8× em modelos como Flux.1-dev, Wan e LTX; a quantização FP8 pode adicionar aproximadamente 1,2× de velocidade.
• O fluxo de trabalho usa spaces.aoti_capture, spaces.aoti_compile e spaces.aoti_apply para patch seguro do pipeline sem manter grandes parâmetros do modelo na memória.
• Atualmente, o ZeroGPU utiliza um MIG slice de 3g.71gb; o slice completo de 7g.141gb deve chegar no final de 2025.
• ZeroGPU funciona com Flash-Attention 3 (FA3) para acelerar ainda mais, com compatibilidade de hardware e driver CUDA essenciais para AoT. Para ler mais, veja o post original no Hugging Face. Fonte

Contexto e antecedentes

Spaces é uma plataforma da Hugging Face que permite aos practitioners publicar apps de demonstração com GPUs poderosas. O ZeroGPU reduz a reserva de GPU ociosa reservando apenas quando há atividade, evitando o desperdício de recursos em tráfego esparso. A abordagem tradicional envolve inicializar o driver NVIDIA ao executar .to(‘cuda’), o que fixa o processo CUDA para sempre. O ZeroGPU, por outro lado, utiliza uma abordagem de iniciação sob demanda: ele bifura o processo, configura CUDA para a tarefa, executa e libera a GPU assim que a tarefa termina. Esse comportamento é habilitado pelo pacote spaces. A AoT, porém, representa uma forma de otimizar a inferência exportando previamente o modelo compilado. Diferente do JIT tradicional, que pode exigir caches de sistema de arquivos para restaurar modelos otimizados, a AoT permite recarregar rapidamente o modelo compilado, diminuindo a sobrecarga de compilação e os tempos de inicialização. Referência ao post original da Hugging Face sobre AoT no ZeroGPU. Fonte

O que há de novo

A estratégia centra-se em compilar antecipadamente o componente transformer (a parte do denoiser em muitos modelos generativos), pois é tipicamente o gargalo de compute. O fluxo envolve várias etapas com o pacote spaces:

• aoti_capture: usado como gerenciador de contexto para interceptar chamadas do transformer, impedindo a execução e registrando entradas simuladas para inferência.
• torch.export.export: exporta o transformer para um PyTorch ExportedProgram, que é um grafo de computação com os parâmetros originais.
• spaces.aoti_compile: envoltório em torno de torch._inductor.aot_compile que gerencia salvar/ carregar o artefato compilado, gerando um compiled_transformer pronto para inferência.
• spaces.aoti_apply: patch seguro do pipeline substituindo forward do transformer pelo modelo compilado, além de liberar parâmetros antigos da memória.
• Envolvimento com @spaces.GPU: para garantir que toda a cadeia seja executada em hardware real, permitindo tuning específico de compilação.
• Integração segura: a substituição direta por compiled_transformer pode perder atributos importantes; spaces.aoti_apply resolve isso, assegurando que o pipeline funcione corretamente. Na prática, rodando a partir do Flux.1-dev e com AoT, observa-se cerca de 1,7× de melhoria de velocidade no Flux.1-Dev. Isso mostra que uma única rodada de compilação AoT já entrega ganhos perceptíveis para demos baseadas em Spaces. A quantização FP8 pode ser combinada com AoT para ganhos adicionais. FP8 requer compute capability CUDA de pelo menos 9.0, compatível com a arquitetura H200 dos ZeroGPU. O TorchAO facilita esse fluxo, com a possibilidade de quantização e configuração de formas dinâmicas durante a exportação (transformer_dynamic_shapes). Em Flux.1-Dev, as entradas de exemplo e as ranges dinâmicas são definidas com base nas formas de hidden_states sob resoluções variadas. Em cenários com dinamismo significativo (como Wan para vídeo), pode ser eficiente compilar modelos por resolução e despachar a versão correta em tempo de execução, mantendo parâmetros compartilhados. FA3 (Flash-Attention 3) também é compatível com AoT, permitindo ainda mais aceleracões nas operações de atenção, embora construir o FA3 a partir do source possa exigir tempo e tuning de hardware. Para quem quer praticidade, há demonstrações completas no zerogpu-aoti-multi.py e em implementações Wan na Space 2.2. Juntas, as capacidades AoT, quantização FP8 e formas dinâmicas oferecem um caminho prático para acelerar demos no Spaces com ZeroGPU.

Por que isso importa (impacto para desenvolvedores/empresas)

A AoT com ZeroGPU Spaces facilita demonstrações mais ágeis sem depender de GPUs sempre ligadas. A redução de latência de inicialização é particularmente valiosa para fluxos de trabalho baseados em demos, que precisam responder rapidamente a usuários. Arquiteturalmente, compilar e reutilizar o artefato AoT permite menor dependência de inicializações CUDA frequentes, o que melhora a utilização de recursos para cargas intermitentes. Empresas que dependem de Spaces podem se beneficiar de:

• Demos mais rápidas e ciclos de iteração mais curtos.
• Melhor aproveitamento de recursos, reduzindo o tempo ocioso de GPUs.
• Possibilidades de throughput maior com quantização FP8 e formas dinâmicas.
• Estratégias de várias resoluções quando o dynamismo é alto, mantendo parâmetros compartilhados. As informações de implementação vêm do post originalmente publicado pela Hugging Face; consulte a fonte para detalhes. Fonte

Detalhes técnicos ou Implementação

A implementação descrita envolve passos concretos dentro de um fluxo ZeroGPU Spaces:

• Preparar entradas de exemplo com spaces.aoti_capture para representar tensores de inferência real. Esse passo garante que as formas de entrada e os ranges sejam adequados para exportação.
• Exportar o transformer para um PyTorch ExportedProgram com torch.export.export. O programa exportado serve de base para a compilação AoT.
• Compilar com spaces.aoti_compile, que envolve aot_compile do PyTorch Inductor, com salva e carregamento preguiçoso do artefato compilado. O resultado é um compiled_transformer pronto para inferência.
• Aplicar patch seguro com spaces.aoti_apply para substituir o forward do transformer pelo modelo compilado e eliminar parâmetros antigos para evitar picos de memória.
• Envolver tudo em uma função @spaces.GPU para garantir que a compilação se beneficie do hardware real e do tuning baseado em micro-benchmarks.
• Inferência com AoT: o transformador compilado é reintroduzido no pipeline, permitindo reutilização instantânea em operações subsequentes, sem novas exportações. Desempenho e escolhas de design:
• Dinâmica de formas e quantização FP8: a combinação de AoT com FP8 pode entregar ganhos adicionais, especialmente em geração de imagens e vídeos. FP8 requer compute capability 9.0+.
• Formas dinâmicas: use transformer_dynamic_shapes para especificar quais entradas são dinâmicas. Em Flux.1-Dev, definiu-se um intervalo de resoluções latentes para orientar a exportação.
• Wan e várias resoluções: quando o dinamismo é alto, pode-se compilar por resolução e despachar a versão apropriada conforme necessário.
• FA3: Flash-Attention 3 funciona com AoT para acelerar mecanismos de atenção; o processo de compilação pode demorar alguns minutos e dependerá do hardware. Observações sobre hardware e capacidades:
• MIG slices: o ZeroGPU atualmente aloca um MIG slice de 3g.71gb; full slice de 7g.141gb chegará no fim de 2025.
• Compatibilidade: toda a pipeline de AoT depende de hardware CUDA real. Em resumo, a emulação de CUDA fora do @spaces.GPU não é suficiente para a compilação realmente otimizada. Para quem quiser começar, os exemplos Flux.1-Dev, Wan e LTX mostram ganhos de desempenho práticos com AoT. A leitura do post original da Hugging Face oferece visão detalhada do método. Fonte

Perguntas frequentes (FAQ)

• Q: O que é AoT neste contexto?\n A: AoT é a compilação antecipada que exporta um modelo compilado e o recarrega instantaneamente em novos processos, reduzindo sobrecarga e latência de inicialização.
• Q: Como habilito AoT em um ZeroGPU Space?\n A: Execute aoti_capture para capturar entradas, exporte com torch.export.export, compile com spaces.aoti_compile e aplique com spaces.aoti_apply dentro de uma função @spaces.GPU.
• Q: Quais modelos mais se beneficiam com AoT?\n A: Modelos generativos cujo componente transformer/denoiser é o gargalo (exemplos citados: Flux.1-dev, Wan, LTX) veem ganhos significativos de inferência.
• Q: É possível usar quantização FP8 com AoT?\n A: Sim. A quantização FP8 com TorchAO pode trazer ganhos adicionais, desde que a máquina suporte compute capability 9.0+.
• Q: Quais são considerações de memória?\n A: A aplicação com spaces.aoti_apply evita manter parâmetros originais em memória, reduzindo o risco de OOM durante a inferência.

Referências

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