beautifulsoups.dev
Published on

Transformers 로컬 LLM 멈춤 - KV 캐시 누수·파편화 해결

Authors

로컬에서 transformers로 LLM을 서빙하다 보면, 처음엔 빠르다가 시간이 지날수록 토큰 생성이 급격히 느려지거나(초당 토큰 수 급락), 심하면 프로세스가 멈춘 것처럼 보이는 현상을 겪습니다. 겉으로는 GPU 사용률이 떨어지고 VRAM이 애매하게 남아 있는데도 CUDA out of memory가 나거나, 반대로 VRAM이 계속 증가하는데 회수가 안 되는 모습이 관찰되기도 합니다.

이 글은 그 증상을 KV 캐시(Key/Value cache) 누수 혹은 파편화 관점에서 설명하고, transformers 기반 로컬 LLM에서 안정적으로 장시간 돌리기 위한 설정/코드 패턴을 정리합니다.

또한 “캐시”를 시스템적으로 다루는 관점은 RAG 캐시 계층 설계와도 맞닿아 있습니다. 필요하다면 Production RAG 벡터 DB 캐시 계층 설계와 튜닝도 함께 참고하면 좋습니다.

1) KV 캐시가 왜 문제를 만들까

KV 캐시 기본 동작

디코딩(오토리그레시브 생성)에서 매 토큰마다 self-attention은 과거 토큰들의 K/V를 다시 계산하면 비용이 큽니다. 그래서 use_cache=True일 때 모델은 레이어별로 K/V를 저장해두고 다음 토큰에서 재사용합니다. 이게 KV 캐시입니다.

KV 캐시의 메모리 사용량은 대략 다음에 비례합니다.

  • batch_size
  • num_layers
  • num_heads 또는 hidden_size
  • sequence_length(프롬프트 길이 + 생성 길이)
  • dtype(fp16, bf16, fp32)

즉 “길게 대화할수록, 컨텍스트가 커질수록” 캐시는 커집니다. 여기까지는 정상입니다.

누수처럼 보이는 현상 1: 컨텍스트가 계속 누적되는 설계

챗봇에서 대화 히스토리를 계속 붙여 넣으면, 각 요청의 sequence_length가 점점 커지고 KV 캐시도 그만큼 커집니다. 이건 누수라기보다 업무 로직이 컨텍스트를 무한히 키우는 구조입니다.

누수처럼 보이는 현상 2: GPU 메모리 캐싱 allocator + 파편화

PyTorch는 성능을 위해 GPU 메모리를 캐싱합니다. torch.cuda.empty_cache()를 호출해도 “드라이버에 즉시 반환”이 아닌 경우가 많고, 다양한 크기의 텐서가 반복 생성/해제되면 파편화(fragmentation) 로 인해 “총량은 남는데 큰 연속 블록이 없어 할당 실패”가 발생할 수 있습니다.

누수처럼 보이는 현상 3: 그래프/텐서 참조가 남는 진짜 누수

다음과 같은 경우는 실제로 메모리 회수를 막을 수 있습니다.

  • past_key_values를 전역 리스트에 저장하거나, 세션 객체에 계속 누적
  • generate 결과(특히 output_scores, output_attentions, output_hidden_states)를 켠 채로 로그/리플레이 용도로 쌓음
  • torch.no_grad() 또는 torch.inference_mode() 없이 추론을 반복해 그래프가 남음

2) 증상별 빠른 체크리스트

아래 질문에 “예”가 많을수록 KV 캐시/파편화 이슈 가능성이 큽니다.

  • 요청이 반복될수록 nvidia-smi에서 VRAM 사용량이 계단식으로 증가한다
  • VRAM 총량은 남아 있는데도 OOM이 난다(특히 긴 프롬프트에서)
  • 토큰 생성 속도가 시간이 갈수록 지속적으로 떨어진다
  • 배치 크기나 max_new_tokens를 줄이면 일시적으로 정상화된다
  • 동일 프로세스에서 모델을 내렸다 올리면(재시작) 즉시 해결된다

3) 관측: “무엇이 커지는지”부터 분리하자

GPU 메모리 관측 코드

아래는 요청 전후로 GPU 메모리 변화를 찍어 보는 최소 코드입니다.

import torch

def gpu_mem(tag: str = ""):
    if not torch.cuda.is_available():
        return
    torch.cuda.synchronize()
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserv = torch.cuda.memory_reserved() / 1024**2
    peak = torch.cuda.max_memory_allocated() / 1024**2
    print(f"[{tag}] allocated={alloc:.1f}MB reserved={reserv:.1f}MB peak={peak:.1f}MB")

# 사용 예
# gpu_mem("before")
# ... inference ...
# gpu_mem("after")
  • allocated는 실제 텐서가 점유 중인 메모리
  • reserved는 PyTorch 캐싱 allocator가 확보해 둔 풀

allocated가 계속 증가하면 참조가 남는 누수 가능성이 높고, reserved만 커지면 파편화/캐싱 정책 이슈일 가능성이 큽니다.

긴 컨텍스트가 원인인지 확인

대화 히스토리 길이를 로그로 남기세요.

# input_ids 길이 기록
seq_len = input_ids.shape[-1]
print("seq_len=", seq_len)

seq_len이 요청마다 늘어난다면 “KV 캐시가 커져서 느려지는 것”은 정상적인 결과일 수 있습니다. 이때는 요약/윈도잉이 정답입니다.

4) 해결 1: 추론 모드 고정으로 그래프 누수 차단

추론은 반드시 torch.inference_mode()를 사용하세요. no_grad()보다 더 공격적으로 autograd 관련 오버헤드를 줄이고, 실수로 그래프가 남는 상황을 줄여줍니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "your-model"

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)
model.eval()

@torch.inference_mode()
def generate_once(prompt: str):
    inputs = tok(prompt, return_tensors="pt").to(model.device)
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        use_cache=True,
        return_dict_in_generate=False,
    )
    return tok.decode(out[0], skip_special_tokens=True)

추가로 다음 옵션은 웬만하면 끄세요.

  • output_attentions=True
  • output_hidden_states=True
  • return_dict_in_generate=True를 쓰더라도, 결과를 장기간 저장하지 않기

5) 해결 2: 컨텍스트 윈도잉(슬라이딩) 또는 요약

채팅 서버에서 “대화 전체를 계속 붙이기”는 KV 캐시를 계속 키웁니다. 가장 효과적인 해법은 토큰 예산을 정하고 자르는 것입니다.

def truncate_to_budget(input_ids, budget: int):
    # budget을 넘으면 앞부분을 자름
    if input_ids.shape[-1] <= budget:
        return input_ids
    return input_ids[:, -budget:]

@torch.inference_mode()
def chat_turn(history_text: str, user_text: str, budget: int = 4096):
    prompt = history_text + "\nUser: " + user_text + "\nAssistant:"
    inputs = tok(prompt, return_tensors="pt")
    input_ids = truncate_to_budget(inputs["input_ids"], budget)
    inputs = {"input_ids": input_ids.to(model.device)}

    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        use_cache=True,
    )
    return tok.decode(out[0], skip_special_tokens=True)

요약을 섞는 전략(예: 최근 N턴은 원문 유지, 오래된 내용은 요약으로 치환)은 장시간 세션에서 특히 효과적입니다.

6) 해결 3: max_new_tokens와 배치 전략을 “안전한 상한”으로

KV 캐시는 생성 길이에도 비례합니다. “답이 길어질 수 있는 요청”을 그대로 허용하면 특정 요청 하나가 VRAM을 크게 흔들어 파편화를 유발할 수 있습니다.

  • max_new_tokens 상한을 보수적으로
  • 스트리밍 응답을 하더라도 내부 생성 상한은 유지
  • 동시 요청 배치(batch_size)는 VRAM과 컨텍스트 예산을 고려해 제한

서버 관점에서는 “최대 프롬프트 토큰 + 최대 생성 토큰”을 하나의 SLA로 묶어 관리하는 게 안전합니다.

7) 해결 4: 파편화 완화(allocator 설정과 워밍업)

PyTorch CUDA allocator 설정

파편화가 의심되면 PYTORCH_CUDA_ALLOC_CONF로 분할 정책을 조정해 볼 수 있습니다. 아래는 흔히 쓰는 예시입니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,garbage_collection_threshold:0.8"
  • max_split_size_mb를 낮추면 큰 블록을 잘게 쪼개는 것을 제한해 파편화를 줄이는 데 도움이 될 수 있습니다.
  • garbage_collection_threshold는 캐시 정리 트리거에 영향을 줍니다.

환경과 워크로드에 따라 최적값이 달라 A/B 테스트가 필요합니다.

워밍업으로 “최대 형태”를 먼저 할당

서버 시작 시 대표적인 최대 컨텍스트 길이로 한 번 워밍업을 수행하면, 이후 비슷한 크기의 할당이 재사용되어 파편화가 줄어드는 경우가 있습니다.

@torch.inference_mode()
def warmup(max_len: int = 4096):
    dummy = "hello " * 1000
    inputs = tok(dummy, return_tensors="pt", truncation=True, max_length=max_len).to(model.device)
    _ = model.generate(**inputs, max_new_tokens=8, do_sample=False, use_cache=True)

warmup()

8) 해결 5: 세션별 KV 캐시를 “명시적으로” 다루기

일부 구현은 generate 호출을 매번 새로 하는 대신, 직접 past_key_values를 받아 다음 토큰 생성에 이어붙이기도 합니다. 이때 세션 캐시를 잘못 관리하면 실제 누수가 됩니다.

안티패턴: 세션 객체에 past_key_values를 무한 보관

  • 세션이 종료되지 않거나
  • 세션 수가 늘어나는 환경에서

past_key_values를 계속 들고 있으면 VRAM이 계속 잠식됩니다.

권장: 세션 TTL과 토큰 예산

  • 세션에 TTL을 두고 만료 시 캐시 폐기
  • 세션별 최대 컨텍스트/생성 토큰 예산을 두고 초과 시 요약 또는 리셋

아래는 “세션 만료 시 참조 제거”의 최소 예시입니다.

import time

class SessionCache:
    def __init__(self, ttl_sec: int = 600):
        self.ttl_sec = ttl_sec
        self.store = {}  # session_id -> (past_key_values, last_ts)

    def get(self, sid):
        item = self.store.get(sid)
        if not item:
            return None
        pkv, ts = item
        if time.time() - ts > self.ttl_sec:
            del self.store[sid]
            return None
        return pkv

    def set(self, sid, pkv):
        self.store[sid] = (pkv, time.time())

    def drop(self, sid):
        if sid in self.store:
            del self.store[sid]

만약 파편화가 심하고 세션 수가 많다면, “세션 캐시 자체를 GPU가 아니라 CPU로 옮기기”는 보통 성능상 손해가 커서 권장하진 않습니다. 대신 세션 수를 제한하거나, 모델 서버를 수평 확장하는 편이 낫습니다.

9) 해결 6: torch.compile과 attention 구현 선택(부작용 점검)

환경에 따라 torch.compile 또는 특정 attention 커널(SDPA, FlashAttention 등)이 메모리 패턴에 영향을 줍니다.

  • 어떤 조합은 메모리를 더 아끼지만, 특정 길이에서 워크스페이스가 커져 순간 피크가 튈 수 있음
  • 커널 선택이 바뀌면 텐서 할당 패턴도 바뀌어 파편화 양상이 달라질 수 있음

따라서 “멈춤”이 재현된다면 다음을 바꿔가며 원인을 좁히세요.

  • torch.backends.cuda.enable_flash_sdp(True/False)
  • torch.backends.cuda.enable_mem_efficient_sdp(True/False)
  • torch.backends.cuda.enable_math_sdp(True/False)

주의: 이 설정은 PyTorch 버전, GPU 아키텍처에 따라 동작이 다를 수 있습니다.

10) 최후의 수단: 주기적 프로세스 리사이클링

파편화는 “완벽히 제어하기 어려운” 영역이 있습니다. 특히 다양한 길이의 요청이 섞이는 공개형 서비스라면, 장시간 운영에서 파편화가 누적될 수 있습니다.

그럴 때는 다음 운영 패턴이 현실적인 안전장치가 됩니다.

  • 워커 프로세스를 여러 개 띄우고
  • 일정 요청 수 또는 일정 시간마다 워커를 graceful하게 교체

이 방식은 웹 서버/잡 워커 운영에서 흔한 전략이고, 누수/파편화에 강합니다. 누수 진단 관점은 Go 고루틴 누수 5분 진단 - pprof·채널닫기처럼 “참조가 남는지”를 먼저 확인하는 접근과 유사합니다.

11) 실전 구성 예시: 안전한 로컬 LLM 루프

아래 예시는 “긴 컨텍스트 제한 + 추론 모드 + 관측 + 최소 결과만 보관”을 한 번에 적용한 형태입니다.

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128,garbage_collection_threshold:0.8")

model_id = "your-model"

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)
model.eval()

BUDGET = 4096

def gpu_mem(tag: str = ""):
    if not torch.cuda.is_available():
        return
    torch.cuda.synchronize()
    a = torch.cuda.memory_allocated() / 1024**2
    r = torch.cuda.memory_reserved() / 1024**2
    print(f"[{tag}] alloc={a:.0f}MB reserv={r:.0f}MB")

def truncate_ids(input_ids, budget: int):
    if input_ids.shape[-1] <= budget:
        return input_ids
    return input_ids[:, -budget:]

@torch.inference_mode()
def answer(prompt: str, max_new_tokens: int = 256):
    gpu_mem("before")
    inputs = tok(prompt, return_tensors="pt")
    input_ids = truncate_ids(inputs["input_ids"], BUDGET).to(model.device)

    out = model.generate(
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        use_cache=True,
        return_dict_in_generate=False,
    )

    text = tok.decode(out[0], skip_special_tokens=True)
    gpu_mem("after")
    return text

if __name__ == "__main__":
    while True:
        q = input("prompt: ").strip()
        if q in {"quit", "exit"}:
            break
        print(answer(q))

이 루프에서 alloc이 계속 증가한다면, 코드 외부(로깅/세션 저장/리스트 누적)에서 텐서 참조가 남는지 확인해야 합니다. 반대로 alloc은 안정적인데 reserv만 커지고 OOM이 난다면 파편화/피크 메모리 문제일 확률이 큽니다.

12) 정리: 어떤 처방이 어떤 원인에 맞나

  • 컨텍스트가 계속 늘어남: 윈도잉/요약/토큰 예산이 1순위
  • allocated가 계속 증가: 참조 누수 차단, 결과/스코어 저장 금지, inference_mode 강제
  • reserved만 커지고 큰 요청에서 OOM: 파편화 완화, allocator 설정, 워밍업, 요청 길이 상한
  • 운영 중 간헐적 멈춤: 워커 리사이클링으로 안전장치

로컬 LLM “멈춤”은 한 가지 원인으로만 설명되지 않는 경우가 많습니다. 하지만 관측 지표를 allocated와 reserved, 그리고 seq_len으로 쪼개서 보면 대부분의 케이스에서 원인이 빠르게 드러납니다. 그 다음은 토큰 예산과 캐시 수명(TTL)을 설계해, KV 캐시가 ‘성능을 위한 캐시’로 남도록 만드는 것이 핵심입니다.

Discuss on Twitter • View on GitHub