Published on

Ray Serve로 FastAPI LLM 서빙 GPU OOM 잡기

Authors

FastAPI로 LLM 추론 API를 만들면 처음엔 잘 돌아가다가, 트래픽이 조금만 늘어도 갑자기 GPU OOM이 터지는 경우가 많습니다. 특히 다음 조건이 겹치면 재현이 쉽습니다.

  • 단일 프로세스에서 모델 1개를 올리고, 요청마다 generate를 바로 호출
  • FastAPI 동시성(uvicorn worker, async)과 모델 추론 동시성이 섞여 제어 불가
  • 입력 길이(프롬프트)와 출력 길이(max tokens)가 요청마다 크게 변동
  • 스트리밍 응답, 대기열 증가로 한 번에 여러 요청이 GPU에 쌓임

Ray Serve는 “모델 실행을 담당하는 워커(Actor)와 HTTP 인입을 분리”하고, 워커 단위로 GPU를 고정 할당하며, 동시성/배치/오토스케일을 통제할 수 있어 OOM을 구조적으로 줄이기에 좋습니다. 이 글은 FastAPI LLM 서빙에서 흔한 OOM 패턴을 짚고, Ray Serve로 바꾸면서 어떤 설정을 어디에 걸어야 하는지 실전 관점으로 정리합니다.

관련해서 OOM과 지연(P95) 튜닝을 더 깊게 보고 싶다면 아래 글도 같이 참고하면 좋습니다.


FastAPI 단독 서빙에서 OOM이 나는 대표 원인

1) “동시 요청 수”가 곧 “동시 GPU 실행 수”가 되는 구조

FastAPI는 HTTP 레벨에서 동시성을 잘 처리합니다. 문제는 모델 추론이 GPU 메모리를 많이 쓰는 작업인데, 이를 별도의 큐/세마포어 없이 병렬로 실행하면 GPU에 다음이 동시에 올라갑니다.

  • KV cache(디코딩 단계에서 누적)
  • attention 중간 텐서
  • 샘플링/로짓 텐서

요청이 2배가 되면 메모리도 거의 선형으로 늘어나는 경우가 많아, 어느 순간 임계점을 넘으면 OOM이 납니다.

2) 입력/출력 길이 변동이 메모리 피크를 만든다

LLM 추론에서 메모리 사용량은 대략 다음 변수에 민감합니다.

  • 배치 크기(동시 요청 수)
  • 프롬프트 토큰 수
  • 생성 토큰 수
  • 모델 크기, dtype(fp16/bf16)

FastAPI 단독 구조에선 “긴 요청이 섞였을 때” 피크가 훨씬 커지고, 그 피크가 곧 OOM으로 연결됩니다.

3) 파편화(fragmentation)와 캐시 정책

PyTorch는 CUDA allocator를 쓰며, 반복적인 할당/해제를 겪으면 파편화가 누적될 수 있습니다. 특히 스트리밍 응답이나 타임아웃/취소가 잦으면 메모리 패턴이 더 불규칙해집니다.


해결 전략: Ray Serve로 “GPU 메모리의 상한”을 만들기

Ray Serve로 바꾸면 다음 3가지를 구조적으로 분리할 수 있습니다.

  1. HTTP 인입(라우팅/인증/리밋)과 모델 실행을 분리
  2. 모델 실행 워커(Replica) 단위로 GPU를 고정 할당
  3. 워커 내부 동시성, 배치, 큐 길이를 제한

즉, “트래픽이 늘어도 GPU에서 동시에 실행되는 추론 수를 제한”하고, 나머지는 큐에 쌓이게 만들어 OOM 대신 지연 증가로 바꾸는 것이 핵심입니다.


아키텍처: FastAPI는 얇게, Ray Serve가 추론을 책임진다

권장 패턴은 다음과 같습니다.

  • FastAPI(또는 Ray Serve의 HTTP)에서 요청 검증/스키마/레이트리밋
  • Ray Serve Deployment가 모델을 로드하고 추론 수행
  • max_concurrent_queries로 워커당 동시 실행 제한
  • @serve.batch로 마이크로배칭 적용(가능한 경우)

FastAPI를 반드시 유지해야 한다면 “FastAPI app을 Ray Serve ingress로 감싸는 방식”이 가장 깔끔합니다.


예제: Ray Serve + FastAPI(ingress) + GPU 워커로 OOM 방지

아래 코드는 다음을 포함합니다.

  • Ray Serve Deployment가 GPU 1장을 점유
  • 워커당 동시 요청 수를 제한
  • 입력 길이/출력 길이 상한을 강제
  • 배치(옵션)로 처리량 개선
# serve_app.py

from typing import List, Optional

import ray
from ray import serve
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

app = FastAPI()

class GenerateRequest(BaseModel):
    prompt: str = Field(..., min_length=1, max_length=20000)
    max_new_tokens: int = Field(256, ge=1, le=1024)
    temperature: float = Field(0.7, ge=0.0, le=2.0)
    top_p: float = Field(0.9, ge=0.0, le=1.0)

class GenerateResponse(BaseModel):
    text: str

@serve.deployment(
    ray_actor_options={"num_gpus": 1},
    # 워커(Replica)당 동시에 처리할 쿼리 수를 제한
    max_concurrent_queries=2,
)
@serve.ingress(app)
class LLMService:
    def __init__(self, model_id: str = "meta-llama/Llama-2-7b-chat-hf"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="cuda",
        )
        self.model.eval()

        # 선택: 추론 중 그래프/캐시로 인한 메모리 증가를 줄이기 위한 설정
        torch.set_grad_enabled(False)

    @app.post("/v1/generate", response_model=GenerateResponse)
    async def generate(self, req: GenerateRequest):
        # 토큰 길이 상한(프롬프트가 너무 길면 OOM 위험)
        inputs = self.tokenizer(
            req.prompt,
            return_tensors="pt",
            truncation=True,
            max_length=4096,
        ).to("cuda")

        try:
            with torch.inference_mode():
                out = self.model.generate(
                    **inputs,
                    max_new_tokens=req.max_new_tokens,
                    do_sample=req.temperature > 0,
                    temperature=req.temperature,
                    top_p=req.top_p,
                    use_cache=True,
                )

            text = self.tokenizer.decode(out[0], skip_special_tokens=True)
            return GenerateResponse(text=text)

        except torch.cuda.OutOfMemoryError as e:
            # OOM을 500으로 터뜨리기보다, 429/503으로 변환해 재시도/백오프로 유도
            raise HTTPException(status_code=503, detail="GPU OOM: try later") from e

        finally:
            # 요청 단위로 캐시를 강제로 비우는 것은 지연을 늘릴 수 있어 신중히 사용
            # 파편화가 심각한 환경에서만 제한적으로 고려
            pass

llm_app = LLMService.bind()

if __name__ == "__main__":
    ray.init()
    serve.run(llm_app)

핵심은 max_concurrent_queries입니다. 이 값이 사실상 “GPU에서 동시에 돌아갈 수 있는 추론 수의 상한”이 됩니다. OOM이 났던 환경이라면 우선 1 또는 2로 낮게 시작하고, 안정화 후 올리는 접근이 안전합니다.


마이크로배칭으로 OOM을 줄이는 이유(그리고 주의점)

배칭은 처리량을 올리는 방법으로 알려져 있지만, “무제한 동시성” 대신 “제어된 배치”로 바꾸면 OOM을 줄이는 데도 도움이 됩니다.

  • 나쁜 동시성: 요청 8개가 제각각 generate를 시작해 KV cache가 8개 쌓임
  • 좋은 배치: 짧은 시간창에 들어온 요청을 묶어 1번의 forward로 처리(가능한 범위)

Ray Serve는 @serve.batch 데코레이터로 배칭을 지원합니다. 다만 생성 모델의 배칭은 구현 난이도가 있고, 단순히 generate를 묶는다고 효율이 항상 좋아지는 것은 아닙니다(길이 패딩, 디코딩 루프 등). 그래도 “동시 실행 수를 제어하면서” 큐잉을 만들고 싶다면 유용합니다.

배칭 예시는 아래처럼 형태를 잡을 수 있습니다.

from ray import serve

@serve.deployment(ray_actor_options={"num_gpus": 1}, max_concurrent_queries=8)
class BatchedLLM:
    def __init__(self):
        ...

    @serve.batch(max_batch_size=4, batch_wait_timeout_s=0.01)
    async def _batched_generate(self, prompts: list[str], max_new_tokens: list[int]):
        # 실제로는 길이 정렬, 패딩, attention mask 처리 등이 필요
        # 생성 모델 배칭은 프레임워크(vLLM 등)를 쓰는 편이 더 안정적일 수 있음
        ...

    async def __call__(self, request):
        data = await request.json()
        return await self._batched_generate(data["prompt"], data["max_new_tokens"])

배칭은 “OOM을 없애는 만능키”가 아니라, 동시성 제어와 함께 쓰는 최적화 옵션으로 보는 게 맞습니다.


OOM을 막는 운영 파라미터 체크리스트

1) 워커(Replica)당 GPU를 고정하고, 동시성을 제한한다

  • ray_actor_options={"num_gpus": 1}로 GPU 점유
  • max_concurrent_queries로 워커당 동시 요청 수 상한
  • 트래픽이 늘면 Replica를 늘려 수평 확장

2) 입력/출력 길이 상한을 API 계약으로 못 박는다

  • max_length(프롬프트 토큰) 상한
  • max_new_tokens 상한
  • 너무 긴 요청은 400으로 거절

이게 없으면 “가끔 들어오는 초장문 요청 1개”가 전체를 OOM으로 무너뜨립니다.

3) 큐 길이와 타임아웃을 설계한다

동시성을 제한하면 큐가 생깁니다. 큐가 무한정 늘어나면 결국 지연이 폭발하고, 클라이언트 재시도 폭풍이 2차 장애를 만듭니다.

  • 서버 측 타임아웃(예: 30초)
  • 클라이언트 재시도 정책(지수 백오프, 지터)

레이트리밋/재시도 설계는 LLM API에서 특히 중요합니다. 429 대응 패턴은 아래 글이 참고됩니다.

4) 파편화가 의심되면 allocator 설정을 점검한다

환경변수로 CUDA allocator 동작을 바꿔 파편화를 줄이는 경우가 있습니다.

# 예시: 파편화 완화에 도움이 되는 경우가 있음
export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,garbage_collection_threshold:0.8"

이 설정은 워크로드에 따라 득실이 갈립니다. 반드시 부하 테스트로 확인하세요.


Ray Serve에서 자주 겪는 함정 4가지

1) FastAPI async라고 해서 GPU 호출이 자동으로 안전해지지 않는다

async def는 I/O에 유리할 뿐, GPU 메모리 상한을 만들어주지 않습니다. 동시 요청이 들어오면 결국 동시에 generate가 실행될 수 있습니다. Ray Serve의 max_concurrent_queries 같은 “명시적 상한”이 필요합니다.

2) Replica 수를 늘리면 OOM이 줄어들 수도, 늘어날 수도 있다

  • 장점: 요청이 분산되어 워커당 동시 실행이 줄어 안정적
  • 단점: 각 Replica가 모델을 하나씩 들고 있어 GPU 메모리가 복제됨

즉, “하나의 GPU에 Replica를 여러 개” 띄우는 건 대개 위험합니다. 보통은 GPU 1장당 Replica 1개가 기본값입니다.

3) 스트리밍은 메모리보다 “동시 연결”과 “취소 처리”가 더 문제다

스트리밍 자체가 메모리를 크게 늘리진 않더라도, 연결이 오래 유지되면서 워커 점유 시간이 늘고 큐가 길어질 수 있습니다. 이때 취소된 요청의 정리(클라이언트 disconnect)가 제대로 안 되면 메모리/세션이 누수처럼 보일 수 있습니다.

4) OOM이 나면 프로세스가 불안정해질 수 있다

PyTorch OOM은 예외 처리로 복구 가능한 경우도 있지만, 반복되면 성능이 급격히 나빠지거나 파편화가 심해질 수 있습니다. 운영에선 “OOM이 나지 않게” 만드는 게 1순위고, 그 다음이 “OOM 시 빠르게 격리/재시작”입니다.

Kubernetes에서 롤링/카나리로 안전하게 바꾸는 전략이 필요하다면 아래 글도 도움이 됩니다.


실전 권장값(시작점)

모델/GPU마다 다르지만, OOM을 줄이기 위한 시작점은 다음처럼 잡는 편이 안전합니다.

  • GPU 1장당 Replica 1개
  • max_concurrent_queries=1 또는 2
  • 프롬프트 토큰 상한: 2048~4096
  • max_new_tokens 상한: 256~512
  • 과부하 시 응답: 503(서버 혼잡) + 클라이언트 백오프

이 상태에서 부하 테스트로 P95를 확인하고, 처리량이 부족하면 다음 순서로 조정합니다.

  1. max_concurrent_queries를 아주 조금씩 증가
  2. 배칭(가능하면) 도입
  3. Replica 수를 늘려 수평 확장(추가 GPU 필요)

결론: OOM을 “코드 트릭”이 아니라 “상한 설계”로 없애기

FastAPI 단독 서빙에서의 GPU OOM은 대부분 “통제되지 않는 동시 추론”에서 시작합니다. Ray Serve로 옮기면 GPU를 점유하는 워커를 분리하고, 워커당 동시성 상한을 강제하며, 큐잉과 스케일링을 운영 레벨에서 다룰 수 있어 OOM을 구조적으로 줄일 수 있습니다.

정리하면 다음 3가지만 먼저 적용해도 효과가 큽니다.

  • ray_actor_options={"num_gpus": 1}로 GPU 격리
  • max_concurrent_queries로 동시 추론 상한
  • 입력/출력 길이 상한을 API 계약으로 고정

그 다음 단계로 배칭과 오토스케일, 그리고 429/503 기반의 재시도 설계를 더하면 “OOM 없이 버티는” LLM API에 가까워집니다.