메모리 절약

메모리 절약 (Conserving Memory)

큰 모델을 로드하다 보면 어느 순간 컴퓨터가 숨이 차는 듯한 느낌, 즉 메모리 부족(Out Of Memory, OOM) 오류를 만나게 돼요. 특히 로컬에서 거대한 LLM을 돌리려 할 때 아주 흔한 문제죠. 다행히 vLLM에는 이런 부담을 줄여주는 여러 옵션이 마련되어 있어요. 하나씩 짚어볼게요.

출처: vLLM 공식 문서 — conserving_memory

텐서 병렬화 (Tensor Parallelism, TP)

가장 직관적인 방법은 모델을 여러 GPU에 나눠 담는 텐서 병렬화(tensor_parallel_size) 예요. 옵션 이름 그대로, 모델의 텐서를 잘라 여러 GPU에 분산시켜 메모리 부담을 나누는 거죠.

다음 코드는 모델을 2개의 GPU에 걸쳐 나누는 예시입니다.

from vllm import LLM

llm = LLM(model="ibm-granite/granite-3.1-8b-instruct", tensor_parallel_size=2)

⚠️ 주의: vLLM이 CUDA를 올바르게 초기화할 수 있도록, vLLM 초기화 전에 torch.cuda.set_device 같은 관련 함수를 호출하지 마세요. 그렇지 않으면 RuntimeError: Cannot re-initialize CUDA in forked subprocess 같은 오류를 만날 수 있어요. 사용할 디바이스를 제어하고 싶다면 CUDA_VISIBLE_DEVICES 환경 변수를 설정하세요.

📌 참고: 텐서 병렬화를 켜면 각 프로세스가 모델 전체를 읽고 다시 잘라내는 방식이라, 디스크 읽기 시간이 오히려 길어질 수 있어요 (텐서 병렬화 크기에 비례). examples/offline_inference/save_sharded_state.py 스크립트로 체크포인트를 샤딩(분할)된 형태로 미리 변환해 두면, 변환 과정은 시간이 걸리지만 이후 로딩은 훨씬 빨라져요. 그러면 모델 로딩 시간이 텐서 병렬화 크기와 무관하게 일정하게 유지됩니다.

양자화 (Quantization)

양자화된 모델은 정밀도를 약간 희생하는 대신 메모리를 훨씬 적게 차지해요. 정적 양자화(statically quantized)된 모델은 HF Hub에서 내려받아(예: Red Hat AI에서 제공하는 유명 모델들) 별다른 추가 설정 없이 바로 사용할 수 있습니다.

동적 양자화(dynamic quantization)도 quantization 옵션으로 지원되며, 자세한 내용은 양자화 문서에서 확인할 수 있어요.

컨텍스트 길이와 배치 크기 (Context length and batch size)

메모리 사용을 더 줄이는 또 하나의 확실한 방법은, 모델이 다루는 범위 자체를 제한하는 것이에요. max_model_len 옵션으로 컨텍스트 길이를, max_num_seqs 옵션으로 최대 배치 크기를 제한하면 그만큼 KV 캐시 등에 쓰이는 메모리가 줄어듭니다.

from vllm import LLM

llm = LLM(model="adept/fuyu-8b", max_model_len=2048, max_num_seqs=2)

CUDA 그래프 줄이기 (Reduce CUDA Graphs)

vLLM은 기본적으로 추론을 최적화하기 위해 CUDA 그래프를 사용하는데, 이 그래프가 GPU 메모리를 추가로 차지해요. compilation_config를 조정하면 추론 속도와 메모리 사용 사이의 균형을 맞출 수 있습니다.

from vllm import LLM
from vllm.config import CompilationConfig, CompilationMode

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    compilation_config=CompilationConfig(
        mode=CompilationMode.VLLM_COMPILE,
        # 기본적으로 최대 max_num_seqs까지 올라갑니다
        cudagraph_capture_sizes=[1, 2, 4, 8, 16],
    ),
)

그래프 캡처를 아예 끄고 싶다면 enforce_eager 플래그를 쓰면 됩니다.

from vllm import LLM

llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", enforce_eager=True)

캐시 크기 조정 (Adjust cache size)

만약 CPU RAM이 부족하다면 다음 옵션들을 시도해 볼 수 있어요.

(CPU 백엔드 전용) VLLM_CPU_KVCACHE_SPACE 환경 변수로 KV 캐시의 크기를 설정할 수 있어요 (기본값 4 GiB).
(멀티모달 모델 전용) 이미지/비디오 표현을 위한 캐시 크기를 별도로 조절할 수 있어요.

메모리 절약