PyTorch CUDA 메모리 관리

PyTorch CUDA 메모리 관리 이해하기

딥러닝 모델을 GPU로 학습하다 보면 “CUDA out of memory” 오류를 접하는 일이 많습니다. 한정된 GPU 메모리를 효율적으로 사용하지 못하면 학습 중간에 메모리가 부족해지기 때문입니다. 이번 글에서는 PyTorch의 CUDA 메모리 관리 개념을 살펴보겠습니다. CUDA 메모리의 기본 개념부터 PyTorch의 메모리 관리 전략, 멀티 GPU 환경에서의 메모리 관리, 그리고 메모리 최적화 기법과 디버깅 방법 등을 알아봅니다. 

1. CUDA 메모리 개요

CUDA (에서 사용하는) 메모리란 NVIDIA GPU에서 사용하는 전용 메모리(VRAM)를 말합니다. CPU 메모리와 분리되어 있으므로, GPU에서 연산하려면 데이터를 명시적으로 GPU 메모리로 복사해야 합니다. C/C++ CUDA API에서는 cudaMalloc 함수를 통해 GPU 메모리를 할당하고 cudaFree로 해제합니다. 이는 CPU에서의 malloc/free와 유사하게 동작하지만, 중요한 차이점이 있습니다. GPU 메모리 할당/해제는 CPU 메모리보다 훨씬 느리고 비용이 높습니다
 
예를 들어 cudaMalloc이나 cudaFree 호출은 내부적으로 GPU와 동기화를 발생시켜 지연을 유발하는데, 이 비용은 일반적인 malloc/free에 비해 몇 배에서 수십 배 이상 크다고 알려져 있습니다. 따라서 GPU 메모리를 빈번하게 할당하거나 해제하면 연산 속도가 크게 떨어질 수 있습니다. GPU 연산은 비동기적으로 이루어지기 때문에, CPU가 GPU 작업을 지시하고 곧바로 다음 작업을 수행할 수 있습니다. 하지만 cudaFree와 같은 메모리 해제 호출은 GPU가 모든 연산을 완료할 때까지 동기화(synchronization)를 일으킵니다.

이로 인해 CPU와 GPU 간 병렬성이 깨지고, 전체 학습 속도가 저하될 수 있습니다. 이러한 이유로 딥러닝 프레임워크들은 GPU 메모리 관리를 효율화하기 위한 자체 전략을 사용합니다. PyTorch에서는 이런 성능 문제를 완화하기 위해 CUDA 메모리 할당/해제 호출을 최소화하고 이미 확보한 메모리를 재사용하는 방향으로 설계되어 있습니다.

2. PyTorch의 메모리 관리 전략

PyTorch CUDA Caching Allocator: PyTorch는 CUDA 캐싱 할당자(Caching Allocator)를 사용하여 GPU 메모리를 관리합니다. 간단히 말해, 한 번 할당한 GPU 메모리 블록을 바로바로 cudaFree로 반환하지 않고 캐시(pool)에 보관해 둔 뒤 재활용하는 기법입니다. 이렇게 하면 매번 새로 GPU 메모리를 할당하는 것을 피하여 할당 속도를 높이고, cudaFree 호출로 인한 동기화 지연을 줄일 수 있습니다.

PyTorch 프로그램이 실행되는 동안 한때 N 바이트의 GPU 메모리를 썼다면, 내부적으로 그 N 바이트를 캐시에 유지해 두었다가 다음번에 재사용합니다. 이 때문에 del로 텐서를 지워도 즉시 OS에 메모리를 반환하지 않아 nvidia-smi에는 사용 중인 메모리로 나타나는 현상이 발생합니다. PyTorch는 캐시 된 메모리를 새로운 텐서 할당에 바로 활용하므로, 실제로는 메모리가 효율적으로 재활용되고 있습니다.
 
메모리 풀과 단편화 방지: PyTorch 캐싱 할당자는 내부적으로 서로 다른 크기의 메모리 블록 풀(pool)을 관리하여 메모리 단편화(fragmentation)를 줄입니다. 예를 들어 1MB 미만 크기의 작은 할당들은 별도의 풀에서 관리하여, 큰 덩어리 메모리 풀이 잘게 조각나는 것을 피합니다. 새로운 텐서를 할당할 때는 캐시 된 블록 중에서 요청 크기에 가장 근접한 크기의 블록을 찾아 할당하고, 없다면 cudaMalloc로 새로운 블록을 확보한 뒤 풀에 등록합니다. 또한 필요한 경우 큰 블록을 둘로 쪼개거나(splitting), 인접한 빈 블록들을 합치는(merging) 방식으로 내부 단편화를 최소화합니다. 이런 메모리 풀링 기법 덕분에 PyTorch는 GPU 메모리를 오랫동안 할당/해제해도 비교적 안정적으로 사용할 수 있고, 메모리 부족(OOM) 위험을 낮춥니다
 
단, 캐싱 할당자로 인해 사용하지 않는 메모리도 PyTorch 프로세스가 계속 점유하기 때문에, 동일 GPU를 다른 프로그램과 동시에 사용할 때는 주의가 필요합니다.
 
Pinned Memory(고정 메모리) 활용: PyTorch는 핀 메모리를 활용하여 CPU 메모리에서 GPU로의 데이터 전송 속도를 높입니다. 핀 메모리란 페이지아웃되지 않도록 고정된(page-locked) CPU 메모리로서, GPU DMA(Direct Memory Access)가 직접 접근할 수 있는 영역입니다. 일반적인 pageable 메모리는 GPU로 전송하기 전에 운영체제가 임시로 페이지를 고정하거나 복사본을 만드는 추가 작업이 필요하지만, 핀 메모리는 이러한 과정을 생략하고 바로 DMA 전송이 가능하기 때문에 전송이 더 빠르고 효율적입니다.

PyTorch의 Tensor.pin_memory() 메서드를 사용하면 해당 텐서의 데이터를 핀 메모리에 복사해 둘 수 있고, 이렇게 하면 이후. cuda() 또는. to()로 GPU에 복사할 때 non_blocking=True 옵션을 통해 비동기 전송을 수행할 수 있습니다. 이는 데이터 전송과 GPU 연산을 겹쳐 수행할 수 있게 해 주어 전체 학습 속도를 향상합니다. 특히 DataLoader에서 pin_memory=True로 설정하면 매 배치 데이터가 자동으로 핀 메모리에 적재되어 GPU로의 전송이 최적화됩니다. 다만, 핀 메모리는 시스템 RAM 자원을 많이 소모할 수 있고 과도한 경우 오히려 성능 저하를 일으킬 수 있으므로 필요한 경우에만 활용하는 것이 좋습니다.

3. 멀티 GPU 환경에서의 메모리 관리

한 대의 머신에 여러 GPU가 있는 경우, PyTorch는 주로 두 가지 방식으로 데이터를 병렬 처리합니다: DataParallel과 DistributedDataParallel입니다.

DataParallel vs DistributedDataParallel

torch.nn.DataParallel (DP)는 단일 프로세스-멀티 스레드 기반 데이터 병렬화 기법입니다. 하나의 프로세스 내에서 여러 GPU를 사용하며, PyTorch가 미니배치를 GPU들로 나누어 각자 forward 연산을 수행하고, 결과를 다시 모아(loss 계산 등을 위해) 기본 장치에 모읍니다. DP는 구현이 비교적 간단하고 코드 수정이 적지만, 매 iteration마다 모델을 각 GPU로 복제(scatter)하고 결과를 수집(gather)하는 부가 연산이 있습니다. 또한 Python GIL(Global Interpreter Lock)로 인해 멀티스레드 간 병렬 실행이 제약을 받아 성능이 떨어질 수 있습니다. 메모리 측면에서는 각 GPU마다 모델 파라미터 복사본이 필요하므로 GPU당 메모리 사용량은 늘어나지만, 이는 멀티 GPU 병렬화의 본질적인 비용입니다. DP에서는 주로 첫 번째 GPU에 gradient 집계와 parameter update를 수행하므로, GPU 0번에 추가 메모리 부담과 통신 부담이 생길 수 있습니다.
 
torch.nn.parallel.DistributedDataParallel (DDP)는 멀티 프로세스-멀티 GPU 병렬화 방법입니다. 일반적으로 GPU마다 하나씩 개별 프로세스를 띄워 동일한 모델을 로드하고, 각 프로세스가 독립적으로 forward/backward를 수행한 후 NCCL 등의 백엔드를 통해 gradient를 all-reduce(모든 GPU의 gradient 평균)하여 동기화합니다. DDP는 구현상 약간의 설정이 필요하지만, PyTorch에서는 torch.distributed.launch나 torchrun을 통해 쉽게 실행할 수 있습니다. DDP의 큰 장점은 프로세스별 실행으로 GIL 문제를 피하고, 고성능 네트워크 통신(NCCL)으로 gradient를 교환하여 DataParallel보다 훨씬 효율적이라는 것입니다. 단일 머신에서도 DDP가 DP보다 빠르며, 멀티 머신까지 확장 가능한 유연성을 가집니다. 메모리 측면에서는 여전히 GPU당 모델 복사본이 필요하지만, DP처럼 매번 모델을 복제하지 않고 프로세스 시작 시 한 번만 모델을 해당 GPU로 올리므로 중복 할당을 줄이고 메모리 활용을 최적화합니다. 요약하면, 소규모 실험이나 간단한 경우가 아니면 멀티 GPU 훈련 시 DDP 사용이 권장됩니다

GPU 간 통신과 NCCL

텐서 할당과 장치 간 데이터 이동

멀티 GPU 환경에서는 각 GPU가 계산한 결과(특히 gradient)를 서로 주고받는 통신 비용이 중요합니다. NVIDIA의 NCCL (Nvidia Collective Communications Library)은 이러한 다중 GPU 간 통신을 최적화하는 라이브러리입니다.

PyTorch DDP의 기본 통신 백엔드가 바로 NCCL이며, NCCL은 GPU들 간에 고속의 All-Reduce 연산(모든 장치의 값을 모아 합산/평균 후 브로드캐스트) 등을 효율적으로 수행합니다. 예를 들어 4개의 GPU가 있다면, NCCL의 ring-allreduce 알고리즘을 통해 각 GPU가 자신의 gradient를 이웃 GPU로 주고받으며 합산하여 최종 모든 GPU에 동일한 평균 gradient를 얻게 됩니다.

NCCL은 PCIe 또는 NVLink를 통해 GPU 간 DMA 전송을 수행하며, 가능하다면 peer-to-peer GPU Direct 기술을 활용해 중앙 CPU 개입 없이 장치 간 직접 통신을 합니다. 이를 통해 멀티 GPU 훈련 시 발생하는 통신 지연을 최소화하고, 거의 선형에 가까운 속도 향상을 얻을 수 있습니다. 실무에서는 별도로 NCCL을 다룰 필요 없이 PyTorch DDP를 사용하면 내부적으로 최적화된 NCCL 통신이 이뤄집니다.

4. 메모리 최적화 기법

딥러닝 모델 학습 시 메모리를 최적으로 활용하면 더 큰 배치 크기를 사용하거나 모델을 더 깊게 쌓을 수 있습니다. PyTorch에서는 몇 가지 유용한 기법과 도구를 제공하고 있으며, 사용자가 할 수 있는 튜닝 방법도 있습니다.
 
torch.cuda.empty_cache() 활용: PyTorch는 앞서 설명한 대로 GPU 메모리를 캐시로 잡아두기 때문에, 필요에 따라 캐시 된 메모리를 비워주는 함수인 torch.cuda.empty_cache()를 제공합니다. 이 함수를 호출하면 현재 사용되지 않는 여유 캐시 메모리를 모두 해제하여 OS에 반환합니다. 이를 통해 다른 애플리케이션이나 PyTorch 외의 코드가 GPU 메모리를 사용할 수 있도록 할 수 있지만, 이미 할당되어 있는 텐서의 메모리는 free 되지 않습니다. 따라서 empty_cache()는 메모리 누수를 직접 해결해주지는 못하며, PyTorch 내에서 실제로 사용 중인 메모리를 줄여주진 않습니다. 주로 큰 모델을 돌리다가 중간에 대용량의 임시 텐서를 삭제한 후, 남은 캐시를 비워서 nvidia-smi 상의 점유율을 낮추거나, 다음 메모리 할당 시 한꺼번에 큰 연속 공간을 확보해야 할 때 유용합니다. 하지만 empty_cache()를 남용할 필요는 없습니다. PyTorch는 캐시를 재사용하므로 일반적으로 캐시를 비워주지 않아도 필요하면 기존 메모리를 재활용하며, 빈번한 empty_cache() 호출은 오히려 성능에 악영향을 줄 수 있습니다.
 
Mixed Precision Training (FP16/BF16): 혼합 정밀도 훈련은 메모리를 절약하고 연산 속도를 높이는 대표적인 기법입니다. 표준 FP32 대신 **절반 정밀도(FP16)**나 **Brain Floating Point (BF16)**로 연산하면, 각 수가 차지하는 비트 수가 절반으로 줄어들어 메모리 사용량도 대폭 감소합니다. 예를 들어 NVIDIA Amp(Automatic Mixed Precision)을 사용하면 모델의 많은 연산을 FP16으로 수행하여 메모리 사용량을 거의 절반으로 줄일 수 있다고 보고된 바 있습니다. 실제로 AMD의 한 실험에서는 AMP 도입으로 훈련 속도가 46% 향상되고, 메모리 사용량이 50% 가까이 감소했다고 합니다.

메모리 사용 모니터링 (torch.cuda.memory_stats(): PyTorch는 메모리 사용량을 모니터링하기 위한 API를 제공합니다. torch.cuda.memory_allocated()는 현재 텐서들이 사용 중인 활성 메모리 양을 바이트 단위로 보여주고, torch.cuda.memory_reserved()는 캐싱 Allocator가 확보한 총 메모리(사용 중인 + 비어있는)를 나타냅니다. 더 나아가 torch.cuda.memory_stats()를 호출하면 각종 세부 메모리 통계를 딕셔너리로 얻을 수 있습니다. 여기에는 현재 할당 횟수, 해제 횟수, 최대 메모리 사용량, 외부 할당 등 다양한 지표가 포함되어 있어 메모리 병목을 분석하는 데 도움이 됩니다. 예를 들어 max_memory_allocated()를 확인하면 학습 중 피크 시 사용된 GPU 메모리의 최대치가 나오므로, 해당 값과 GPU 총량을 비교해 여유 여부를 판단할 수 있습니다. 또한 torch.cuda.memory_snapshot()를 이용하면 메모리 할당 상태를 찍은 스냅숏을 얻을 수도 있는데, 이를 분석하면 어떤 코드 부분이 많은 메모리를 소비했는지 추적할 수 있습니다. 이러한 툴들을 통해 모델 구조나 입력 배치 크기를 조절하여 메모리를 최적화하는 근거를 얻을 수 있습니다.
 
배치 크기 조정과 Gradient Checkpointing: 가장 원초적인 메모리 최적화 방법은 배치 크기(batch size)를 줄이는 것입니다. 배치를 작게 하면 한 번에 유지해야 할 활성 메모리양(입력, 출력, 중간 활성값 등)이 줄어들어 OOM 위험이 감소합니다. 배치 크기를 줄이는 대신 학습 반복 횟수를 늘리거나 Gradient Accumulation(그래디언트 누적)을 통해 동일한 효과를 내도록 조정합니다. 또한 Gradient Checkpointing 기법을 사용하면 모델의 메모리 사용량을 크게 줄일 수 있습니다. PyTorch의 torch.utils.checkpoint를 통해 일부 레이어의 순전파 중간 결과를 저장하지 않고 필요할 때 역전파 시 재계산하게 할 수 있습니다. 이렇게 하면 저장해야 할 활성화 메모리가 감소하여, 메모리 사용량을 획기적으로 줄일 수 있지만 대신 역전파 시 해당 부분을 다시 계산하므로 연산량이 증가합니다.

5. CUDA 메모리 디버깅 및 문제 해결

마지막으로, CUDA 메모리 부족(OOM) 오류를 만났을 때 대처 방법과 PyTorch가 제공하는 디버깅 도구를 알아보겠습니다. 아무리 최적화를 해도 모델이 크거나 배치가 너무 크면 OOM이 발생할 수 있는데, 이때 원인을 파악하고 해결하는 역량이 중요합니다.
 
OOM 오류 대처법: "RuntimeError: CUDA out of memory. Tried to allocate..."와 같은 오류가 뜨면 우선 메모리 여유분이 있는지 확인합니다. 만약 GPU 전체 용량에 비해 여유 공간이 조금이라도 있다면, 메모리 단편화로 인한 문제일 수도 있습니다 (조각난 작은 빈 공간들은 많지만 필요한 연속 큰 공간을 못 구하는 경우). 이럴 때 torch.cuda.empty_cache()를 호출하여 캐시 메모리를 비워보거나, PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:... 환경변수를 조절해 큰 블록 분할을 제한해 볼 수 있습니다. 그러나 대부분의 OOM은 단순히 모델/데이터가 과하게 큰 경우이므로, 해결책은 배치 크기 축소가 됩니다. 배치를 줄이면 그만큼 한 번에 드는 메모리가 감소하여 OOM이 사라질 수 있습니다. 또 다른 방법으로는 모델의 일부를 FP16으로 변경하거나, 앞서 언급한 Gradient Checkpointing을 적용하여 메모리 사용량을 줄이는 것이 있습니다.

코드 레벨에서는 사용이 끝난 중간 변수에 대해 del tensor로 참조를 삭제하고 gc.collect()를 호출하여 Python이 해당 객체를 메모리에서 해제하도록 유도할 수 있습니다. 다만 이 경우도 PyTorch 캐싱 때문에 바로 OS에 반환되지는 않지만, 참조가 남아있어 메모리를 붙잡는 현상(메모리 누수)을 방지할 수 있습니다. 만약 학습 루프 외부에 GPU 텐서를 리스트 등에 쌓아두고 있지는 않은지 확인하여, 불필요한 저장을 피해야 합니다. 요약하면, OOM이 발생하면 (1) 배치 크기 축소, (2) 불필요한 객체 해제, (3) 정밀도 낮추기 or 체크포인팅 등의 순서로 시도해 보고, 그래도 안 되면 모델 크기 자체를 줄이는 것도 고려해야 합니다.
 
메모리 상태 요약 (torch.cuda.memory_summary()): PyTorch는 메모리 디버깅을 돕기 위해 torch.cuda.memory_summary() 함수를 제공합니다. 이 함수를 호출하면 현재 GPU 메모리 할당자(Caching Allocator)의 상태를 사람이 읽기 쉬운 형태로 출력해 줍니다. 출력에는 각 크기별 블록 풀의 현황, 할당된 블록과 빈 블록의 개수, 누적 할당량/해제량, 최대 사용량 등의 정보가 포함됩니다. 이를 통해 메모리 단편화 정도를 파악할 수 있습니다. 예를 들어 요약에 "inactive_split" 블록이 많이 보인다면, 과거에 큰 블록을 잘게 쪼갠 뒤 남은 조각들이 많다는 의미일 수 있습니다. memory_summary의 결과를 분석함으로써, 특정 크기의 텐서 할당 패턴이 문제를 일으키는지 알 수 있고, 앞서 언급한 환경변수 조정이나 코드 구조 변경의 실마리를 얻을 수 있습니다. 또한 이 함수는 주기적으로 호출하여 학습 도중 메모리 추이를 모니터링하는 용도로도 쓸 수 있습니다. 더 나아가 PyTorch 2.0+ 버전에서는 메모리 스냅숏을 시각화하는 도구(memory_viz)도 제공되어, 웹 브라우저로 메모리 사용 히스토리를 탐색하며 어느 시점에 메모리가 급증했는지 등을 분석할 수 있습니다. 이러한 진단 정보를 활용하면, 메모리 문제를 일으키는 코드 부분을 찾아내어 최적화하거나 리팩터링 하는 데 큰 도움이 됩니다.
 
PyTorch Profiler를 이용한 메모리 분석: PyTorch에는 연산별로 시간과 메모리 사용량을 측정할 수 있는 프로파일러(torch.profiler)가 있습니다. 프로파일러를 사용하면 모델의 어떤 연산(op)들이 많은 메모리를 할당하는지 세부적으로 확인할 수 있습니다. 예를 들어 profiler를 profile_memory=True 옵션과 함께 실행하면, forward 한 번을 도는 동안 각 연산자가 얼마만큼의 CUDA 메모리를 할당 또는 해제했는지 표로 볼 수 있습니다. 이때 self_gpu_memory_usage (혹은 self_cuda_memory_allocated) 같은 지표가 각 연산의 자체 메모리 사용량을 나타내며, 이를 통해 메모리 집약적인 연산이 어디인지 pinpoint 할 수 있습니다. PyTorch Profiler는 TensorBoard와 연계하여 시각적으로 시간·메모리 프로파일을 볼 수도 있으므로, 성능 최적화와 함께 메모리 최적화에도 유용하게 쓰입니다.