[TIL] PyTorch CUDA Out of Memory(OOM) 오류 해결기 (feat. VLM 모델 학습)

개발 환경:

- AWS g5.12xlarge (NVIDIA A10G 24GB x4)

- CUDA 12.8

- PyTorch 2.7.1

1. 문제 사항

Vision-Language Model (VLM)을 학습하는 과정에서 다음과 같은 CUDA OOM 오류가 지속 발생

[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.70 GiB. GPU 1 has a total capacity of 22.07 GiB of which 2.06 GiB is free. Including non-PyTorch memory, this process has 20.00 GiB memory in use. Of the allocated memory 19.39 GiB is allocated by PyTorch, and 149.28 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

2. 원인 분석

• GPU 메모리 부족: 2.70 GiB 할당 요청 중, 사용 가능한 메모리는 2.06 GiB
• PyTorch 메모리 할당: 이미 19.39 GiB가 PyTorch에 의해 사용 중
• 단편화(Fragmentation) 문제: 일부 메모리는 reserved 상태지만 실제 할당되지 않아 할당 실패

3. 해결 과정

1. expandable_segments:True 설정

단편화 완화를 위해 PyTorch 환경변수 추가:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

 

PyTorch는 GPU 메모리를 다음 두 단계로 관리함:

1. 예약(reserved): GPU에서 큰 덩어리의 메모리를 미리 확보
2. 할당(allocated): 이 예약된 영역 중 일부를 실제 연산에 사용하도록 나눔

학습 중 layer별로 GPU 메모리 할당/해제가 반복되면서, 남은 메모리가 작고 쪼개진 조각으로 남게 되는데, 이때 PyTorch가 연속된 일정 크기의 메모리를 요청할 경우, 실제 총 남은 용량은 충분해도 할당이 실패할 수 있음.

PyTorch 2.0 이후에는 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 옵션을 통해 단편화 문제를 줄일 수 있는 개선된 allocator를 사용할 수 있음. 이 옵션은 기존에 고정된 크기의 allocator segment를 사용하는 대신, 동적으로 확장 가능한 segment를 도입함

그 결과 더 유연한 메모리 할당이 가능하고, 단편화에 강함. 또한, 연속 메모리 블록을 찾기 쉬워짐. 결과적으로 메모리 할당 실패(=allocation failure)가 줄어들고, 모델 학습 중 OOM 발생 빈도도 감소하게 됨

2.  gc.collect() + torch.cuda.empty_cache() 사용

CUDA 메모리를 자주 할당(allocate)/해제(free)를 하면 GPU 드라이버 레벨에서 오버헤드가 크기 때문에, PyTorch는 자체 캐싱 메커니즘을 통해 관리함. 

gc.collect() : Python의 Carbage Collector(GC)를 수동으로 실행시켜, 더 이상 참조되지 않는 파이썬 객체를 메모리 해제하는 역할을 함. PyTorch에서 Tensor는 파이썬의 객체이기도 하고, GPU 메모리를 소유한 C++의 객체이기도 함. del tensor를 하더라도 참조 카운트가 0이 아니면 즉시 해제되지 않기 때문에, gc.collect()를 호추랗여 파이썬 인터프리터가 순환 참조된 객체까지도 강제로 해제하게 만드는 역항르 함

torch.cuda.empty_cache() : GPU 메모리 중 PyTorch가 캐시용으로 보유하고 있는 메모리 블록을 해제하여, CUDA 드라이버에 반환하는 역할. 실제 GPU 상의 메뮤리를 줄일 수 있음. 다만, allocated된 메모리는 해제 하지 않고, 캐시용으로 보유하던 메모리만 반환하는 역하을 진행

단독으로 torch.cuda.emtpy_cache()만 호출하면, 아직 참조되고 있는 Tensor 때문에 PyTorch가 여전히 allocated 메모리를 유지함. 따라서 gc.collect()로 객체를 해제하여 참조를 끊고, 메모리를 반환 조건 충적 후 empty_cache()로 GPU 메모리를 실 반환하므로써 OOM 문제를 해결

gc.collect() + torch.cuda.empty_cache() 조합의 경우, 적절한 타이밍에 위치를시켜야 효과적임

• 큰 텐서들이 소멸한 후
• 다음 GPU 연산이 시작되기 전
• 반복적 학습 스텝 간에 일시적 피크 메모리가 문제가 될 때 

나의 경우에는 HuggingFace Traniner.train() 내부에서, 학습이 진행중인 step에서, forward -> bachward로 수행하는 loss.backward() 중 OOM이 발생하였다. 한 step 처리 중 backward pass 계산 그래프가 너무 커져서 GPU VRAM을 초과한 상태였다. 근본적인 해결책은 max_token_length를 줄이거나, gradient checkpointing를  적용해서 해결을 해야 했다. 

먼저, 적용하기 전에 초기 캐시 메모리 및 expandable_segments:True 적용하고, Trainer 인스턴스를 완전히 구성하고 나서, 불필요하게 남아 있을 수 있는 메모리를 수동 정리하여 일단 돌아가게는 했다.

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=None,
        tokenizer=tokenizer,
        data_collator=collator,
    )

    import gc
    gc.collect()
    torch.cuda.empty_cache()

이렇게 정리하여, forward/backward 전에 더 많은 여유 메모리를 확보하여 전체적인 VRAM 사용은 줄이도록 하였다.