diffusers/main/es/optimization/pruna.md

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.

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.

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.

Podemos cargar y evaluar un modelo optimizado usando el EvaluationAgent y pasarlo al 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")

En lugar de comparar el modelo optimizado con el modelo base, también puede evaluar el modelo de diffusers de forma independiente. Esto es útil si desea evaluar el rendimiento del modelo sin la optimización. Puede hacerlo utilizando el envoltorio PrunaModel y ejecutando EvaluationAgent sobre él.

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)

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

• Pruna
• Optimización de Pruna
• Evaluación de Pruna
• Tutoriales de Pruna