AI Toolkit에서 Gradient Checkpointing을 끄면 OOM이 발생하는 이유

Gradient Checkpointing을 끈 이유가 확실하지 않다면, 가장 안전한 규칙은 간단합니다:

다시 켜세요.

AI Toolkit의 대형 이미지 모델에서 gradient_checkpointing을 끄는 것은 안정적인 학습을 OOM 크래시로 바꾸는 가장 빠른 방법 중 하나입니다.

이 가이드에서 설명하는 내용:

Gradient Checkpointing이 하는 일
끄면 왜 VRAM이 갑자기 부족해지는지
어떤 모델 계열이 가장 민감한지
GC 끄기를 테스트해도 괜찮은 경우
속도가 목표라면 대신 무엇을 변경해야 하는지

빠른 답변

Gradient Checkpointing은 기본적으로 메모리 vs 속도 트레이드오프입니다.

켜짐 = 더 낮은 VRAM 피크, 일반적으로 더 안전
꺼짐 = 더 높은 VRAM 피크, 때때로 더 빠름, OOM 발생 가능성이 훨씬 높음

대부분의 사용자에게, 특히 대형 이미지 모델에서는 품질 조절이 아니라 안정성 조절입니다.

빠른 수정 체크리스트

✅ Show Advanced 클릭
✅ train: 아래에서 gradient_checkpointing: true인지 확인
✅ 동일한 Batch Size와 Resolution으로 먼저 재시도
✅ 여전히 OOM이면 Batch Size를 줄이고 가장 높은 Resolution 버킷 제거
✅ 프리뷰 생성이 문제라면 GC를 끄는 대신 Sample Width / Height / Sample Steps 줄이기

1) Gradient Checkpointing이 실제로 하는 일

올바르게 사용하는 데 PyTorch 이론을 깊이 알 필요는 없습니다.

실용적인 이해 방식:

GC 켜짐에서는 AI Toolkit이 중간 활성화 데이터를 메모리에 적게 유지하고, 필요할 때 일부를 재계산합니다
GC 꺼짐에서는 더 많은 데이터가 VRAM에 상주하여, 메모리가 충분하면 학습이 더 빨라질 수 있습니다

괜찮아 보이지만, 다음이 결합되면 문제가 됩니다:

높은 해상도
다중 해상도 버킷
큰 배치 사이즈
비용이 큰 샘플링
Qwen Edit, Z-Image, FLUX / Flex 계열 같은 무거운 아키텍처

그러면 추가 메모리 여유가 매우 빠르게 사라집니다.

2) 끄면 "갑자기" OOM이 발생하는 이유

사용자들은 종종 이렇게 생각합니다:

"토글 하나만 바꿨을 뿐인데."

하지만 그 토글 하나가 학습 중에 유지되는 활성화 메모리 양을 바꿉니다.

GC 켜짐 상태에서 "무겁지만 괜찮았던" 설정이, GC 꺼짐에서는 "아슬아슬하거나 불가능"해집니다.

그래서 실패가 갑작스럽게 느껴집니다:

같은 데이터셋
같은 모델
같은 프리뷰 프롬프트
같은 배치 사이즈
GC만 변경
이제 스텝 2, 스텝 3, 또는 첫 프리뷰에서 OOM

3) GC 꺼짐 시 고위험 모델 계열

AI Toolkit에서 반복적으로 관찰되는 OOM 패턴을 기반으로, 다음 조합은 기본적으로 위험하게 취급해야 합니다:

모델 계열	GC 꺼짐 시 고위험 조건	안전한 첫 시도
Z-Image	큰 배치로 1024–1536 버킷	GC 켜짐, 보수적으로 시작 후 스케일업
Qwen-Edit	큰 배치 / 여러 무거운 조건의 1024 워크플로	GC 켜짐, Batch Size 1, 프리뷰 부하 줄이기
FLUX-dev / Flex 계열 대형 이미지 모델	큰 배치 또는 고해상도 멀티버킷 학습	GC 켜짐, 여유에 따라 Batch Size 1–4

유용한 사고 모델:

Z-Image는 고해상도 + 큰 배치 + GC 꺼짐 조합에서 빠르게 위험해짐
Qwen Edit는 1024 + 무거운 컨디셔닝 + GC 꺼짐 조합에서 빠르게 위험해짐
FLUX / Flex 계열은 큰 배치 + GC 꺼짐 조합에서 빠르게 위험해짐

4) RunComfy AI Toolkit에서 변경하는 위치

현재 RunComfy UI에서 직접 확인할 수 있는 설정:

Batch Size
Gradient Accumulation
Steps
Unload TE
Cache Text Embeddings

gradient_checkpointing은 고급 설정에서 확인하는 것이 가장 쉽습니다.

단계별 방법

작업 열기
Show Advanced 클릭
train: 블록 찾기
설정:

train:
  gradient_checkpointing: true

작업 저장
다른 것은 변경하지 않고 재시도

같은 설정으로 실행되면 GC가 원인이었음이 확인됩니다.

5) 속도가 진짜 목표라면 대신 해야 할 것

많은 사용자가 더 빠른 학습을 위해 GC를 끕니다.

이해할 수 있지만, 보통은 잘못된 첫 번째 속도 실험입니다.

GC 끄기 전에 이것들을 먼저 시도하세요:

A. 프리뷰 샘플링 비용 줄이기

Sample 패널에서:

Width / Height 줄이기
Sample Steps 줄이기
Sample Every 늘리기
일시적으로 Disable Sampling 켜서 학습 안정성 확인

위험한 GC 꺼짐 설정을 추구하는 것보다 실용적인 속도 향상을 얻을 수 있는 경우가 많습니다.

B. 배치는 작게 유지하고, 필요하면 Accumulation으로 스케일

Training 패널에서:

Batch Size는 보수적으로 유지
VRAM 피크의 큰 점프 없이 유효 배치 사이즈를 약간 늘리고 싶을 때만 Gradient Accumulation 증가

C. 가장 높은 버킷을 먼저 제거

Datasets에서:

512 / 768 유지
안정적인 학습 후에만 1024 / 1536 재도입

D. 모델의 저메모리 경로 사용

일부 아키텍처에서는 올바른 방법이 "GC 꺼짐"이 아니라:

low_vram: true
모델 특화 양자화
모델 특화 텍스트 인코더 최적화

일반적인 추측이 아니라 모델 가이드를 따르세요.

6) GC 꺼짐을 테스트해도 괜찮은 경우

GC 꺼짐은 기본값이 아닌 고급 실험으로 취급하세요.

합리적인 조건:

이미 안정적인 학습이 있음
Batch Size가 여전히 보수적
가장 큰 버킷에서 이미 한계에 있지 않음
프리뷰가 가볍거나 비활성화됨
충분한 VRAM 여유가 있음
하나의 변수만 변경, 여러 개가 아님

좋은 테스트 흐름:

GC 켜짐으로 작업 안정화
나머지 설정은 동일하게 유지
GC 끄기
다음을 관찰:

초기 스텝에서의 OOM
샘플링 OOM
간헐적인 버킷 스파이크

이 중 하나라도 발생하면 GC를 다시 켜고 거기서 멈추세요.

7) GC 꺼짐이 아닌 것

Gradient Checkpointing은 다음이 아닙니다:

마법의 품질 향상 스위치
유사도 개선 토글
나쁜 샘플에 대한 올바른 첫 대응
이미 OOM 상태에서의 좋은 첫 실험

샘플이 약하다면 다음을 확인하세요:

데이터셋 품질
캡션
Rank
Steps
프리뷰 일치성

GC 꺼짐이 답이라고 가정하지 마세요.

8) FAQ

GC를 켜면 출력 품질이 떨어지나요?

실제로 대부분의 사용자가 느끼는 차이가 아닙니다.

진짜 트레이드오프는 주로 메모리 vs 속도입니다.

GC를 켰는데 여전히 OOM이 납니다. 이제 어떻게 하죠?

다음 조절 포인트:

Batch Size
가장 큰 Resolution 버킷
프리뷰 샘플링 비용
비디오의 경우 Num Frames

첫 학습을 GC 꺼짐으로 시작해도 되나요?

대부분의 사용자에게: 아니요.

먼저 안정성을 확인한 다음 실험하세요.

한 줄 요약

AI Toolkit에서 OOM이 발생하고 gradient_checkpointing이 꺼져 있다면, 그것부터 고치세요.

GC 켜짐이 안전한 기본값입니다. GC 꺼짐은 고급 속도 실험입니다.