diffusers/v0.43.0/pt/optimization/pruna.md

Pruna

Pruna é um framework de otimização de modelos que oferece diversos métodos de otimização — quantização, poda, cache, compilação — para acelerar a inferência e reduzir o uso de memória. Uma visão geral dos métodos de otimização é mostrada abaixo.

Técnica	Descrição	Velocidade	Memória	Qualidade
batcher	Agrupa várias entradas para serem processadas simultaneamente, melhorando a eficiência computacional e reduzindo o tempo de processamento.	✅	❌	➖
cacher	Armazena resultados intermediários de cálculos para acelerar operações subsequentes.	✅	➖	➖
compiler	Otimiza o modelo com instruções para hardware específico.	✅	➖	➖
distiller	Treina um modelo menor e mais simples para imitar um modelo maior e mais complexo.	✅	✅	❌
quantizer	Reduz a precisão dos pesos e das ativações, diminuindo os requisitos de memória.	✅	✅	❌
pruner	Remove conexões e neurônios menos importantes ou redundantes, resultando em uma rede mais esparsa e eficiente.	✅	✅	❌
recoverer	Restaura o desempenho de um modelo após a compressão.	➖	➖	✅
factorizer	A fatorização agrupa várias multiplicações de matrizes pequenas em uma única operação grande e fundida.	✅	➖	➖
enhancer	Melhora a saída do modelo aplicando algoritmos de pós-processamento, como remoção de ruído ou aumento de escala.	❌	-	✅

✅ (melhora), ➖ (aproximadamente o mesmo), ❌ (piora)

Explore a gama completa de métodos de otimização na documentação do Pruna.

Instalação

Instale o Pruna com o seguinte comando.

pip install pruna

Otimize modelos do Diffusers

Uma ampla variedade de algoritmos de otimização é suportada para modelos do Diffusers, conforme mostrado abaixo.

O exemplo abaixo otimiza black-forest-labs/FLUX.1-dev com uma combinação dos algoritmos factorizer, compiler e cacher. Essa combinação acelera a inferência em até 4,2x e reduz o uso de memória de GPU de pico de 34,7GB para 28,0GB, mantendo praticamente a mesma qualidade de saída.

Consulte a documentação de otimização do Pruna para saber mais sobre as técnicas de otimização usadas neste exemplo.

Comece definindo um SmashConfig com os algoritmos de otimização a serem usados. Para otimizar o modelo, envolva o pipeline e o SmashConfig com smash e, em seguida, use o pipeline normalmente para inferência.

import torch
from diffusers import FluxPipeline

from pruna import PrunaModel, SmashConfig, smash

# load the model
# Try segmind/Segmind-Vega or black-forest-labs/FLUX.1-schnell with a small GPU memory
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

# define the configuration
smash_config = SmashConfig()
smash_config["factorizer"] = "qkv_diffusers"
smash_config["compiler"] = "torch_compile"
smash_config["torch_compile_target"] = "module_list"
smash_config["cacher"] = "fora"
smash_config["fora_interval"] = 2

# for the best results in terms of speed you can add these configs
# however they will increase your warmup time from 1.5 min to 10 min
# smash_config["torch_compile_mode"] = "max-autotune-no-cudagraphs"
# smash_config["quantizer"] = "torchao"
# smash_config["torchao_quant_type"] = "fp8dq"
# smash_config["torchao_excluded_modules"] = "norm+embedding"

# optimize the model
smashed_pipe = smash(pipe, smash_config)

# run the model
smashed_pipe("a knitted purple prune").images[0]

Após a otimização, podemos compartilhar e carregar o modelo otimizado usando o Hugging Face Hub.

# save the model
smashed_pipe.save_to_hub("<username>/FLUX.1-dev-smashed")

# load the model
smashed_pipe = PrunaModel.from_hub("<username>/FLUX.1-dev-smashed")

Avaliar e fazer benchmark de modelos do Diffusers

A Pruna fornece o EvaluationAgent para avaliar a qualidade dos seus modelos otimizados.

Podemos definir as métricas que nos importam, como o tempo total e a taxa de transferência, e o conjunto de dados para avaliação. Podemos definir um modelo e passá-lo para o EvaluationAgent.

Podemos carregar e avaliar um modelo otimizado usando o EvaluationAgent e passá-lo para o Task.

import torch
from diffusers import FluxPipeline

from pruna import PrunaModel
from pruna.data.pruna_datamodule import PrunaDataModule
from pruna.evaluation.evaluation_agent import EvaluationAgent
from pruna.evaluation.metrics import (
    ThroughputMetric,
    TorchMetricWrapper,
    TotalTimeMetric,
)
from pruna.evaluation.task import Task

# define the device
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
smashed_pipe = PrunaModel.from_hub("PrunaAI/FLUX.1-dev-smashed")

# Define the metrics
metrics = [
    TotalTimeMetric(n_iterations=20, n_warmup_iterations=5),
    ThroughputMetric(n_iterations=20, n_warmup_iterations=5),
    TorchMetricWrapper("clip"),
]

# Define the datamodule
datamodule = PrunaDataModule.from_string("LAION256")
datamodule.limit_datasets(10)

# Define the task and evaluation agent
task = Task(metrics, datamodule=datamodule, device=device)
eval_agent = EvaluationAgent(task)

# Evaluate smashed model and offload it to CPU
smashed_pipe.move_to_device(device)
smashed_pipe_results = eval_agent.evaluate(smashed_pipe)
smashed_pipe.move_to_device("cpu")

Em vez de comparar o modelo otimizado com o modelo base, você também pode avaliar o modelo diffusers isolado. Isso é útil se você quiser avaliar o desempenho do modelo sem a otimização. Podemos fazer isso usando o wrapper PrunaModel e executando o EvaluationAgent nele.

import torch
from diffusers import FluxPipeline

from pruna import PrunaModel

# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16
).to("cpu")
wrapped_pipe = PrunaModel(model=pipe)

Agora que você viu como otimizar e avaliar seus modelos, pode começar a usar o Pruna para otimizar os seus próprios modelos. Felizmente, temos muitos exemplos para ajudar você a começar.

Para mais detalhes sobre o benchmarking do Flux, confira o post no blog Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)! e o Space InferBench.

Referência

• Pruna
• Otimização do Pruna
• Avaliação do Pruna
• Tutoriais do Pruna