Published on

Stable Diffusion VRAM OOM 해결 - xFormers·SDPA·VAE 타일링

Authors

서로 다른 GPU에서 Stable Diffusion을 돌리다 보면 가장 흔하게 마주치는 에러가 VRAM OOM(out of memory)입니다. 특히 해상도를 올리거나 배치 크기를 키우거나, 고해상도 업스케일·리파이너·컨트롤넷을 얹는 순간 갑자기 터집니다.

이 글은 OOM을 “무조건 옵션 몇 개 켜기”로 땜질하지 않고, 어디에서 VRAM이 폭발하는지(어텐션, VAE, 중간 피처맵, 캐시)로 나눠서 진단한 뒤 xFormers, SDPA(Scaled Dot Product Attention), VAE 타일링을 중심으로 가장 효과가 큰 처방을 정리합니다.

중간중간 시스템 트러블슈팅 관점도 섞어 설명할 텐데, 원인 분해 방식은 예를 들어 DB에서 병목을 로그로 좁히는 접근과 유사합니다. 관심 있다면 MySQL InnoDB 데드락 로그로 범인 쿼리 찾기처럼 “증상에서 원인으로 내려가는” 글도 함께 참고하면 도움이 됩니다.

VRAM OOM이 나는 대표 지점 3가지

Stable Diffusion 파이프라인에서 VRAM을 크게 잡아먹는 지점은 대략 아래 3곳입니다.

  1. U-Net 어텐션(Attention) 연산

    • 해상도가 커질수록 토큰 수가 늘고, 어텐션 메모리가 급증합니다.
    • SD1.5, SDXL 모두 여기서 가장 자주 OOM이 납니다.
  2. VAE 디코드(잠복 공간에서 RGB로 복원)

    • 최종 이미지를 픽셀 공간으로 복원할 때 큰 텐서가 생깁니다.
    • 고해상도에서 특히 위험합니다.
  3. 부가 기능이 추가하는 그래프

    • ControlNet, IP-Adapter, LoRA 다중 적용, Refiner, Highres fix 등.
    • “기본은 되는데 옵션 하나 켜면 터짐” 패턴이면 여기가 범인인 경우가 많습니다.

이제 각각을 겨냥한 처방으로 들어가겠습니다.

1) 어텐션 메모리 절감: xFormers vs SDPA

xFormers란 무엇이고 왜 OOM에 강한가

xFormers는 메모리 효율적인 어텐션 커널을 제공해, 어텐션 계산 시 필요한 중간 텐서를 줄이는 방식으로 VRAM을 아낍니다. 체감상 “해상도 한 단계 더 올릴 수 있게 해주는” 옵션으로 가장 유명합니다.

다만 환경에 따라 설치·호환 문제가 생길 수 있고(드라이버, CUDA, torch 버전), 특정 GPU에서 속도나 안정성이 다를 수 있습니다.

SDPA란 무엇이고 언제 더 좋은가

SDPA는 PyTorch 2 계열에서 제공하는 scaled_dot_product_attention 기반 최적화입니다. 내부적으로 Flash Attention 계열 커널을 선택할 수 있고, xFormers 없이도 상당한 메모리 절감을 기대할 수 있습니다.

정리하면:

  • xFormers: “옵션 하나로 OOM을 줄이는” 실전형, 다만 설치 난이도나 조합 이슈가 가끔 있음
  • SDPA: PyTorch 네이티브, 최신 torch에서 안정적이고 성능도 좋은 편

diffusers에서 SDPA 켜는 예시

아래는 diffusers 파이프라인에서 SDPA를 활성화하는 전형적인 코드입니다.

import torch
from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

# PyTorch 2 계열에서 SDPA 사용
pipe.unet.set_attn_processor("sdpa")

prompt = "a cinematic portrait, 35mm, shallow depth of field"
image = pipe(prompt, num_inference_steps=30).images[0]
image.save("out.png")

주의할 점:

  • torch 버전이 낮으면 SDPA가 기대만큼 동작하지 않거나 커널 선택이 제한됩니다.
  • 일부 환경에서는 SDPA가 켜진 상태에서 특정 연산이 느려질 수 있으니, OOM 해결이 우선인지 속도가 우선인지 목표를 정해 선택하세요.

xFormers 활성화 예시

xFormers는 파이프라인에 아래처럼 적용합니다.

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

pipe.enable_xformers_memory_efficient_attention()

image = pipe("a watercolor landscape", num_inference_steps=25).images[0]
image.save("out.png")

실무적으로는 “xFormers를 먼저 켜보고, 설치나 충돌이 있으면 SDPA로 간다” 또는 “torch 최신이면 SDPA부터 간다” 식으로 결정하는 경우가 많습니다.

2) VAE가 터진다: VAE 타일링과 슬라이싱

어텐션을 최적화했는데도 고해상도에서 OOM이 난다면, 범인은 VAE 디코드일 확률이 큽니다. 특히 SDXL에서 최종 디코드 시점에 VRAM이 확 튀는 케이스가 자주 보입니다.

VAE 타일링이 해결하는 문제

VAE 타일링은 이미지를 한 번에 디코드하지 않고, 타일 단위로 쪼개서 디코드합니다. 즉 피크 VRAM을 낮추는 대신 약간의 속도 손해가 생길 수 있습니다.

diffusers에서는 보통 아래 옵션을 사용합니다.

  • enable_vae_tiling()
  • enable_vae_slicing()

둘 다 켜는 조합이 “OOM 방지”에 꽤 강력합니다.

import torch
from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")

# VAE 메모리 피크 낮추기
pipe.enable_vae_tiling()
pipe.enable_vae_slicing()

# 어텐션도 함께 최적화
pipe.unet.set_attn_processor("sdpa")

image = pipe(
    "ultra detailed city skyline at night, rain, reflections",
    width=1024,
    height=1024,
    num_inference_steps=30,
).images[0]
image.save("sdxl_1024.png")

타일링의 부작용과 체크 포인트

  • 타일 경계에서 미세한 이음새가 보이는 경우가 있습니다(모델과 설정에 따라 다름).
  • 속도는 느려질 수 있습니다.
  • “어텐션은 괜찮은데 마지막 저장 직전에 죽는다”면 거의 VAE 쪽이므로, 이 옵션이 1순위입니다.

3) 가장 즉효인 안전장치: 해상도·배치·정밀도

xFormers나 SDPA, VAE 타일링은 “최적화”이고, 아래는 “물리 법칙”에 가까운 안전장치입니다.

FP16 또는 BF16 사용

가능하면 float16 또는 bfloat16로 돌리세요.

  • NVIDIA 소비자 GPU에서는 FP16이 흔한 선택입니다.
  • BF16은 일부 GPU에서 더 안정적일 수 있지만, 지원 여부가 갈립니다.
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

배치 크기와 이미지 크기

  • batch_size는 VRAM에 직격탄입니다. OOM이면 1로 내리고 다시 확인하세요.
  • 해상도는 어텐션 메모리를 기하급수로 키웁니다. 512에서 되던 것이 768에서 터지는 건 정상적인 현상입니다.

Gradient checkpointing은 학습에서, 추론에서는 보통 비해당

추론 OOM을 잡는 데는 위에서 말한 어텐션 최적화와 VAE 타일링이 더 직접적입니다.

4) “분명 VRAM 남는데 OOM”처럼 보일 때: 파편화와 캐시

가끔 nvidia-smi로 보면 여유가 있어 보이는데도 OOM이 납니다. 이때는 메모리 파편화 또는 PyTorch 캐시 동작 때문에 “연속 블록”을 못 잡는 상황일 수 있습니다.

PyTorch allocator 설정

환경 변수로 allocator 동작을 바꿔 파편화를 완화하는 방법이 있습니다.

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128

상황에 따라 max_split_size_mb는 64, 128, 256 등으로 조정해 보세요.

캐시 비우기

프로세스 내에서 여러 번 파이프라인을 교체하거나 큰 텐서를 반복 생성한다면, 중간중간 캐시 정리가 도움이 될 때가 있습니다.

import torch

torch.cuda.empty_cache()

다만 이건 “근본 해결”이라기보다는, 장시간 실행하는 서비스에서 피크를 낮추는 보조 수단에 가깝습니다.

5) 우선순위 가이드: 무엇부터 켜야 하나

OOM 해결은 옵션을 무작정 다 켜기보다, 비용 대비 효과가 큰 순서대로 접근하는 게 좋습니다.

  1. 정밀도 낮추기: FP16 또는 BF16
  2. 배치 1로 고정: 그리고 해상도부터 안정화
  3. 어텐션 최적화: SDPA 또는 xFormers
  4. VAE 타일링·슬라이싱: 특히 SDXL 고해상도에서 필수급
  5. 부가 기능 하나씩 추가: ControlNet, Refiner, Highres fix 등을 단계적으로
  6. allocator 튜닝: 파편화 의심 시

이 접근은 운영 환경 트러블슈팅에서도 동일합니다. 예를 들어 CI에서 “캐시가 안 먹는 것 같다”는 막연한 증상을 원인 후보로 쪼개며 체크하듯이, SD OOM도 “어느 단계에서 피크가 생기는지”를 분리해야 빨리 잡힙니다. 비슷한 방식의 점검 글로 Docker 빌드 cache miss? BuildKit 캐시 진단법도 참고할 만합니다.

6) 자동1111(WebUI) 기준 체크리스트

코드가 아니라 WebUI를 쓰는 경우에도 개념은 같습니다.

  • xFormers 옵션 활성화
  • SDPA 사용 옵션이 있다면 비교(버전별로 표기 다름)
  • VAE 관련 최적화(타일링, 슬라이싱 또는 유사 옵션)
  • 해상도와 Highres fix 업스케일 배율 조정
  • ControlNet 여러 개를 동시에 켜면 VRAM이 급증하므로 1개씩 확인

WebUI는 조합이 다양해 “어떤 옵션이 실제로 적용됐는지”가 헷갈릴 때가 많습니다. 그럴 땐 한 번에 하나씩만 바꿔가며 OOM 재현 여부를 확인하는 게 가장 빠릅니다.

7) 결론: OOM은 어텐션과 VAE를 분리해 잡는다

Stable Diffusion VRAM OOM은 대부분 “어텐션 최적화로 해결되는 OOM”과 “VAE 디코드에서 터지는 OOM”으로 나뉩니다.

  • 어텐션 OOM에는 xFormers 또는 SDPA가 가장 큰 효과를 냅니다.
  • 마지막 단계에서 죽거나 고해상도에서 유독 불안정하면 VAE 타일링·슬라이싱이 결정타가 됩니다.
  • 그래도 불안정하면 해상도, 배치, 정밀도를 먼저 안정화하고 부가 기능을 단계적으로 얹으세요.

이 순서대로만 접근해도 “왜 터지는지 모른 채 옵션을 마구 켜는” 시행착오를 크게 줄일 수 있습니다.