추론 가속

Diffusion models are slow at inference because generation is an iterative process where noise is gradually refined into an image or video over a certain number of “steps”. To speedup this process, you can try experimenting with different schedulers, reduce the precision of the model weights for faster computations, use more memory-efficient attention mechanisms, and more.

이러한 기법들을 결합하여 함께 사용하면, 단일 기법만 사용할 때보다 추론 속도를 더 빠르게 할 수 있습니다.

이 가이드에서는 추론을 가속하는 방법을 다룹니다.

모델 데이터 유형

모델 가중치의 정밀도와 데이터 유형은 추론 속도에 영향을 미칩니다. 정밀도가 높을수록 가중치를 로드하는 데 더 많은 메모리가 필요하고, 연산을 수행하는 데 더 많은 시간이 소요되기 때문입니다. PyTorch는 기본적으로 모델 가중치를 float32 또는 전체 정밀도로 로드하므로, 데이터 유형을 변경하는 것은 추론 속도를 빠르게 하는 간단한 방법입니다.

import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

import torch
from diffusers import StableDiffusionXLPipeline

torch.backends.cuda.matmul.allow_tf32 = True

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

스케일된 도트 프로덕트 어텐션

메모리 효율적 어텐션은 추론 속도 및 메모리 사용량을 최적화합니다!

Scaled dot product attention (SDPA)는 FlashAttention, xFormers 및 네이티브 C++ 구현 등 여러 attention backends를 구현합니다. 하드웨어에 가장 최적화된 backends를 자동으로 선택합니다.

PyTorch >= 2.0을 사용하는 경우 SDPA가 기본적으로 활성화되며, 코드에 추가적인 변경 사항은 필요하지 않습니다. 원하시는 경우 다른 attention backends를 실험하여 직접 선택할 수 있습니다. 아래 예제는 torch.nn.attention.sdpa_kernel 컨텍스트 매니저를 사용하여 효율적인 attention을 활성화합니다.

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
  image = pipeline(prompt, num_inference_steps=30).images[0]

torch.compile

torch.compile은 PyTorch 코드와 연산을 최적화된 커널로 컴파일하여 추론을 가속화합니다. Diffusers는 일반적으로 UNet, transformer, VAE와 같이 계산 집약적인 모델을 컴파일합니다.

최대 속도를 위해 다음 컴파일러 설정을 활성화하세요(추가 옵션은 전체 목록을 참조하세요).

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

UNet과 VAE를 로드하고 컴파일합니다. 선택할 수 있는 여러 가지 모드가 있지만, "max-autotune"은 CUDA graph로 컴파일하여 가장 빠른 속도를 최적화합니다. CUDA graphs는 단일 CPU 연산을 통해 여러 GPU 연산을 실행함으로써 오버헤드를 효과적으로 줄여줍니다.

PyTorch 2.3.1에서는 torch.compile의 캐싱 동작을 제어할 수 있습니다. 이는 여러 컴파일 플래그에 대해 그리드 서치를 수행하여 최적의 구성을 찾는 "max-autotune"와 같은 컴파일 모드에서 특히 유용합니다. torch.compile의 Compile Time Caching 튜토리얼에서 자세히 알아보세요.

메모리 레이아웃을 channels_last로 변경하면 메모리와 추론 속도도 최적화됩니다.

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

컴파일은 첫 번째 실행 시 느리지만, 한 번 컴파일된 후에는 훨씬 빠릅니다. 동일한 유형의 추론 작업에만 컴파일된 파이프라인을 사용하도록 하십시오. 다른 이미지 크기로 컴파일된 파이프라인을 호출하면 컴파일이 다시 트리거되며, 이는 느리고 비효율적입니다.

동적 형상 컴파일

더 나은 지원을 위해 항상 PyTorch의 nightly 버전을 사용하도록 하십시오.

torch.compile은 입력 형상과 조건을 추적하며, 이들이 다르면 모델을 다시 컴파일합니다. 예를 들어, 1024x1024 해상도의 이미지로 컴파일된 모델이 다른 해상도의 이미지에 사용되면 재컴파일이 트리거됩니다.

재컴파일링을 피하려면 dynamic=True를 추가하여 조건이 변경될 때 재컴파일링을 피할 수 있는 더 동적인 커널을 생성하도록 시도해 보세요.

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.unet = torch.compile(
    pipeline.unet, fullgraph=True, dynamic=True
)

use_duck_shape=False를 지정하면 컴파일러가 동일한 크기의 입력을 동일한 기호 변수로 표현할지 여부를 결정합니다. 자세한 내용은 이 댓글을 확인하세요.

모든 모델이 기본적으로 동적 컴파일의 이점을 누릴 수 있는 것은 아니며, 일부 모델은 변경이 필요할 수 있습니다. AuraFlowPipeline 구현을 개선하여 동적 컴파일의 이점을 누릴 수 있도록 한 이 PR을 참고하세요.

Diffusers 모델에서 동적 컴파일이 예상대로 작동하지 않는 경우, 자유롭게 이슈를 열어 주세요.

지역별 컴파일

Regional compilation은 모델의 작고 반복적으로 실행되는 블록, 일반적으로 transformer layer만 컴파일하여 콜드 스타트 지연 시간을 줄이고, 이후 모든 발생 위치에서 컴파일된 산출물을 재사용할 수 있게 합니다. 많은 diffusion 아키텍처의 경우, 이는 전체 그래프 컴파일과 동일한 런타임 속도 향상을 제공하면서 컴파일 시간을 8~10배 줄여줍니다.

아래와 같이 transformer 모델과 같은 컴포넌트에 torch.compile을 래핑하는 헬퍼인 compile_repeated_blocks() 메서드를 사용합니다.

# pip install -U diffusers
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

# compile only the repeated transformer layers inside the UNet
pipeline.unet.compile_repeated_blocks(fullgraph=True)

새 모델에 대한 지역 컴파일을 활성화하려면, 컴파일할 블록의 클래스 이름(문자열)을 포함하는 _repeated_blocks 속성을 모델 클래스에 추가하세요.

class MyUNet(ModelMixin):
    _repeated_blocks = ("Transformer2DModel",)  # ← compiled by default

지역별 컴파일 예제가 더 필요하면 PR을 참고하세요.

Accelerate에는 모델에서 컴파일할 후보 블록을 자동으로 선택하는 compile_regions 메서드도 있습니다. 나머지 그래프는 별도로 컴파일됩니다. 컴파일할 블록을 지정하거나 컴파일 플래그를 조정할 수 있는 옵션이 많지 않으므로, 빠른 실험에 유용합니다.

# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

compile_repeated_blocks()는 의도적으로 명시적으로 설계되었습니다. _repeated_blocks에 반복할 블록을 나열하면 헬퍼는 해당 블록만 컴파일합니다. 이는 한 줄의 코드로 캐시 재사용에 대한 예측 가능한 동작과 쉬운 추론을 제공합니다.

그래프 브레이크

fullgraph=True를 torch.compile에 지정하는 것이 중요합니다. 이를 통해 기반 모델에 그래프 브레이크가 발생하지 않도록 보장할 수 있으며, 성능 저하 없이 torch.compile의 이점을 활용할 수 있습니다. UNet과 VAE의 경우, 반환 변수에 접근하는 방식이 달라집니다.

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

GPU 동기화

step() 함수는 denoiser가 예측을 할 때마다 scheduler에서 호출되며, sigmas 변수는 인덱싱됩니다. GPU에 배치되면 CPU와 GPU 간의 통신 동기화 때문에 지연이 발생합니다. denoiser가 이미 컴파일된 경우 이 지연이 더 뚜렷하게 나타납니다.

일반적으로 통신 동기화 및 지연을 피하기 위해 sigmas는 CPU에 유지되어야 합니다.

Refer to the torch.compile and Diffusers: A Hands-On Guide to Peak Performance blog post for maximizing performance with torch.compile for diffusion models.

벤치마크

컴파일된 파이프라인의 추론 지연 시간 및 메모리 사용량 데이터를 확인하려면 diffusers/benchmarks 데이터셋을 참조하세요.

diffusers-torchao 저장소에는 Flux와 CogVideoX의 컴파일된 버전 벤치마킹 결과도 포함되어 있습니다.

커널

Kernels는 Hub에서 최적화된 compute kernel을 빌드, 배포 및 로드하기 위한 라이브러리입니다. RMSNorm, GEGLU, RoPE, AdaLN과 같은 연산을 위한 attention kernel 및 사용자 정의 CUDA kernel을 지원합니다.

Diffusers Pipeline Integration 가이드는 커널을 add cuda-kernels 스킬과 통합하는 방법을 보여줍니다. 이 스킬을 사용하면 Claude나 Codex와 같은 에이전트가 특정 모델과 사용자의 하드웨어를 대상으로 하는 커스텀 커널을 작성할 수 있습니다.

Install the add cuda-kernels skill to teach an agent how to write a kernel. The Custom kernels for all from Codex and Claude blog post covers this in more detail.

예를 들어, add cuda-kernels 스킬로 생성된 커스텀 RMSNorm 커널에 torch.compile을 적용하면 H100에서 LTX-Video 생성 속도가 1.43배 빨라집니다.

동적 양자화

Dynamic quantization은 정밀도를 낮추어 더 빠른 수학 연산을 가능하게 함으로써 추론 속도를 향상시킵니다. 이 특정 양자화 방식은 고정된 스케일링 계수를 사용하는 대신 런타임의 데이터에 기반하여 활성화(activation)의 스케일링 방식을 결정합니다. 그 결과, 스케일링 계수가 데이터에 더 정확하게 맞춰집니다.

아래 예시에서는 torchao 라이브러리를 사용하여 UNet과 VAE에 동적 int8 양자화를 적용합니다.

Diffusers의 torchao 통합 사용법에 대해 자세히 알아보려면 torchao 문서를 참조하세요.

최대 속도를 위해 컴파일러 태그를 구성합니다.

import torch
from torchao import apply_dynamic_quant
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

동적 양자화에서 이점을 얻지 못하는 UNet과 VAE의 일부 선형 레이어를 dynamic_quant_filter_fn으로 필터링합니다.

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

Fused projection matrices

fuse_qkv_projections 메서드는 실험적이며 지원 범위가 주로 Stable Diffusion 파이프라인으로 제한됩니다. 다른 파이프라인에서 이를 활성화하는 방법에 대해서는 이 PR을 참고하세요.

Q, K, V 투영 행렬로 표현되는 세 개의 부분 공간으로 입력이 투영되는 attention 블록에서, 이러한 투영은 일반적으로 개별적으로 계산됩니다. 하지만 이를 가로로 결합하여 단일 행렬로 만들고, 한 단계에서 투영을 수행할 수 있습니다. 이렇게 하면 입력 투영의 행렬 곱셈 크기가 증가하고, 양자화의 효과도 향상됩니다.

pipeline.fuse_qkv_projections()

리소스

• Presenting Flux Fast: Making Flux go brrr on H100s 블로그 포스트를 읽어 flux-fast의 레시피를 사용하여 TorchInductor와 AOTInductor와 이 모든 최적화를 결합하여 약 2.5배의 속도 향상을 달성하는 방법에 대해 자세히 알아보세요.

  이 레시피는 AMD 하드웨어와 Flux.1 Kontext Dev를 지원합니다.

• torch.compile 사용 시 성능을 극대화하려면 torch.compile 및 Diffusers: 피크 성능을 위한 실전 가이드 블로그 글을 읽어보세요.