Pruna

Pruna — это фреймворк оптимизации моделей, предлагающий различные методы оптимизации — квантование, обрезание, кэширование, компиляцию — для ускорения вывода и снижения потребления памяти. Общий обзор методов оптимизации приведен ниже.

✅ (улучшение), ➖ (примерно то же), ❌ (ухудшение)

Ознакомьтесь с полным перечнем методов оптимизации в документации Pruna.

Установка

Установите Pruna с помощью следующей команды.

pip install pruna

Оптимизация моделей Diffusers

Для моделей Diffusers поддерживается широкий спектр алгоритмов оптимизации, как показано ниже.

Пример ниже оптимизирует black-forest-labs/FLUX.1-dev с помощью комбинации алгоритмов factorizer, compiler и cacher. Эта комбинация ускоряет инференс до 4,2 раз и снижает пиковое потребление памяти GPU с 34,7 ГБ до 28,0 ГБ, при этом практически не изменяя качество вывода.

См. документацию Pruna optimization, чтобы узнать больше об оптимизационных техниках, использованных в этом примере.

Начните с определения SmashConfig с алгоритмами оптимизации, которые вы хотите использовать. Чтобы оптимизировать модель, оберните пайплайн и SmashConfig в smash, а затем используйте пайплайн для инференса как обычно.

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]

После оптимизации мы можем поделиться оптимизированной моделью и загрузить её с помощью 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")

Оценка и бенчмаркинг моделей Diffusers

Pruna предоставляет EvaluationAgent для оценки качества ваших оптимизированных моделей.

Мы можем указать метрики, которые нас интересуют, например, общее время и пропускную способность, а также набор данных для оценки. Мы можем определить модель и передать её в EvaluationAgent.

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")

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)

Теперь, когда вы узнали, как оптимизировать и оценивать свои модели, вы можете начать использовать Pruna для оптимизации собственных моделей. К счастью, у нас есть множество примеров, которые помогут вам начать работу.

Для получения подробной информации о бенчмаркинге Flux ознакомьтесь с публикацией в блоге Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)! и Space InferBench.

Справочник

• Pruna
• Оптимизация Pruna
• Оценка Pruna
• Учебные материалы Pruna