Transformers 로컬 LLM 스트리밍 끊김·지연 해결

로컬에서 transformers 기반 LLM을 띄우고 stream=True 느낌의 토큰 스트리밍을 붙였는데, 출력이 뚝뚝 끊기거나 한참 있다가 한 번에 쏟아지는 경우가 자주 있습니다. 특히 Gradio/Streamlit/FastAPI 같은 웹 레이어를 얹으면 체감이 더 심해지고, GPU 사용률은 낮은데 응답은 느린 “이상한” 상태가 나오기도 합니다.

이 글은 스트리밍 끊김·지연을 원인별로 분해하고, 가장 효과가 큰 순서대로 고치는 방법을 제공합니다. 핵심은 단순히 모델을 바꾸는 게 아니라, 다음 4가지 병목을 구분하는 것입니다.

• 생성 루프가 토큰을 “만드는 속도”가 느리다 (순수 추론 병목)
• 토큰은 만들어지는데 “전달”이 늦다 (스트리머/스레드/이벤트루프 병목)
• 전달은 되는데 “클라이언트 렌더링”이 늦다 (UI/네트워크 버퍼링)
• 부하가 걸리면 간헐적으로 멈춘다 (메모리/스케줄링/GC/OOM)

아래 체크리스트를 위에서 아래로 적용하면 대부분 개선됩니다.

1) 먼저 증상을 3가지로 분류하기

A. “한참 무응답 후 한 번에 출력”

• 스트리밍이 아니라 버퍼링이 일어나는 전형적인 패턴입니다.
• 원인 후보: TextStreamer 대신 TextIteratorStreamer 미사용, 웹 프레임워크가 flush를 안 함, print 버퍼링, SSE/WebSocket 구현 문제.

B. “토큰이 나오긴 하는데 일정 주기로 멈칫”

• 원인 후보: CPU로 일부 연산이 떨어짐, torch.compile 워밍업, KV 캐시/메모리 재할당, Python GIL 경쟁, 백그라운드 스레드 스케줄링.

C. “처음 몇 토큰은 빠른데 점점 느려짐”

• 원인 후보: KV 캐시가 커지면서 메모리 대역폭 병목, max_new_tokens 과다, 샘플링 파라미터로 연산 증가, 컨텍스트 길이 증가.

이 분류가 중요한 이유는, A는 “전달” 문제고 B/C는 “생성” 문제인 경우가 많아서 처방이 달라지기 때문입니다.

2) Transformers에서 올바른 스트리밍 패턴 사용하기

transformers에서 가장 흔한 실수는 생성을 메인 스레드에서 돌리고, 스트리밍을 같은 흐름에서 처리하려다 이벤트 루프가 막히는 것입니다. 안정적인 기본형은 아래처럼 TextIteratorStreamer + 별도 스레드입니다.

import threading
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer

model_id = "meta-llama/Llama-3.1-8B-Instruct"  # 예시

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)
model.eval()

prompt = "한국어로 로컬 LLM 스트리밍이 끊기는 이유를 설명해줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

streamer = TextIteratorStreamer(
    tokenizer,
    skip_prompt=True,
    skip_special_tokens=True,
)

gen_kwargs = dict(
    **inputs,
    streamer=streamer,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
    top_p=0.9,
    use_cache=True,
)

t = threading.Thread(target=model.generate, kwargs=gen_kwargs)
# 데몬으로 두면 프로세스 종료 시 정리되지만, 서버에서는 join 전략을 명확히 하세요.
t.daemon = True

t.start()

for text in streamer:
    # 여기서 stdout, SSE, WebSocket 등으로 즉시 flush
    print(text, end="", flush=True)

t.join()

포인트는 다음과 같습니다.

• TextStreamer는 편하지만 반복자로 토큰을 뽑기 어려워, 웹 스트리밍에선 TextIteratorStreamer가 유리합니다.
• print(..., flush=True)는 로컬 콘솔에서도 “한 번에 쏟아짐”을 줄입니다.
• model.generate는 블로킹이므로 메인 스레드에서 돌리면 웹 이벤트루프가 멈출 수 있습니다.

3) FastAPI에서 SSE로 끊김 없이 전달하기

웹에서 “한참 있다가 한 번에 출력”이 나오면, 모델이 느린 게 아니라 전송 레이어가 flush를 못 하는 경우가 많습니다. SSE는 구현이 단순하지만, 프록시/서버 설정에 따라 버퍼링이 생길 수 있습니다.

import json
import threading
from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

@app.get("/sse")
def sse():
    prompt = "SSE로 토큰을 스트리밍하는 예시를 보여줘."

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        use_cache=True,
    )

    t = threading.Thread(target=model.generate, kwargs=gen_kwargs)
    t.daemon = True
    t.start()

    def event_gen():
        for chunk in streamer:
            data = json.dumps({"delta": chunk}, ensure_ascii=False)
            # SSE 포맷: data: ...\n\n
            yield f"data: {data}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(event_gen(), media_type="text/event-stream")

추가로 서버/프록시에서 아래를 점검하세요.

• Nginx 사용 시 proxy_buffering off;가 없으면 SSE가 뭉쳐서 나갈 수 있습니다.
• Gunicorn/uvicorn 조합에 따라 워커/타임아웃이 스트리밍과 충돌할 수 있습니다.
• 네트워크 타임아웃 설계는 스트리밍에서 특히 중요합니다. 요청이 길어질수록 중간에 끊길 확률이 커지므로 데드라인 전파 관점도 같이 보세요: gRPC 타임아웃 지옥 탈출 - 데드라인 전파 설계

4) 생성 속도 자체를 올려 “끊김처럼 보이는 지연” 줄이기

스트리밍이 매끄럽게 보이려면 토큰 생성 속도가 최소한 인간이 읽는 속도보다 빨라야 합니다. 아래는 체감 개선이 큰 순서입니다.

4.1 torch.inference_mode()와 eval() 확인

학습 모드나 그래디언트 트래킹이 켜져 있으면 쓸데없는 오버헤드가 생깁니다.

model.eval()

with torch.inference_mode():
    out = model.generate(**gen_kwargs)

4.2 use_cache=True로 KV 캐시 사용

대부분 기본값이지만, 파이프라인/래퍼에 따라 꺼지는 경우가 있습니다. KV 캐시가 꺼지면 토큰이 늘어날수록 매 스텝이 더 비싸져서 “점점 느려짐”이 심해집니다.

4.3 dtype과 디바이스 매핑 정리

• GPU가 있다면 torch_dtype=torch.float16 또는 bfloat16.
• device_map="auto"는 편하지만, 일부 레이어가 CPU로 떨어지면 주기적 스톨이 생길 수 있습니다.

CPU 오프로딩이 발생했는지 간단히 확인하려면 모델 파라미터 디바이스를 훑어보세요.

from collections import Counter

devs = Counter(str(p.device) for p in model.parameters())
print(devs)

cpu가 섞여 있다면, VRAM이 부족해 오프로딩 중일 가능성이 큽니다. 이 경우 해결책은 다음 중 하나입니다.

• 더 작은 모델/짧은 컨텍스트
• 4bit/8bit 양자화
• KV 캐시 메모리 최적화

4.4 양자화로 VRAM 스왑 방지 (끊김의 주범)

VRAM이 모자라면 페이지 폴트나 CPU 오프로딩으로 “몇 초 멈춤”이 발생합니다. bitsandbytes 4bit는 체감이 크게 좋아지는 경우가 많습니다.

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
)
model.eval()

주의할 점:

• 4bit는 메모리는 줄지만, 환경에 따라 속도가 항상 빨라지진 않습니다. 하지만 “오프로딩으로 인한 끊김”을 막는 데 매우 효과적입니다.
• 윈도우/드라이버/쿠다 조합에 따라 설치 이슈가 있을 수 있습니다.

4.5 샘플링 파라미터가 과하게 비싼지 점검

top_k, typical_p, repetition_penalty 등은 조합에 따라 연산량이 늘 수 있습니다. “스트리밍 체감”이 목적이면 우선 단순하게 시작하세요.

권장 시작값:

• do_sample=True, temperature=0.7, top_p=0.9
• 너무 긴 max_new_tokens는 피하고, 필요하면 클라이언트에서 이어받기

5) 파이프라인(pipeline)을 쓰면 왜 더 끊기나

transformers.pipeline("text-generation")는 편하지만 내부에서 전처리/후처리와 배치 처리를 하면서 토큰 단위 flush가 늦어지는 경우가 있습니다. 특히 스트리밍이 중요한 서비스라면 model.generate를 직접 호출하는 편이 제어가 쉽습니다.

다만 파이프라인을 꼭 써야 한다면, 스트리머를 명시하고 스레딩 구조를 동일하게 가져가세요.

6) 서버에서 동시 요청이 늘 때: “스트리밍이 서로를 막는” 문제

로컬 LLM 서버는 보통 GPU 1장에 생성 루프 1개가 가장 안정적입니다. 동시 요청을 무리하게 처리하면 다음이 발생합니다.

• GPU 컨텍스트 스위칭으로 토큰 생성이 불규칙해짐
• VRAM 압박으로 오프로딩 발생
• Python 레벨 락 경합으로 스트리밍 전달 지연

실전 처방은 “동시 실행”이 아니라 “동시 접수 + 큐잉 + 공정 스케줄링”입니다.

6.1 요청 큐와 단일 워커(또는 소수 워커)

• GPU당 워커 1개
• 나머지는 대기열
• 스트리밍은 워커가 토큰을 생성하는 즉시 전달

이 구조는 KServe나 서빙 플랫폼에서도 동일한 튜닝 포인트로 이어집니다. 운영에서 503이나 readiness 문제가 섞이면 아래 글도 같이 참고하세요.

• KServe GPU 추론 503, readiness·오토스케일 튜닝

7) “끊김”이 사실 OOM 전조인 경우

간헐적 멈춤 후 프로세스가 재시작되거나, 응답이 중간에 끊기면 OOM이 섞여 있을 수 있습니다. 특히 컨테이너 환경에서 OOMKilled가 나면 스트리밍은 당연히 끊깁니다.

• 컨테이너 로그에서 종료 코드 확인
• dmesg 또는 쿠버네티스 이벤트 확인

쿠버네티스에서 원인 진단 루틴은 아래 글과 유사합니다.

• EKS Pod CrashLoopBackOff? OOMKilled 진단법

8) 디버깅 체크리스트: 어디서 지연이 생기는지 계측하기

체감만으로는 원인을 착각하기 쉽습니다. 최소한 아래 3개 타임스탬프를 찍어보세요.

• 요청 수신 시각
• 첫 토큰이 스트리머에서 나온 시각
• 클라이언트가 첫 chunk를 받은 시각

import time

t0 = time.time()

# 요청 수신
print("recv", time.time() - t0)

# generate는 스레드로 시작

def event_gen():
    first = True
    for chunk in streamer:
        if first:
            print("first_token", time.time() - t0)
            first = False
        yield chunk

• first_token이 늦으면 추론이 느린 것
• first_token은 빠른데 클라이언트가 늦으면 전송/버퍼링 문제

9) 결론: 가장 효과 좋은 처방 순서

1. TextIteratorStreamer + 별도 스레드로 생성과 전달 분리
2. SSE/WebSocket에서 버퍼링 설정 제거, flush 보장
3. eval() + torch.inference_mode() + use_cache=True 확인
4. GPU 오프로딩 여부 확인 후, 필요하면 4bit/8bit 양자화로 VRAM 안정화
5. 동시성은 큐잉으로 풀고 GPU 워커 수를 보수적으로
6. OOM/Crash가 섞이면 먼저 안정성부터 잡기

로컬 LLM 스트리밍의 “끊김”은 대개 모델 자체보다 서빙 구조와 메모리 압박에서 시작합니다. 위 순서대로 적용하면, 같은 모델로도 체감이 확 달라집니다.