Pruna

Pruna adalah kerangka kerja optimasi model yang menawarkan berbagai metode optimasi - kuantisasi, pemangkasan, penyangga, kompilasi - untuk mempercepat inferensi dan mengurangi penggunaan memori. Gambaran umum dari metode-metode optimasi tersebut ditunjukkan di bawah ini.

✅ (memperbaiki), ➖ (sekitar sama), ❌ (memburuk)

Jelajahi berbagai metode optimasi dalam dokumentasi Pruna.

Instalasi

Install Pruna dengan perintah berikut.

pip install pruna

Optimalkan model Diffusers

Berbagai algoritma optimasi didukung untuk model Diffusers seperti yang ditunjukkan di bawah ini.

Contoh di bawah ini mengoptimalkan black-forest-labs/FLUX.1-dev dengan kombinasi algoritma factorizer, compiler, dan cacher. Kombinasi ini mempercepat inferensi hingga 4,2x dan memangkas penggunaan memori GPU puncak dari 34,7GB menjadi 28,0GB, sambil tetap mempertahankan kualitas output yang hampir sama.

Lihat dokumentasi Pruna optimization untuk mempelajari lebih lanjut tentang teknik optimasi yang digunakan dalam contoh ini.

Mulai dengan mendefinisikan SmashConfig dengan algoritma optimasi yang akan digunakan. Untuk mengoptimalkan model, bungkus pipeline dan SmashConfig dengan smash lalu gunakan pipeline seperti biasa untuk inferensi.

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]

Setelah optimasi, kita dapat membagikan dan memuat model yang telah dioptimalkan menggunakan 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")

Evaluasi dan uji kinerja model Diffusers

Pruna menyediakan EvaluationAgent untuk mengevaluasi kualitas model yang telah Anda optimalkan.

Kita dapat metrik yang kita pedulikan, seperti total waktu dan throughput, serta dataset untuk dievaluasi. Kita dapat mendefinisikan model dan melewatkan ke 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)

Sekarang setelah Anda melihat cara mengoptimalkan dan mengevaluasi model Anda, Anda dapat mulai menggunakan Pruna untuk mengoptimalkan model Anda sendiri. Untungnya, kami memiliki banyak contoh untuk membantu Anda memulai.

Untuk detail lebih lanjut tentang benchmarking Flux, lihatlah blog post Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)! dan Space InferBench.

Referensi

• Pruna
• Optimasi Pruna
• Evaluasi Pruna
• Tutorial Pruna