Pruna

Pruna, çıkarımı hızlandırmak ve bellek kullanımını azaltmak için kuantizasyon, budama, önbellekleme ve derleme gibi çeşitli optimizasyon yöntemleri sunan bir model optimizasyon çerçevesidir. Optimizasyon yöntemlerinin genel bir özeti aşağıda gösterilmektedir.

✅ (iyileştirir), ➖ (yaklaşık aynı), ❌ (kötüleştirmir)

Pruna belgelerinde optimizasyon yöntemlerinin tam kapsamını keşfedin.

Kurulum

Pruna'yı aşağıdaki komutla yükleyin.

pip install pruna

Diffusers modellerini optimize et

Aşağıda gösterildiği gibi, Diffusers modelleri için geniş bir yelpazede optimizasyon algoritmaları desteklenmektedir.

Aşağıdaki örnek, black-forest-labs/FLUX.1-dev modelini factorizer, compiler ve cacher algoritmalarının bir kombinasyonuyla optimize eder. Bu kombinasyon, çıkarımı (inference) en fazla 4.2x hızlandırırken tepe GPU bellek kullanımını 34.7GB'den 28.0GB'a düşürür ve tüm bunları neredeyse aynı çıktı kalitesini koruyarak gerçekleştirir.

Bu örnekte kullanılan optimizasyon teknikleri hakkında daha fazla bilgi edinmek için Pruna optimizasyonu belgelerine başvurun.

Optimizasyon algoritmalarını kullanmak için bir SmashConfig tanımlayarak başlayın. Modeli optimize etmek için pipeline'ı ve SmashConfig'i smash ile sarmalayın ve ardından pipeline'ı çıkarım için normal şekilde kullanı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]

Optimizasyonun ardından, optimize edilmiş modeli Hugging Face Hub üzerinden paylaşabilir ve yükleyebiliriz.

# 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 modellerini değerlendirin ve performanslarını ölçün

Pruna, optimize ettiğiniz modellerin kalitesini değerlendirmek için EvaluationAgent sağlar.

Toplam süre ve işleme hızı gibi önem verdiğimiz metrikleri ve değerlendirme yapılacak veri setini seçebiliriz. Bir model tanımlayıp EvaluationAgent nesnesine iletebiliriz.

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)

Modellerinizi nasıl optimize edeceğinizi ve değerlendireceğinizi gördüğünüze göre, artık Pruna'yı kendi modellerinizi optimize etmek için kullanmaya başlayabilirsiniz. Neyse ki, başlamanıza yardımcı olacak birçok örneğimiz var.

Flux için benchmarking hakkında daha fazla ayrıntı için Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)! blog yazısına ve InferBench Space'ine göz atın.

Referans

• Pruna
• Pruna optimizasyonu
• Pruna değerlendirme
• Pruna eğitimleri