Pruna

Pruna es un marco de optimización de modelos que ofrece diversos métodos de optimización: cuantización, poda, caché y compilación, para acelerar la inferencia y reducir el uso de memoria. A continuación se muestra una visión general de los métodos de optimización.

Técnica	Descripción	Velocidad	Memoria	Calidad
`batcher`	Agrupa varias entradas para procesarlas simultáneamente, mejorando la eficiencia computacional y reduciendo el tiempo de procesamiento.	✅	❌	➖
`cacher`	Almacena los resultados intermedios de los cálculos para acelerar las operaciones posteriores.	✅	➖	➖
`compiler`	Optimiza el modelo con instrucciones para hardware específico.	✅	➖	➖
`distiller`	Entrena un modelo más pequeño y sencillo para imitar a un modelo más grande y complejo.	✅	✅	❌
`quantizer`	Reduce la precisión de los pesos y las activaciones, lo que disminuye los requisitos de memoria.	✅	✅	❌
`pruner`	Elimina conexiones y neuronas menos importantes o redundantes, lo que da como resultado una red más dispersa y eficiente.	✅	✅	❌
`recoverer`	Restaura el rendimiento de un modelo después de la compresión.	➖	➖	✅
`factorizer`	La factorización agrupa varias pequeñas multiplicaciones de matrices en una única operación grande fusionada.	✅	➖	➖
`enhancer`	Mejora la salida del modelo aplicando algoritmos de posprocesamiento, como la reducción de ruido o el aumento de escala.	❌	-	✅

✅ (mejora), ➖ (aproximadamente igual), ❌ (empeora)

Explore el rango completo de métodos de optimización en la documentación de Pruna.

Instalación

Instala Pruna con el siguiente comando.

pip install pruna

Optimizar modelos de Diffusers

Se admite una amplia gama de algoritmos de optimización para los modelos de Diffusers, como se muestra a continuación.

Overview of the supported optimization algorithms for diffusers models

El ejemplo siguiente optimiza black-forest-labs/FLUX.1-dev con una combinación de los algoritmos factorizer, compiler y cacher. Esta combinación acelera la inferencia hasta 4,2 veces y reduce el uso de memoria GPU pico de 34,7 GB a 28,0 GB, todo ello manteniendo prácticamente la misma calidad de salida.

Consulte la documentación de optimización de Pruna para obtener más información sobre las técnicas de optimización utilizadas en este ejemplo.

Optimization techniques used for FLUX.1-dev showing the combination of factorizer, compiler, and cacher algorithms

Comience definiendo un SmashConfig con los algoritmos de optimización a utilizar. Para optimizar el modelo, envuelva el pipeline y el SmashConfig con smash y luego use el pipeline de forma habitual para la inferencia.

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]

Tras la optimización, podemos compartir y cargar el modelo optimizado usando 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")

Evaluar y medir el rendimiento de los modelos de Diffusers

Pruna proporciona el EvaluationAgent para evaluar la calidad de tus modelos optimizados.

Podemos especificar las métricas que nos interesan, como el tiempo total y el rendimiento, y el conjunto de datos sobre el que evaluar. Podemos definir un modelo y pasarlo al EvaluationAgent.

modelo optimizado

modelo independiente

Ahora que has visto cómo optimizar y evaluar tus modelos, puedes empezar a usar Pruna para optimizar los tuyos. Afortunadamente, tenemos muchos ejemplos para ayudarte a empezar.

Para más detalles sobre el benchmarking de Flux, consulta la publicación de blog Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)! y el Space InferBench.

Referencia

Update on GitHub