Inferência rápida de LoRA para Flux com Diffusers e PEFT

LoRA adapters fornecem uma grande flexibilidade para modelos de todos os formatos e tamanhos. No contexto de geração de imagens, eles permitem estilos diferentes, personagens distintos e muito mais. Às vezes, também podem ser usados para reduzir a latência de inferência. Assim, sua importância é crucial, especialmente ao personalizar e ajustar modelos. Neste post, tomamos o modelo Flux.1-Dev para geração de imagem a partir de texto por sua popularidade e adoção generalizadas, e mostramos como otimizar a velocidade de inferência usando LoRA (~2,3x). O Flux.1-Dev possui mais de 30 mil adapters treinados com ele (conforme reportado na plataforma Hub da Hugging Face). Portanto, sua importância para a comunidade é significativa. Observe que, embora demonstremos ganhos de velocidade com Flux, nossa crença é de que nossa receita é genérica o suficiente para ser aplicada a outros modelos. Se você não quiser esperar pelo código, confira o repositório de código acompanhante. Ao servir LoRAs, é comum realizar hotswapping (trocar LoRAs diferentes) entre si. Um LoRA altera a arquitetura do modelo base. Além disso, LoRAs podem ser diferentes entre si – cada um pode ter ranks diferentes e alvos de camadas diferentes para adaptação. Para acomodar essas propriedades dinâmicas das LoRAs, é necessário adotar medidas para garantir que as otimizações aplicadas sejam robustas. Por exemplo, podemos aplicar o torch.compile a um modelo carregado com uma LoRA específica para obter ganhos de velocidade na inferência. No entanto, assim que trocarmos a LoRA por outra (com uma configuração potencialmente diferente), encontraremos problemas de recompilação, o que reduz a velocidade de inferência. Também é possível fundir os parâmetros LoRA nos parâmetros base do modelo, executar a compilação e desfazer a fusão ao carregar novas LoRAs. No entanto, essa abordagem volta a enfrentar o problema da recompilação a cada inferência, devido a mudanças em nível de arquitetura. Nossa receita de otimização considera os cenários acima para ser o mais realista possível. A seguir, os componentes-chave da nossa receita de otimização: Observe que, entre os itens acima, a quantização FP8 é loss em termos de precisão, mas frequentemente oferece o trade-off mais eficaz entre velocidade e memória. Embora tenhamos testado a receita principalmente com GPUs NVIDIA, ela deve funcionar em GPUs AMD também. Em posts anteriores de blog ( post 1 e post 2 ), já discutimos os benefícios de usar os três primeiros componentes da nossa receita de otimização. Aplicá-los um a um envolve apenas algumas linhas de código: O processador FA3 vem de aqui . Os problemas surgem quando tentamos trocar LoRAs em um transformador de difusão compilado ( pipe.transformer ) sem acionar recomplição. Normalmente, carregar e descarregar LoRAs exigiria recompilação, o que anularia qualquer vantagem de velocidade obtida pela compilação. Felizmente, existe uma forma de evitar a recompilação. Ao passar hotswap=True , o diffusers deixará a arquitetura do modelo inalterada e apenas troca os pesos do adaptador LoRA, o que não requer recompilação. (Como lembrete, a primeira chamada de pipe será lenta, pois o torch.compile é um compilador just-in-time. No entanto, as chamadas subsequentes devem ser significativamente mais rápidas.) Isso, em geral, permite a troca de LoRAs sem recompilação, mas há limitações: para mais informações sobre hotswapping no Diffusers e suas limitações, consulte a seção de hotswapping da documentação. Os benefícios desse fluxo tornam-se evidentes quando observamos a latência de inferência sem usar compilação com hotswapping. Perguntas rápidas : A receita de otimização discutida até agora pressupõe acesso a uma GPU poderosa como a H100. No entanto, o que fazer quando temos apenas GPUs de consumo, como a RTX 4090? Vamos descobrir. Flux.1-Dev (sem LoRA) usando o tipo de dados Bfloat16 requer ~33GB de memória para rodar. Dependendo do tamanho do módulo LoRA e sem usar nenhuma otimização, esse footprint de memória pode aumentar ainda mais. Ao longo do restante desta seção, consideramos uma máquina RTX 4090 como nosso banco de testes. Primeiro, para permitir a execução ponta a ponta do Flux.1-Dev, podemos aplicar offloading de CPU, movendo componentes que não são necessários para a computação atual para a CPU, liberando mais memória no acelerador. Fazendo isso, conseguimos rodar todo o pipeline em ~22GB em 35,403 segundos em uma RTX 4090. Ativar a compilação pode reduzir a latência para 31,205 segundos (1,12x de ganho). Em termos de código, são apenas algumas linhas: Note que não aplicamos a quantização FP8 aqui porque ela não é suportada com offloading de CPU e compilação (thread de suporte). Portanto, apenas aplicar quantização FP8 ao Flux Transformer não é suficiente para mitigar o problema de exaustão de memória, também. Nessa situação, decidimos removê-lo. Assim, para aproveitar a quantização FP8, precisamos encontrar uma forma de fazê-lo sem offloading da CPU. Para Flux.1-Dev, se quantizarmos também o codificador de texto T5, devemos ser capazes de carregar e executar todo o pipeline em 24GB. Abaixo, segue uma comparação dos resultados com e sem quantização do codificador de texto T5 (NF4 da bitsandbytes). Como podemos observar, quantizar o codificador T5 não traz perda de qualidade significativa. Combinando o codificador de texto T5 quantizado e o Flux Transformer quantizado com FP8 e torch.compile, obtemos resultados razoáveis – 9,668 segundos frente a 32,27 segundos (aproximadamente 3,3x de ganho) sem queda perceptível de qualidade. É possível gerar imagens com 24 GB de VRAM mesmo sem quantizar o codificador T5, mas isso tornaria o pipeline de geração um pouco mais complicado. Agora temos uma forma de executar todo o pipeline Flux.1-Dev com quantização FP8 em uma RTX 4090. Podemos aplicar a receita de otimização previamente estabelecida para otimizar a inferência LoRA no mesmo hardware. Como FA3 não é suportado na RTX 4090, manteremos a receita de otimização a seguir com a quantização do T5 adicionada recentemente. Na tabela abaixo, apresentamos os números de latência de inferência com diferentes combinações dos componentes acima aplicados. Observações rápidas : Para habilitar hotswapping sem acionar recompilação, dois obstáculos precisam ser superados. Primeiro, o fator de escala LoRA precisa ser convertido em tensores do PyTorch a partir de floats, o que é relativamente simples. Segundo, o formato dos pesos LoRA precisa ser preenchido com o maior formato requerido. Deste modo, os dados nos pesos podem ser substituídos sem reatribuir todo o atributo. É por isso que o argumento max_rank é crucial. Ao preenchermos os valores com zeros, os resultados permanecem inalterados, embora a computação fique mais lenta dependendo do tamanho do preenchimento. Como não são adicionados novos atributos LoRA, também é exigido que cada LoRA após a primeira só aponte para as mesmas camadas, ou um subconjunto, das camadas-alvo da primeira LoRA. Assim, escolha a ordem de carregamento com sabedoria. Se LoRAs atingirem camadas deslocadas, existe a possibilidade de criar uma LoRA fictícia que aponte para a união de todas as camadas-alvo. Para ver os detalhes desta implementação, consulte o arquivo hotswap.py no PEFT. Este post descreve uma receita de otimização para inferência rápida de LoRA com Flux, demonstrando ganhos significativos de velocidade. Nossa abordagem combina Flash Attention 3, torch.compile e quantização FP8, mantendo a capacidade de hotswapping sem recompilação. Em GPUs de alto desempenho como a H100, essa configuração otimizada oferece aproximadamente 2,23x de ganho em relação ao baseline. Para GPUs de consumo, especificamente a RTX 4090, lidamos com limitações de memória introduzindo quantização NF4 no codificador de texto T5 e recorrendo à compilação regional. Esta receita abrangente alcançou um ganho substancial de ~2,04x, tornando a inferência LoRA no Flux viável e eficiente mesmo com memória VRAM limitada. A visão-chave é que, ao gerenciar cuidadosamente a compilação e a quantização, os benefícios do LoRA podem ser plenamente realizados em diferentes configurações de hardware. Esperamos que as receitas deste post inspirem você a otimizar seus casos de uso baseados em LoRA, beneficiando-se de inferência mais rápida. Abaixo está uma lista dos recursos importantes citados ao longo deste post: Mais artigos do nosso blog Por linoyts 2 de janeiro de 2024 • 70 Por derekl35 19 de junho de 2025 • 83 Gostei do seu post! Posso perguntar qual versão do PyTorch você utilizou para os resultados de teste? PyTorch nightly. A boy with baloon · Inscreva-se ou faça login para comentar