Stable Diffusion VRAM OOM 해결 - xFormers·Triton 최적화

Stable Diffusion을 돌리다 보면 가장 흔하게 마주치는 에러가 CUDA out of memory, 즉 VRAM OOM입니다. 해상도를 조금만 올리거나 배치 크기를 늘리면 바로 터지고, 같은 설정인데도 어떤 날은 되고 어떤 날은 안 되는 것처럼 보이기도 합니다.

이 글에서는 OOM을 단순히 해상도 탓으로 돌리지 않고, 어떤 텐서가 VRAM을 잡아먹는지를 기준으로 원인을 쪼갠 뒤, 실전에서 효과가 큰 xFormers 메모리 효율 어텐션과 Triton 기반 커널 최적화를 중심으로 해결 전략을 정리합니다.

참고로 LoRA를 섞어 쓰는 환경이라면, 메모리뿐 아니라 결과 품질 이슈도 같이 터질 수 있습니다. LoRA 병합 후 색감이 깨지는 문제는 별도 글에서 다뤘습니다: Stable Diffusion LoRA 병합 후 색감 깨짐 해결

VRAM OOM이 나는 진짜 지점: U-Net의 어텐션과 활성화 메모리

Stable Diffusion 파이프라인에서 VRAM을 크게 쓰는 축은 대략 다음입니다.

• 모델 파라미터(가중치): U-Net, VAE, 텍스트 인코더 가중치
• 활성화(activations): 특히 U-Net의 중간 피처맵, 어텐션의 Q/K/V 및 attention matrix
• 옵티마이저 상태(학습 시): Adam 계열의 모멘트 등(추론이면 해당 없음)
• 캐시/워크스페이스: cuDNN, 커널 오토튜닝, allocator fragmentation

추론(inference)에서 OOM의 주범은 보통 U-Net의 활성화 메모리입니다. 이 중에서도 어텐션은 해상도에 따라 메모리 스케일이 급격히 커집니다.

• 일반적인 “naive attention”은 attention matrix 크기가 토큰 수의 제곱에 비례합니다.
• 이미지 latent의 토큰 수는 대략 공간 크기에 비례하므로, 해상도를 올리면 메모리 폭증이 발생합니다.

여기서 xFormers의 memory efficient attention이 강력합니다. attention matrix를 통째로 만들지 않고, 타일링/스트리밍 방식으로 계산해 피크 VRAM을 크게 줄입니다.

먼저 확인할 것: OOM이 진짜 VRAM 부족인지, 단편화인지

같은 설정인데도 간헐적으로 OOM이 난다면, 절대량 부족보다 메모리 단편화(fragmentation) 가능성이 큽니다. PyTorch CUDA allocator가 큰 연속 블록을 못 잡아 실패하는 케이스입니다.

빠른 진단 코드

아래는 현재 프로세스에서 VRAM 사용량을 대략 확인하는 예시입니다.

import torch

def vram_report(tag=""):
    torch.cuda.synchronize()
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    max_alloc = torch.cuda.max_memory_allocated() / 1024**2
    print(f"[{tag}] allocated={allocated:.1f}MB reserved={reserved:.1f}MB max_alloc={max_alloc:.1f}MB")

vram_report("start")
# inference here
vram_report("after")

• allocated는 실제 텐서가 쓰는 메모리
• reserved는 allocator가 잡아둔 풀(pool)

reserved가 과도하게 커지고 allocated는 줄었는데도 OOM이 난다면 단편화 가능성이 있습니다.

단편화 완화 팁

• 프로세스를 오래 살려두는 서버형 워크로드라면, 주기적으로 워커를 재시작하거나 요청 단위로 프로세스를 분리하는 것도 실전에서 효과가 큽니다.
• PyTorch 환경변수로 allocator 동작을 조정할 수 있습니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,garbage_collection_threshold:0.9"

max_split_size_mb는 큰 블록을 잘게 쪼개는 전략을 완화해 단편화를 줄이는 데 도움이 되는 경우가 있습니다.

xFormers: 메모리 효율 어텐션으로 피크 VRAM 줄이기

xFormers는 Stable Diffusion 계열에서 사실상 “OOM 방지의 1순위 옵션”입니다. 핵심은 다음입니다.

• 어텐션 연산을 메모리 효율 구현으로 바꿔 attention matrix의 물리적 생성 부담을 줄임
• 결과적으로 피크 VRAM이 낮아져 더 높은 해상도, 더 큰 배치, 더 긴 프롬프트가 가능해짐

설치(일반적인 경우)

환경에 따라 다르지만, 보통은 PyTorch CUDA 버전에 맞는 xFormers 휠을 설치합니다.

pip install -U xformers

설치 후에는 실제로 xFormers가 활성화되었는지 로그나 설정으로 확인해야 합니다. UI나 런처마다 플래그가 다르지만, diffusers를 직접 쓴다면 아래처럼 확인/설정하는 흐름이 일반적입니다.

diffusers에서 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(
    prompt="a photo of a cat",
    num_inference_steps=30,
    guidance_scale=7.5,
).images[0]

주의할 점

• 일부 조합에서는 xFormers가 속도를 올리기도 하지만, 환경에 따라 속도가 비슷하거나 약간 느려질 수도 있습니다. 목적이 OOM 회피라면 충분히 가치가 큽니다.
• 드라이버, CUDA, PyTorch, xFormers 버전 매칭이 어긋나면 빌드나 런타임에서 문제가 날 수 있습니다. 특히 소스 빌드가 필요한 상황에서는 컴파일 시간이 길고 실패 확률도 올라갑니다.

Triton: 커널 최적화로 속도와 메모리를 동시에 잡는 지점

Triton은 GPU 커널을 Python 레벨에서 작성/컴파일할 수 있게 해주는 프레임워크로, 최근 PyTorch 생태계에서 다양한 최적화의 기반으로 쓰입니다.

Stable Diffusion 관점에서 Triton이 체감되는 지점은 크게 두 가지입니다.

• 어텐션/레이어 정규화/MLP 같은 핵심 연산에 대한 더 나은 커널
• 일부 경로에서 불필요한 메모리 이동을 줄여 피크 VRAM과 시간을 함께 낮춤

다만 “Triton을 설치하면 자동으로 OOM이 해결된다” 같은 단순한 그림은 아닙니다. 보통은 다음 중 하나로 들어옵니다.

• PyTorch가 내부적으로 Triton 커널을 쓰는 경로를 타게 됨
• xFormers나 FlashAttention 류가 내부에서 Triton을 사용
• torch.compile 또는 특정 최적화 라이브러리가 Triton을 활용

Triton 설치

pip install -U triton

환경에 따라 PyTorch가 요구하는 Triton 버전이 정해져 있을 수 있으니, 충돌이 나면 PyTorch 릴리스 노트를 우선 확인하는 편이 안전합니다.

xFormers와 Triton을 써도 OOM이면: 우선순위 체크리스트

여기부터는 “기술적으로 가장 효과가 큰 순서”에 가깝게 정리합니다.

1) FP16 또는 BF16으로 추론

가중치와 활성화를 절반 정밀도로 내리면 VRAM이 즉시 줄어듭니다.

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

• NVIDIA Ampere 이후는 BF16도 선택지입니다.
• 다만 모든 연산이 완벽히 FP16으로 내려가는 것은 아니고, 일부는 FP32로 남을 수 있습니다.

2) VAE를 타일링 또는 슬라이싱

VAE 디코딩은 고해상도에서 VRAM을 크게 먹습니다. diffusers 기준으로는 아래 옵션이 흔합니다.

pipe.enable_vae_slicing()
pipe.enable_vae_tiling()

• 타일링은 속도 손해가 있을 수 있지만, 고해상도에서 OOM을 막는 데 매우 유효합니다.

3) Attention slicing

xFormers가 불가능하거나 여전히 빡빡하면 attention을 슬라이스해서 피크를 줄일 수 있습니다.

pipe.enable_attention_slicing("max")

• 속도는 느려질 수 있습니다.

4) 해상도만이 아니라 latent 크기를 의식하기

Stable Diffusion은 이미지 공간이 아니라 latent 공간에서 U-Net이 돌아갑니다. 그러나 해상도가 커지면 latent의 공간 크기도 커지고, 결국 U-Net 활성화가 폭증합니다.

• 512에서 768로 올리는 것은 단순 1.5배가 아니라, 면적 기준으로 더 큰 증가입니다.
• OOM이 나면 해상도를 “조금” 내리는 것만으로도 피크 VRAM이 크게 줄 수 있습니다.

5) 배치 크기, 동시성부터 줄이기

서버에서 여러 요청을 동시에 처리하면, 배치가 1이어도 동시 실행으로 인해 VRAM이 겹칩니다.

• 워커 수를 줄이거나
• 큐잉을 도입하거나
• GPU당 동시 실행을 1로 제한

이런 운영적 조치가 가장 확실한 해결책인 경우가 많습니다.

실전 구성 예시: “8GB VRAM에서 최대한 버티기” 세팅

아래는 diffusers 기준으로, 비교적 보수적으로 VRAM을 아끼는 조합입니다.

import torch
from diffusers import StableDiffusionPipeline

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

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

# 1) xFormers
pipe.enable_xformers_memory_efficient_attention()

# 2) VAE 메모리 절약
pipe.enable_vae_slicing()
pipe.enable_vae_tiling()

# 3) 필요 시 attention slicing까지
# pipe.enable_attention_slicing("max")

prompt = "highly detailed portrait photo"
image = pipe(prompt, num_inference_steps=28, guidance_scale=7).images[0]

여기서도 OOM이 나면, 다음 순서로 조정하는 것이 보통 가장 효율적입니다.

• 해상도 낮추기
• 스텝 줄이기(스텝은 주로 시간에 영향, 일부 메모리에도 간접 영향)
• attention slicing 켜기
• VAE 타일링 유지

자주 하는 오해: “Triton을 켜면 VRAM이 무조건 준다”

Triton은 커널 레벨 최적화 도구라서, 어떤 연산은 메모리를 덜 쓰도록 바뀔 수 있지만, 모든 경로가 그런 것은 아닙니다. 다음을 기억하면 디버깅이 쉬워집니다.

• OOM은 피크 메모리 문제다
• 커널 최적화는 피크를 낮추는 경우도 있지만, 주로 속도 개선에 더 직접적일 수 있다
• 피크 메모리는 대체로 어텐션, 고해상도 U-Net 활성화, VAE 디코딩에서 터진다

즉, Triton은 “xFormers로도 부족할 때의 추가 카드”이거나 “속도까지 챙기고 싶을 때의 카드”로 보는 편이 현실적입니다.

운영 관점 팁: OOM을 장애로 만들지 않기

서비스 형태로 Stable Diffusion을 운영하면, OOM은 단순 에러가 아니라 프로세스 크래시나 워커 재시작으로 이어져 장애가 됩니다. 이때는 모델 최적화뿐 아니라 운영 전략이 중요합니다.

• 요청당 최대 해상도, 최대 스텝, 최대 배치 제한을 명시
• 동시성 제한과 큐잉
• 워커 재시작 정책과 헬스체크

Triton Inference Server를 쓰는 경우라면, 헬스체크와 워커 튜닝이 503과 맞물리는 케이스가 많습니다. 관련 운영 튜닝은 아래 글이 맥락상 도움이 됩니다: NVIDIA Triton 배포서 503? 헬스체크·워커 튜닝

마무리: OOM 해결의 정석은 “피크 VRAM을 낮추는 것”

Stable Diffusion의 VRAM OOM은 대부분 “총량 부족”이라기보다, 특정 구간에서 순간적으로 치솟는 피크 VRAM 때문에 발생합니다. 그래서 해결도 다음 원칙을 따르면 빠릅니다.

• 1순위: xFormers 메모리 효율 어텐션으로 어텐션 피크를 내린다
• 2순위: FP16 또는 BF16, VAE slicing 또는 tiling으로 큰 덩어리를 쪼갠다
• 3순위: attention slicing, 해상도 및 동시성 제한으로 피크를 제어한다
• 보너스: Triton 기반 최적화로 속도와 일부 메모리 이동을 개선한다

이 조합으로도 부족하다면, 결국은 더 큰 VRAM의 GPU로 가거나, 모델을 더 작은 것으로 바꾸거나, 해상도 정책을 조정하는 운영적 합의가 필요합니다. 그래도 대부분의 환경에서는 위 설정만으로 “자주 터지던 OOM”을 “관리 가능한 수준”으로 내릴 수 있습니다.