Pruna

Pruna là một framework tối ưu hóa mô hình cung cấp nhiều phương pháp tối ưu hóa khác nhau - lượng tử hóa, cắt tỉa, lưu trữ tạm, biên dịch - để tăng tốc suy luận và giảm mức sử dụng bộ nhớ. Tổng quan chung về các phương pháp tối ưu hóa được hiển thị bên dưới.

✅ (cải thiện), ➖ (xấp xỉ như nhau), ❌ (suy giảm)

Khám phá đầy đủ các phương pháp tối ưu hóa trong tài liệu Pruna.

Cài đặt

Cài đặt Pruna bằng lệnh sau.

pip install pruna

Tối ưu hóa các mô hình Diffusers

Một loạt các thuật toán tối ưu hóa được hỗ trợ cho các mô hình Diffusers như được hiển thị bên dưới.

Ví dụ bên dưới tối ưu hóa black-forest-labs/FLUX.1-dev bằng cách kết hợp các thuật toán factorizer, compiler và cacher. Sự kết hợp này tăng tốc suy luận lên tới 4,2 lần và giảm mức sử dụng bộ nhớ GPU cao nhất từ 34,7GB xuống còn 28,0GB, trong khi vẫn duy trì chất lượng đầu ra gần như không thay đổi.

Tham khảo tài liệu tối ưu hóa Pruna để tìm hiểu thêm về các kỹ thuật tối ưu hóa được sử dụng trong ví dụ này.

Bắt đầu bằng cách định nghĩa một SmashConfig với các thuật toán tối ưu hóa cần sử dụng. Để tối ưu hóa mô hình, bọc pipeline và SmashConfig bằng smash rồi sử dụng pipeline như bình thường cho suy luận.

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]

Sau khi tối ưu hóa, chúng ta có thể chia sẻ và tải mô hình đã được tối ưu hóa bằng cách sử dụng 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")

Đánh giá và đo hiệu năng các mô hình Diffusers

Pruna cung cấp EvaluationAgent để đánh giá chất lượng các mô hình đã được tối ưu của bạn.

Chúng ta có thể theo dõi các chỉ số quan tâm, chẳng hạn như tổng thời gian và thông lượng, cũng như bộ dữ liệu để đánh giá. Chúng ta có thể định nghĩa một mô hình và truyền nó vào 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)

Giờ đây, khi bạn đã biết cách tối ưu và đánh giá các mô hình của mình, bạn có thể bắt đầu sử dụng Pruna để tối ưu các mô hình của riêng mình. May mắn thay, chúng tôi có nhiều ví dụ để giúp bạn bắt đầu.

Để biết thêm chi tiết về việc benchmark Flux, hãy xem bài đăng trên blog Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)! và Space InferBench.

Tham khảo

• Pruna
• Tối ưu hóa Pruna
• Đánh giá Pruna
• Các hướng dẫn của Pruna