Deixe os ZeroGPU Spaces mais rápidos com compilação ahead-of-time (AoT) do PyTorch

Visão geral

ZeroGPU permite usar hardware Nvidia H200 poderoso no Hugging Face Spaces sem manter a GPU ocupada o tempo todo. Embora seja eficiente para demonstrações, nem sempre o CUDA é aproveitado ao máximo. A compilação ahead-of-time (AoT) do PyTorch resolve isso ao exportar um modelo compilado para recarregar instantaneamente em novos processos, resultando em demonstrações mais rápidas e experiência mais suave. Na prática, a AoT oferece ganhos de performance entre 1,3× e 1,8× em modelos como Flux.1-dev, Wan e LTX quando usada em ZeroGPU Spaces. A AoT é especialmente valiosa para o modelo de execução de ZeroGPU, que é baseado em tarefas rápidas e criação de processos sob demanda. Em vez de compilar a cada execução (o que é custoso, pois os processos são criados para cada tarefa de GPU), a AoT pré-compila e reutiliza o gráfico otimizado. Isso torna viável usar as ferramentas TorchInductor de alto desempenho em um Space sem atrasos de inicialização a frio. O artigo descreve como integrar AoT aos ZeroGPU Spaces, com truques avançados como quantização FP8 e formas dinâmicas, oferecendo demos funcionais para experimentar. Os MIG slices do H200 (3g.71gb) estão atuais (7g.141gb deverá chegar no fim de 2025). A abordagem mantém a memória do programa leve e garante que a GPU seja alocada apenas quando necessário.

Principais recursos

AoT exporta e recarrega: Compile um modelo uma vez e reutilize o programa compilado em diferentes processos/task, reduzindo a sobrecarga de inicialização.
Integração com Spaces: Utilize o pacote spaces com o decorador GPU para habilitar o fluxo AoT no ZeroGPU.
Ferramentas AoT: Use spaces.aoti_capture para capturar entradas de exemplo, spaces.aoti_compile para gerar o binário AoT, e spaces.aoti_apply para ajustar o pipeline com segurança.
Infra PyTorch: AoT usa torch._inductor.aot_compile e depende de um PyTorch ExportedProgram compilado.
Quantização: A quantização FP8 de pós-treinamento pode ser combinada com AoT para ganhos adicionais (requisito: compute capability 9.0+).
Formas dinâmicas: A exportação pode ser configurada para tratar determinadas entradas como formas dinâmicas, permitindo lidar com resoluções de imagem/vídeo variadas.
Modelo por resolução: Para geração de vídeo altamente dinâmica (Wan), é possível compilar um modelo por resolução mantendo os parâmetros compartilhados.
Compatibilidade FA3: Flash Attention 3 pode ser usado dentro de fluxos AoT para mais melhorias de desempenho.
Hardware: ZeroGPU roda em H200s, com opções MIG como 3g.71gb disponíveis (7g.141gb virá em breve).
Demonstrações de velocidade: Um exemplo mostra melhoria de cerca de 1,7× no Flux.1-dev; ganhos de 1,3×–1,8× são citados para outros modelos.

O artigo também aponta links para TorchAO e tutoriais PyTorch AoTInductor para mais contexto.

Casos de uso comuns

Demonstrações ao vivo e showcases em Hugging Face Spaces, onde a atividade do usuário é esporádica e o tempo ocioso da GPU é caro.
Tarefas de geração de imagem e vídeo (por exemplo, famílias Flux e Wan) onde a latência é crítica para demonstrações interativas.
Cenários que exigem inicialização rápida de modelos compilados entre espaços/task, evitando o overhead de compilação por tarefa.
Fluxos que combinam AoT com quantização FP8 e formas dinâmicas para maximizar o throughput em hardware H200.
Geração de vídeo com várias resoluções, despachando o modelo previamente compilado correto mantendo parâmetros compartilhados.

Configuração & instalação

O artigo descreve como habilitar AoT dentro de ZeroGPU Spaces aproveitando o pacote spaces e as ferramentas de PyTorch Inductor. O fluxo envolve interceptar entradas, exportar o módulo para um PyTorch ExportedProgram, compilar com AoT, e substituir o transformer na pipeline com a utilidade de patching de forma segura. Observa-se que uma GPU real é necessária durante a compilação devido ao ajuste dependente de hardware. Etapas-chave (conceituais):

Envolver a computação alvo dentro de uma função com @spaces.GPU para garantir gerenciamento de dispositivo e para facilitar o ciclo de vida Just-In-Time do ZeroGPU.
Usar spaces.aoti_capture para interceptar entradas de exemplificação do transformer.
Exportar o transformer para um PyTorch ExportedProgram usando torch.export.export.
Compilar o programa exportado com spaces.aoti_compile, que envolve torch._inductor.aot_compile e gerenciamento de salvamento/carregamento.
Aplicar patch com spaces.aoti_apply para substituir forward do transformer pela versão compilada e liberar parâmetros antigos da memória.
Executar inferência dentro da função GPU para garantir compatibilidade de hardware e benefício do gráfico AoT compilado. Observação: o texto enfatiza que comandos exatos são específicos do ambiente e que é necessário uso de GPU real durante a compilação.

# Comandos exatos não são fornecidos na fonte. Veja o artigo para o fluxo completo.
# Exemplos de placeholder (não citados literalmente na fonte):
# 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)
# # rodar inferência aqui

Início rápido (visão geral executável)

Identifique o componente transformer pesado no seu pipeline de difusão ou geração (ex.: bloco de denoising/transformer no Flux.1-dev).
Envolva o código de demonstração em uma função Spaces.GPU para ativar o comportamento Just-In-Time do ZeroGPU.
Capture entradas com aoti_capture, exporte com torch.export.export, compile com aoti_compile e aplique com aoti_apply.
Rode uma inferência de teste para verificar os ganhos de velocidade e garantir que o caminho compilado está sendo utilizado. O artigo aponta que o transformer compilado resultou em cerca de 1,7× de melhoria em Flux.1-dev e que quantização FP8 pode adicionar cerca de 1,2× a mais quando apropriado.

Vantagens e desvantagens

Vantagens
Reduções significativas na latência de inicialização/compile para ZeroGPU Spaces.
Separação clara entre exportação do modelo, compilação e execução em tempo de run.
Possibilidade de combinar com quantização FP8 e formas dinâmicas para ganhos adicionais.
Compatibilidade com FA3 para obter mais acelerações em modelos de geração.
Funciona com a estratégia de alocação sob demanda do ZeroGPU, maximizando a eficiência de recursos.
Desvantagens
Requer GPU real durante a compilação, devido ao tuning dependente de hardware.
Alguns modelos podem exigir variantes compiladas por resolução para melhores resultados.
A configuração envolve várias etapas (captura, exportação, compilação, aplicação) que precisam estar sincronizadas entre spaces.

Alternativas (comparação rápida)

| Abordagem | Prós | Contras |---|---|---| | AoT com as ferramentas spaces.aoti_* | Recarga rápida entre processos; evita compilação por execução; suporta formas dinâmicas e FP8 | Requer GPU durante a compilação; integração envolve várias APIs; precisa de separação de componentes |PyTorch torch.compile (JIT) | Caminho mais simples em ambientes padrão; compila na primeira execução e armazena no filesystem | Em ZeroGPU, cada tarefa inicia um novo processo, tornando o cache de FS menos eficaz; atraso de cold-start |Inferência não AoT | Simples; sem complexidade de export/compile | Maior overhead de inicialização em ZeroGPU |