diffusers/main/ko/optimization/pruna.md

Pruna

Pruna는 추론을 가속화하고 메모리 사용량을 줄이기 위해 양자화, 프루닝, 캐싱, 컴파일 등 다양한 최적화 방법을 제공하는 모델 최적화 프레임워크입니다. 최적화 방법의 일반적인 개요는 아래에 나와 있습니다.

기법	설명	속도	메모리	품질
batcher	여러 입력을 한 그룹으로 묶어 동시에 처리함으로써 계산 효율성을 높이고 처리 시간을 줄입니다.	✅	❌	➖
cacher	계산의 중간 결과를 저장하여 후속 연산의 속도를 높입니다.	✅	➖	➖
compiler	특정 하드웨어에 대한 지침으로 모델을 최적화합니다.	✅	➖	➖
distiller	더 크고 복잡한 모델을 모방하도록 더 작고 단순한 모델을 학습합니다.	✅	✅	❌
quantizer	가중치와 활성화의 정밀도를 낮추어 메모리 요구량을 줄입니다.	✅	✅	❌
pruner	중요도가 낮거나 중복된 연결과 뉴런을 제거하여 더 희소하고 효율적인 네트워크를 생성합니다.	✅	✅	❌
recoverer	압축 후 모델의 성능을 복원합니다.	➖	➖	✅
factorizer	Factorization는 여러 작은 행렬 곱셈을 하나의 큰 융합된 연산으로 배치합니다.	✅	➖	➖
enhancer	노이즈 제거 또는 업스케일링과 같은 후처리 알고리즘을 적용하여 모델 출력을 향상시킵니다.	❌	-	✅

✅ (개선), ➖ (대략 동일), ❌ (악화)

Pruna 문서에서 최적화 기법의 전체 범위를 살펴보세요.

설치

다음 명령어로 Pruna를 설치합니다.

pip install pruna

Diffusers 모델 최적화

아래에 표시된 바와 같이 Diffusers 모델에는 다양한 최적화 알고리즘이 지원됩니다.

아래 예시에서는 factorizer, compiler, cacher 알고리즘의 조합을 사용하여 black-forest-labs/FLUX.1-dev를 최적화합니다. 이 조합은 추론 속도를 최대 4.2배까지 가속화하고, 피크 GPU 메모리 사용량을 34.7GB에서 28.0GB로 줄여주며, 출력 품질은 거의 동일하게 유지합니다.

이 예제에서 사용된 최적화 기법에 대해 자세히 알아보려면 Pruna 최적화 문서를 참조하세요.

최적화 알고리즘을 사용하여 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를 제공합니다.

총 시간과 처리량과 같이 관심 있는 metrics와 평가할 데이터셋을 정의할 수 있습니다. 모델을 정의하여 EvaluationAgent에 전달할 수 있습니다.

EvaluationAgent를 사용하여 최적화된 모델을 로드하고 평가한 다음, 이를 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")

최적화된 모델을 기본 모델과 비교하는 대신, 단독 diffusers 모델을 평가할 수도 있습니다. 이는 최적화 없이 모델의 성능을 평가하려는 경우에 유용합니다. PrunaModel 래퍼를 사용하고 해당 래퍼에서 EvaluationAgent를 실행하여 이를 수행할 수 있습니다.

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)! 블로그 글과 InferBench Space를 참고하세요.

참고

• Pruna
• Pruna 최적화
• Pruna 평가
• Pruna 튜토리얼