BentoML GPU 서빙에서 VRAM 누수 잡는 법

서빙이 멀쩡히 돌아가는데도 VRAM이 요청을 처리할수록 조금씩 올라가 결국 OOM으로 죽는 경우가 있습니다. 특히 BentoML 같은 모델 서빙 프레임워크에서 GPU를 붙이면, 애플리케이션 레벨에서는 “메모리 누수”처럼 보이지만 실제로는 캐시/그래프/레퍼런스 유지/배치 누적/워커 생명주기 같은 복합 요인인 경우가 많습니다.

이 글에서는 BentoML GPU 서빙에서 VRAM 증가를 재현 가능하게 만들고, 관측 지표를 고정한 뒤, 원인별로 차단하는 코드 패턴을 정리합니다. (PyTorch 기준이지만 다른 프레임워크에도 응용 가능합니다.)

참고로, “누수”는 꼭 메모리 객체가 영원히 해제되지 않는 버그만 뜻하지 않습니다. GPU에서는 다음도 모두 “누수처럼 보이는 현상”을 만듭니다.

• CUDA allocator의 캐싱 정책으로 인해 nvidia-smi 상 사용량이 내려가지 않음
• inference가 아닌 그래프가 쌓이거나(autograd), 텐서 참조가 남아 해제가 지연됨
• 동적 shape로 인해 매 요청마다 새로운 workspace/커널 캐시가 늘어남
• BentoML 워커/프로세스 모델과 모델 로딩 위치가 맞지 않아 중복 적재

1) 먼저 “진짜 누수”인지 “캐시”인지 구분하기

GPU 문제를 잡을 때 가장 먼저 할 일은 관측 기준을 통일하는 것입니다. nvidia-smi는 “현재 점유”를 보여주지만, PyTorch는 내부적으로 메모리를 캐시합니다. 그래서 텐서를 해제해도 nvidia-smi가 즉시 내려가지 않을 수 있습니다.

PyTorch에서는 최소 아래 3가지를 같이 봐야 합니다.

• torch.cuda.memory_allocated() : 텐서가 실제로 점유 중인 메모리
• torch.cuda.memory_reserved() : allocator가 잡아둔 캐시 포함 예약 메모리
• torch.cuda.max_memory_allocated() : 피크 추적(요청별로 리셋하면 유용)

관측 유틸리티 코드

아래는 BentoML 서비스 코드에 그대로 넣어도 되는 “요청 단위 VRAM 스냅샷” 예시입니다.

import os
import time
import torch


def gpu_mem_snapshot(tag: str = "") -> dict:
    if not torch.cuda.is_available():
        return {"tag": tag, "cuda": False}

    torch.cuda.synchronize()
    return {
        "tag": tag,
        "cuda": True,
        "allocated_mb": round(torch.cuda.memory_allocated() / 1024**2, 2),
        "reserved_mb": round(torch.cuda.memory_reserved() / 1024**2, 2),
        "max_allocated_mb": round(torch.cuda.max_memory_allocated() / 1024**2, 2),
        "ts": time.time(),
        "pid": os.getpid(),
    }

요청 처리 전후로 이 값을 로그로 남기면, 다음을 구분할 수 있습니다.

• allocated_mb가 계속 증가한다: 실제 텐서/그래프/참조가 남는 진짜 누수 가능성
• allocated_mb는 안정적이고 reserved_mb만 증가한다: 캐시/fragmentation/동적 shape 영향

2) BentoML에서 VRAM이 새는 흔한 구조적 원인

BentoML 서빙에서 GPU 메모리 문제가 생기는 지점은 크게 4군데입니다.

1. 모델이 워커마다 중복 로딩됨(프로세스 수와 로딩 위치 문제)
2. inference인데도 autograd 그래프가 쌓임(no_grad 누락)
3. 요청 스코프 밖으로 GPU 텐서 참조가 새어 나감(전역 캐시, 리스트 누적, 로거/트레이서)
4. 동적 shape, 토크나이저/전처리 방식, 배치 전략 때문에 allocator 캐시가 커짐

BentoML 자체가 누수를 만든다기보다는, 서빙 코드가 장기 실행 프로세스에서 안전하지 않은 패턴을 쓰면 누수가 빨리 드러납니다.

3) “autograd 그래프 누적” 차단: inference_mode를 기본값으로

가장 흔한 실수는 inference 경로에서 torch.no_grad() 또는 torch.inference_mode()를 빼먹는 것입니다. 특히 Hugging Face 모델을 감싸거나, 전처리/후처리에서 텐서를 조작하다가 그래프가 연결되면 요청이 쌓일수록 VRAM이 증가합니다.

안전한 기본 패턴

import bentoml
import torch

@bentoml.service(resources={"gpu": 1})
class MyService:
    def __init__(self):
        self.device = "cuda"
        self.model = ...  # 모델 로드
        self.model.to(self.device)
        self.model.eval()

    @bentoml.api
    def predict(self, x: list[float]) -> list[float]:
        before = gpu_mem_snapshot("before")

        with torch.inference_mode():
            t = torch.tensor(x, device=self.device)
            y = self.model(t)
            out = y.detach().float().cpu().tolist()

        after = gpu_mem_snapshot("after")
        # 로그 시스템에 before/after를 남기기
        return out

핵심은 다음 3줄입니다.

• self.model.eval()
• with torch.inference_mode():
• 결과를 즉시 cpu()로 옮기고 Python 객체로 변환

detach()는 습관처럼 넣어도 좋지만, inference_mode만 제대로 쓰면 대부분 충분합니다.

4) “GPU 텐서 참조가 요청 밖으로 남는” 패턴 제거

VRAM 누수의 본질은 보통 “해제되어야 할 텐서가 살아있다”입니다. 서빙 코드에서 자주 발생하는 누수 트리거는 다음과 같습니다.

• 전역 리스트/캐시에 결과 텐서를 저장
• 디버깅 로그에서 텐서를 그대로 들고 있음
• 예외 처리에서 텐서를 포함한 객체를 전역으로 보관
• 스트리밍/제너레이터 구현에서 yield 이후 참조가 남음

나쁜 예: 전역에 GPU 텐서 캐시

CACHE = []

@bentoml.api
def predict(x):
    t = torch.tensor(x, device="cuda")
    y = model(t)
    CACHE.append(y)  # GPU 텐서를 전역에 보관
    return y[0].item()

좋은 예: 캐시가 필요하면 CPU로 내리기

CACHE = []

@bentoml.api
def predict(x):
    with torch.inference_mode():
        t = torch.tensor(x, device="cuda")
        y = model(t)
        y_cpu = y.detach().float().cpu()
    CACHE.append(y_cpu)  # CPU 텐서/넘파이/파이썬 객체만 저장
    return float(y_cpu[0].item())

추가로, “요청별 디버그를 위해 텐서를 저장”해야 한다면 개수 제한을 걸어야 합니다.

from collections import deque

DEBUG_BUFFER = deque(maxlen=100)

5) empty_cache()는 치료제가 아니다: 언제 쓰고 언제 버려야 하나

torch.cuda.empty_cache()는 “PyTorch가 잡고 있는 캐시를 OS로 반환”하려고 시도합니다. 하지만 아래 특성이 있습니다.

• 성능을 떨어뜨릴 수 있음: 다음 요청에서 다시 할당 비용 발생
• 누수를 고치지 못함: allocated가 증가하는 진짜 누수에는 효과 없음
• nvidia-smi 숫자만 내려가서 착시를 줄 수 있음

권장 전략은 다음과 같습니다.

• 기본적으로는 쓰지 않는다
• 다만 동적 shape가 심하고 트래픽 패턴이 들쭉날쭉해서 캐시가 과도하게 커지는 경우, 운영 정책으로 제한적으로 사용한다
  • 예: 요청 N건마다 1회
  • 예: reserved_mb가 특정 임계치를 넘을 때만

임계치 기반 캐시 비우기(운영용)

def maybe_empty_cache(reserved_mb_threshold: int = 20000):
    if not torch.cuda.is_available():
        return
    reserved_mb = torch.cuda.memory_reserved() / 1024**2
    if reserved_mb > reserved_mb_threshold:
        torch.cuda.empty_cache()

이건 “누수 방지”가 아니라 “캐시 폭주 완화”입니다. 진짜 누수는 반드시 원인을 제거해야 합니다.

6) 동적 shape와 배치 전략이 만드는 VRAM 증가(누수처럼 보임)

LLM, diffusion, 멀티모달 모델 서빙에서는 입력 길이/해상도가 요청마다 달라 동적 shape가 됩니다. 이때 커널 선택, workspace, 메모리 풀 사용이 변하면서 reserved가 점진적으로 커질 수 있습니다.

대응책은 다음 중 하나입니다.

• 입력을 버킷팅해서 shape 다양성을 줄임(예: 토큰 길이 버킷)
• 최대 길이/해상도 제한
• BentoML에서 배치 설정을 하되, 배치가 커질 때 VRAM 피크가 올라가는지 관측

BentoML의 배치 기능을 쓴다면, 배치 크기 증가가 max_allocated 피크를 밀어 올려 결국 캐시가 커지는지 확인해야 합니다.

7) BentoML 워커/프로세스 모델: “중복 로딩”을 의심하라

VRAM이 서서히가 아니라 시작부터 예상보다 크게 먹고, 워커 수를 늘리면 VRAM이 워커 수만큼 증가한다면 “누수”가 아니라 프로세스별 모델 중복 적재일 가능성이 큽니다.

체크리스트:

• GPU 1장인데 워커를 여러 개 띄우고 있지 않은가
• 모델 로딩이 프로세스 시작 시점마다 실행되는 구조인가
• 멀티프로세싱 start method가 spawn이라 매번 새로 로딩되는가

운영에서는 보통 GPU 1장당 프로세스 1개(또는 MIG/멀티 GPU 전략)로 단순화하는 게 안전합니다. 멀티 워커가 필요하면 GPU가 아니라 CPU 쪽(전처리)만 스케일링하고, GPU inference는 단일 워커에서 배치로 처리하는 편이 안정적입니다.

8) 누수를 “재현 가능한 테스트”로 만들기: 반복 호출 + 스냅샷

운영에서만 재현되는 VRAM 증가는 잡기 어렵습니다. 가장 좋은 방법은 동일 입력을 1000번 반복하면서 allocated와 reserved의 추세를 로그로 남기는 것입니다.

간단한 부하 스크립트(클라이언트)

import time
import requests

URL = "http://localhost:3000/predict"

for i in range(2000):
    r = requests.post(URL, json={"x": [1.0, 2.0, 3.0]})
    r.raise_for_status()
    if i % 50 == 0:
        print("ok", i)
    time.sleep(0.01)

여기서 서버 로그에 gpu_mem_snapshot 전후를 남기면, “요청 1건당 증가량”이 보입니다. 증가량이 선형이면 참조 누수일 확률이 높고, 어느 지점에서 plateau가 생기면 캐시/fragmentation일 가능성이 큽니다.

9) 예외 경로에서 누수: 실패 요청이 VRAM을 쌓는다

성공 케이스만 테스트하면 놓치는 게 예외 경로입니다. 예외가 발생했을 때 GPU 텐서가 포함된 객체를 로깅/리턴/리트라이 큐에 넣어버리면 누수가 됩니다.

• 예외 로깅 시 텐서를 문자열로 강제 변환하거나 CPU로 내리기
• 재시도 큐에는 원본 입력만 넣고, 중간 텐서는 넣지 않기

재시도/백오프 설계 관점은 GPU 누수와 직접 주제는 다르지만, “실패 트래픽이 시스템을 악화시킨다”는 구조는 동일합니다. 레이트리밋 재시도 설계 글도 함께 보면 운영 안정성에 도움이 됩니다.

• OpenAI Responses API 429 레이트리밋 재시도 설계

10) 운영에서의 최종 방어선: 워커 재시작과 메모리 상한

원인을 제거했더라도, GPU 서빙은 드라이버/커널 캐시/서드파티 라이브러리 영향으로 장기 실행 시 메모리 변동이 있을 수 있습니다. 그래서 마지막 방어선으로 “유출이 감지되면 프로세스를 교체”하는 전략을 둡니다.

• 일정 요청 수 처리 후 graceful restart
• allocated_mb 또는 reserved_mb가 임계치 초과 시 헬스체크 실패 처리 후 재기동
• 쿠버네티스에서는 liveness probe와 결합

쿠버네티스에서 재시작 루프가 생기면 원인 추적이 중요합니다. 아래 글의 접근(이벤트, 로그, 리소스, OOM/exit code 추적)이 GPU 서빙에도 그대로 적용됩니다.

• EKS CrashLoopBackOff 진단 - Pod 재시작 원인 추적

11) 실전 체크리스트(가장 많이 해결되는 순서)

1. inference 경로에 torch.inference_mode() 적용했는가
2. 출력/중간 텐서를 요청 스코프 밖으로 들고 있지 않은가(전역 캐시, 디버그 버퍼)
3. 결과는 즉시 cpu()로 내리고 파이썬 타입으로 변환하는가
4. 동적 shape가 과도한가(버킷팅/제한/배치 전략)
5. 워커 수만큼 모델이 중복 로딩되는 구조인가
6. 예외 경로에서 텐서를 로깅/큐잉하지 않는가
7. allocated와 reserved를 분리 관측했는가

12) 결론: “VRAM 누수”는 관측부터가 반이다

BentoML GPU 서빙에서 VRAM이 증가할 때, 먼저 allocated와 reserved를 분리해서 봐야 합니다. 그 다음으로는 inference_mode 적용, GPU 텐서 참조가 요청 밖으로 새지 않게 차단, 동적 shape 완화, 워커/프로세스 구조 점검 순으로 접근하면 대부분 해결됩니다.

마지막으로, 누수는 코드 한 줄로 끝나는 문제가 아니라 “장기 실행 프로세스의 생명주기 관리” 문제이기도 합니다. 비정상 증가를 감지하고 재시작으로 격리하는 운영 안전장치까지 갖추면, GPU 서빙을 훨씬 안정적으로 운영할 수 있습니다.