Stable Diffusion VRAM OOM - xFormers·Tiling·FP8 최적화

Stable Diffusion을 돌리다 보면 가장 흔히 마주치는 에러가 CUDA OOM, 즉 VRAM 부족입니다. 특히 고해상도(1024x1024 이상), 배치 증가, ControlNet 다중 사용, 업스케일러 추가 같은 옵션을 켜는 순간 메모리 사용량이 기하급수적으로 튀면서 터집니다.

이 글은 단순히 “해상도 낮추세요”가 아니라, **왜 OOM이 발생하는지(어떤 텐서가 VRAM을 잡아먹는지)**를 기준으로, xFormers(메모리 효율 어텐션), Tiling(타일 기반 VAE/UNet 처리), **FP8(저정밀 연산/저장)**까지 실제로 적용 가능한 최적화 루트를 정리합니다.

또한 OOM은 GPU만의 문제가 아니라 시스템 전체 메모리 정책과도 연결됩니다. 리눅스 환경에서 프로세스가 갑자기 죽는다면 커널 레벨에서 무엇이 일어났는지 함께 확인하는 것이 좋습니다: 리눅스 OOM Killer로 프로세스 죽음 진단·방지

VRAM OOM이 나는 구조: 어디서 메모리가 터지나

Stable Diffusion의 핵심은 대략 다음 블록으로 구성됩니다.

• Text Encoder(CLIP): 프롬프트를 임베딩으로 변환
• UNet(확산 모델 본체): 단계(steps)마다 latent를 업데이트
• VAE: latent를 이미지로 디코딩(또는 이미지에서 latent 인코딩)
• 옵션 모듈: ControlNet, LoRA, IP-Adapter, 업스케일러 등

VRAM 사용량을 크게 만드는 주범은 보통 아래입니다.

1. UNet의 어텐션(Attention) 중간 텐서
2. 고해상도 latent: 해상도가 커질수록 latent 텐서도 커짐
3. 배치(batch size)와 CFG: 동시에 들고 있는 텐서 수 증가
4. VAE 디코딩 시 피크 메모리: 최종 이미지로 풀 때 순간 피크가 큼

대략적인 감각을 잡기 위한 규칙은 이렇습니다.

• 해상도를 2x로 키우면 픽셀 수는 4x가 되고, 중간 텐서도 유사하게 커집니다.
• batch size를 2로 올리면 대부분 텐서가 2배로 증가합니다.
• ControlNet을 여러 개 붙이면 UNet에 들어가는 조건 텐서가 늘어 VRAM이 더 듭니다.

즉, OOM은 특정 단계에서 순간적으로 피크가 치솟는 패턴이 많고, 그 피크를 줄이는 게 핵심입니다.

0단계: OOM을 재현 가능하게 만들기(로그·모니터링)

최적화는 “바꿨더니 되는 것 같음”이 아니라, 어떤 옵션이 VRAM 피크를 줄였는지 확인하면서 진행해야 합니다.

nvidia-smi로 피크 관찰

아래처럼 1초 단위로 VRAM을 찍어두면, 어떤 구간에서 피크가 생기는지 감이 옵니다.

watch -n 1 nvidia-smi

PyTorch 메모리 스냅샷(선택)

파이프라인 코드 레벨에서 확인할 수 있다면 다음이 유용합니다.

import torch

print("allocated", torch.cuda.memory_allocated() / 1024**2, "MB")
print("reserved ", torch.cuda.memory_reserved() / 1024**2, "MB")
print(torch.cuda.memory_summary())

allocated는 실제 텐서가 사용 중인 메모리, reserved는 캐싱까지 포함한 예약량입니다. OOM은 보통 reserved가 커서 나는 것처럼 보이기도 하지만, 실제로는 큰 텐서를 새로 할당하려다 실패하는 경우가 많습니다.

1단계: 가장 먼저 먹히는 처방 5가지(부작용 적음)

아래는 품질 영향이 거의 없거나 체감이 적으면서도 효과가 큰 순서입니다.

1) --xformers 또는 SDPA 사용

xFormers는 메모리 효율 어텐션 구현으로, 어텐션 중간 텐서를 덜 들고 가도록 최적화합니다. Stable Diffusion에서 가장 흔한 “한 방에 해결” 옵션입니다.

• Automatic1111: 실행 옵션에 --xformers
• ComfyUI: 노드/설정에서 xFormers 또는 PyTorch SDPA 사용

주의할 점은 환경에 따라 xFormers 설치가 까다롭고, 특정 GPU/드라이버 조합에서 성능이나 안정성이 달라질 수 있다는 것입니다.

2) FP16 또는 BF16 강제

대부분 환경에서 FP16은 기본이지만, 간혹 FP32로 돌아가거나 일부 모듈이 FP32로 승격되면 VRAM이 급증합니다.

Diffusers 예시:

import torch
from diffusers import StableDiffusionPipeline

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

Ampere 이상이면 BF16도 고려할 수 있습니다.

3) Attention slicing

어텐션을 슬라이스로 쪼개 처리해 피크 메모리를 줄입니다. 속도는 조금 느려질 수 있습니다.

pipe.enable_attention_slicing("auto")

4) VAE slicing

VAE 디코딩 시 피크를 줄이는 옵션입니다.

pipe.enable_vae_slicing()

5) CPU offload(최후의 보루에 가깝지만 강력)

GPU VRAM이 작다면 모델 일부를 CPU로 내립니다. 속도는 크게 손해 볼 수 있습니다.

pipe.enable_model_cpu_offload()

2단계: xFormers 제대로 이해하고 적용하기

xFormers가 효과적인 이유는 어텐션 계산에서 발생하는 거대한 행렬(QK^T)과 softmax 중간 결과를 전부 저장하지 않고, 더 효율적인 커널로 처리하기 때문입니다.

적용 체크리스트

• PyTorch 버전과 CUDA 버전이 맞는지
• GPU 아키텍처가 지원되는지
• 실행 시 실제로 xFormers 경로를 타는지(로그 확인)

만약 xFormers가 불안정하거나 설치가 어렵다면, PyTorch의 SDPA(Scaled Dot Product Attention) 경로가 대안입니다. 최신 PyTorch에서는 SDPA만으로도 상당한 메모리 절감이 됩니다.

3단계: Tiling으로 피크 메모리 깨기(VAE tiling 중심)

고해상도에서 OOM이 터지는 대표 구간이 VAE 디코딩입니다. latent를 최종 RGB로 풀어내는 과정에서 순간적으로 큰 텐서가 생깁니다.

Tiling은 이미지를 타일로 쪼개서 처리하므로, **한 번에 필요한 활성화 메모리(activation)**를 줄입니다.

Diffusers에서 VAE tiling

Diffusers는 VAE tiling을 지원합니다.

pipe.enable_vae_tiling()

• 장점: VRAM 피크가 크게 감소
• 단점: 타일 경계 아티팩트 가능(대개 오버랩/설정으로 완화)

UNet tiling은 어디에 쓰나

UNet 자체를 타일링하는 방식도 존재하지만, 구현/워크플로우에 따라 난이도가 높습니다. 실무적으로는 다음 순서가 효율적입니다.

1. VAE tiling으로 “디코딩 피크” 제거
2. 그래도 부족하면 해상도/배치/ControlNet 개수 조정
3. 그 다음이 UNet 쪽 더 강한 최적화

4단계: FP8 최적화의 현실적인 포인트

FP8은 이름만 들으면 “VRAM을 반으로 줄인다” 같은 기대를 하게 되지만, 실제 적용은 다음 조건에 크게 좌우됩니다.

• 하드웨어 지원: Hopper(H100) 계열은 FP8에 최적화
• 프레임워크 지원: Transformer Engine, 최신 CUDA/드라이버, PyTorch 설정
• 어디에 FP8을 쓰는가: 가중치 저장만 FP8인지, 연산까지 FP8인지

Stable Diffusion에서 FP8을 실전적으로 쓰는 케이스는 보통 다음으로 나뉩니다.

1. 가중치(weight) FP8 저장/로드로 VRAM 절감
2. 일부 연산을 FP8로 수행해 속도/메모리 최적화

다만 FP8은 수치 안정성 이슈가 생길 수 있어, 모델/스케줄러/옵션 조합에 따라 밴딩이나 디테일 손실이 보일 수 있습니다. 그래서 운영 관점에서는 아래처럼 접근하는 편이 안전합니다.

• 1차 목표: xFormers + FP16(BF16) + VAE tiling으로 해결
• 2차 목표: 그래도 안 되면 FP8 가중치 또는 더 강한 오프로딩

FP8 가중치 로딩 예시(개념 코드)

환경마다 구현이 달라 “그대로 복붙”은 어렵지만, 대략 이런 형태로 구성됩니다.

# 개념 예시: 실제로는 사용 라이브러리(TE, optimum, custom loader)에 따라 다릅니다.
# 핵심은 "weights를 더 낮은 정밀도로" 들고 가서 VRAM 상주량을 줄이는 것.

import torch

def load_weights_low_precision(model, ckpt_path):
    state = torch.load(ckpt_path, map_location="cpu")
    model.load_state_dict(state)
    model.to(dtype=torch.float16)  # 또는 fp8 지원 경로
    return model

FP8은 “한 줄 옵션”으로 끝나는 경우가 드물고, 커널/드라이버/모델 호환성까지 같이 봐야 합니다.

5단계: OOM을 줄이는 운영용 체크리스트(우선순위)

아래는 실제로 문제를 빠르게 줄이는 우선순위입니다.

1. xFormers 또는 SDPA 활성화
2. FP16 또는 BF16 고정
3. VAE slicing + VAE tiling 적용
4. Attention slicing 적용
5. 해상도 조정: 1024에서 896 또는 768로 내려도 체감 품질이 크게 안 깨지는 경우가 많음
6. batch size를 1로 고정
7. ControlNet 개수 줄이기 또는 해상도 낮은 프리프로세서 사용
8. CPU offload / sequential offload

여기서 중요한 운영 팁은, “옵션을 한 번에 여러 개 바꾸지 말고” 하나씩 적용해 VRAM 피크 변화를 관찰하는 것입니다.

6단계: 그래도 터지면? GPU OOM과 시스템 OOM을 분리

간혹 사용자는 CUDA OOM만 보고 있지만, 실제로는 CPU RAM이 부족하거나 스왑 정책 때문에 프로세스가 죽는 경우도 있습니다. 특히 CPU offload를 켰을 때 이런 일이 늘어납니다.

• CUDA OOM: 보통 Python 예외로 명확히 보임
• 시스템 OOM: 프로세스가 갑자기 종료되거나 killed 로그가 남음

리눅스에서 시스템 OOM을 의심한다면 커널 로그를 확인해 원인을 분리하세요. 자세한 진단 흐름은 아래 글이 도움이 됩니다.

• 리눅스 OOM Killer로 프로세스 죽음 진단·방지

또한 컨테이너 환경에서 재시작 루프가 난다면 OOM으로 프로세스가 죽고 재기동되는 패턴일 수 있습니다. 쿠버네티스라면 이 글의 점검 항목도 같이 보면 좋습니다.

• Kubernetes CrashLoopBackOff 원인 7가지·즉시복구

실전 예시: Diffusers 파이프라인에 “OOM 방지 세트” 적용

아래는 Diffusers에서 흔히 쓰는 조합을 한 번에 묶은 예시입니다.

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")

# 1) 메모리 효율 옵션들
pipe.enable_attention_slicing("auto")
pipe.enable_vae_slicing()
pipe.enable_vae_tiling()

# 2) 정말 부족할 때만(속도 손해 큼)
# pipe.enable_model_cpu_offload()

prompt = "a photo of a futuristic city at night, ultra detailed"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("out.png")

이 세트(attention slicing, vae slicing, vae tiling)는 품질 영향이 비교적 적으면서 피크 VRAM을 잘 낮춰줍니다. 여기에 xFormers 또는 SDPA가 더해지면, 같은 GPU에서도 한 단계 높은 해상도를 시도할 여지가 생깁니다.

결론: 정답은 한 가지가 아니라 “피크 제거” 전략

Stable Diffusion VRAM OOM은 대부분 “총량 부족”이라기보다 특정 구간 피크가 원인입니다. 그래서 해결도 순서가 있습니다.

• 1순위: xFormers(또는 SDPA)로 어텐션 메모리부터 줄이기
• 2순위: VAE tiling으로 고해상도 디코딩 피크 제거
• 3순위: FP16/BF16 고정 + slicing 조합으로 안정화
• 4순위: FP8은 환경이 받쳐줄 때, 마지막 최적화 카드로 신중히 적용

이 흐름대로 적용하면 “VRAM이 작은 GPU에서도 돌아가게 만들기”와 “같은 GPU에서 더 높은 해상도/옵션을 열기” 둘 다 현실적으로 달성할 수 있습니다.