[LLM] VLLM로 모델 load시, 비정상적으로 GPU 메모리를 사용할 때 해결법

langchain의 LlamaCPP를 사용하다가. 테스트로 vllm을 사용했을 때 평소 13GB정도 먹었던 애가 느닷없이 70GB를 넘게 잡아먹었다.

먼저 VLLM(Very Large Language Models)은 대규모 언어 모델의 효율적인 추론을 위해 설계된 시스템이다. 특히, VLLM은 모델 추론 중에 반복적으로 참조되는 데이터의 캐싱을 위해 키-값 캐시(KV Cache)를 사용한다. KV Cache는 트랜스포머 모델의 각 레이어에서 생성된 키(Key)와 값(Value) 텐서를 저장하여, 동일한 입력에 대해 반복적인 계산을 피하고 성능을 최적화한다.

KV Cache는 다음과 같은 장점을 제공한다:

1. 추론 속도 향상: 반복되는 계산을 줄여 추론 속도를 크게 향상시킨다.
2. 메모리 사용 최적화: 캐싱을 통해 필요한 메모리 용량을 줄이고, 효율적인 메모리 관리를 가능하게 한다.

그러나, KV Cache를 사용하는 VLLM은 GPU 메모리를 많이 소모할 수 있다. 특히, VLLM은 기본적으로 전체 GPU 메모리의 90%를 할당하여 모델 추론과 KV Cache를 처리한다. 이는 VLLM이 최대 성능을 발휘하도록 돕지만, 다른 작업과 병행 시 GPU 메모리 부족을 초래할 수 있다.

예를 들어, 13B LLaMA 모델을 Hugging Face 라이브러리에서 로드할 때 GPU 메모리 사용량은 약 26GB이다. 그러나 동일한 모델을 VLLM에서 로드할 때는 만약 A100( 메모리 약 81GB )사용 시 GPU 메모리 사용량이 약 73GB에 이른다. 이는 VLLM이 전체 GPU 메모리의 90%를 할당하기 때문이다.

 

해결 방법

KV Cache을 처리할 GPU의 사용량을 gpu_memory_utilization 설정으로 줄일 수 있다.

from vllm import VLLMModel

gpu_memory_utilization = 0.75  # 예시로 75%로 설정

# 모델 로드
model = VLLMModel.from_pretrained("facebook/llama-13b", gpu_memory_utilization=gpu_memory_utilization)

input_text = "세종대왕이 초콜릿을 보고 한 말은? '가나다' 이 드립이 재밌어?"
inputs = model.tokenizer(input_text, return_tensors="pt")

with torch.no_grad():
    with torch.cuda.amp.autocast():
        outputs = model(**inputs)