강화학습 포트폴리오 모델 학습 중 발생한 PyTorch OOM(Out of Memory) 메모리 누수 해결기

최근 심층 강화학습(Deep Reinforcement Learning) 기반의 동적 자산 배분 모델을 고도화하는 과정에서 아주 골치 아픈 문제를 겪었습니다. 로컬 환경에서 에피소드(Episode)를 몇 번 돌릴 때는 멀쩡하던 파이썬(Python) 학습 스크립트가, 클라우드 서버에 올려 대규모 백테스트 데이터를 먹이자마자 메모리가 미친 듯이 치솟더니 결국 프로세스가 강제 종료(Killed)되는 현상이 발생했습니다.

오늘은 퀀트 투자 모델링이나 딥러닝 파이프라인을 구축할 때 흔히 마주치는 PyTorch 기반의 메모리 누수(Memory Leak) 원인과 해결 방법을 짧게 공유하고자 합니다.

🚨 문제 상황: RAM 용량이 우주로 돌파하다

에러 로그를 확인해보니 전형적인 OOM(Out of Memory) 에러였습니다. RuntimeError: CUDA out of memory (또는 Linux OOM Killer 개입)

자산 배분 모델의 특성상 과거 20년 치의 금융 틱 데이터를 Replay Buffer에 담아두고 샘플링해야 하므로 메모리를 많이 쓰는 것은 당연하지만, 학습 에포크(Epoch)가 진행될수록 메모리 사용량이 계단식으로 계속 누적해서 상승하는 것이 문제였습니다. 이는 명백한 ‘메모리 누수’ 증상이었습니다.

🔍 원인 분석: 텐서(Tensor)에 숨겨진 연산 그래프

코드를 한 줄씩 디버깅해 본 결과, 원인은 모델의 학습 상태(Loss값, Reward 등)를 로그로 남기기 위해 파이썬 리스트(List)에 데이터를 쌓아두는 부분에 있었습니다.

Python

# 문제가 된 기존 코드 (오류 원인)
epoch_losses = []
for data in dataloader:
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    
    # 여기서 메모리 누수 발생!
    epoch_losses.append(loss)

PyTorch에서 loss는 단순한 숫자(Float) 값이 아닙니다. 역전파(Backpropagation) 연산을 위해 해당 Loss가 도출되기까지의 모든 신경망 연산 그래프(Computational Graph) 이력을 꼬리표처럼 달고 있는 무거운 객체입니다. 위 코드처럼 loss 객체 자체를 리스트에 계속 append 해버리면, 가비지 컬렉터(Garbage Collector)가 지난 에포크의 연산 그래프를 메모리에서 해제하지 못하고 계속 쥐고 있게 됩니다. 결국 에포크가 돌 때마다 거대한 연산 그래프가 램(RAM)에 차곡차곡 쌓여 OOM이 터진 것입니다.

🛠️ 해결 과정: .item()과 .detach()의 활용

해결책은 허무할 정도로 간단했습니다. 연산 그래프의 연결 고리를 끊어주고, 순수한 스칼라(Scalar) 숫자값만 추출해서 저장하도록 코드를 수정했습니다.

Python

# 수정한 코드 (정상 동작)
epoch_losses = []
for data in dataloader:
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    
    # .item()을 사용하여 순수 float 값만 추출하여 저장
    epoch_losses.append(loss.item()) 

추가로, Replay Buffer에 특정 상태(State) 값을 밀어 넣을 때도 기울기(Gradient) 계산이 필요 없는 데이터라면 .detach().cpu().numpy()를 명시적으로 호출하여 연산 그래프로부터 분리하고 메모리 부담을 최소화했습니다.

💡 실무 인사이트

딥러닝 기반의 알고리즘 트레이딩이나 인프라를 다루다 보면, 로직 자체의 수학적 완결성 못지않게 ‘컴퓨팅 자원의 효율적 관리’가 뼈저리게 중요함을 느낍니다. 사소한 .item() 메서드 하나를 누락한 대가로 클라우드 인스턴스가 터져버리는 것을 보며, 메모리 프로파일링의 중요성을 다시금 새기게 되었습니다. 강화학습 모델을 구축하시는 분들은 루프문 내부에서 텐서를 컬렉션에 담을 때 반드시 연산 그래프 분리 여부를 확인하시기 바랍니다.

이 글은 기존 자산 배분/강화학습 주제와 자연스럽게 결을 같이 하면서도, ‘기술적 오류 해결(Troubleshooting)’이라는 완전히 독립적이고 실용적인 성격을 띱니다.

댓글 남기기