What is vLLM?
최근 성능이 좋은 LLM의 경우 오픈 소스로 배포되는 경우가 많은데, 모델의 훈련된 parameter가 배포되었다면, 사용자나 application을 개발하는 입장에서는 훈련된 parameter를 받아와서 추론 (inference) 서비스를 잘 실행하는 것이 중요하다. 서버 환경에서 여러 사용자 요청이 동시 다발적으로 들어오는 환경에서 추론 서비스를 제공해야 하는 경우 scheduling, batching, memory management 등 시스템적으로 고려해야 할 요소가 많고, 특정 요청만을 GPU에서 돌리는 경우라도 GPU가 여러 개 있으면 어떻게 잘 나눠서 사용할 것인지, 메모리가 부족하면 어떻게 할 것인지, 모델에서 정의된 수학적 연산을 실제 GPU 하드웨어에서 어떻게 잘 맞춤으로 돌릴지 등 고려할 요소가 많다. 이렇듯 시스템 및 하드웨어 관점에서 LLM 서빙을 최적화하기 위한 기법이 많이 연구되어 왔는데, vLLM은 이런 여러 기법들을 적용해 LLM 서빙을 효율적으로 할 수 있도록 만들어놓은 LLM serving framework이고, 코드는 github에 모두 공개되어 있다. 원래는 SOSP'23의 PagedAttention 이라는 유명한 논문에서 처음으로 소개된 KV 캐시의 페이징 관리 기법 및 이를 이용한 attention 연산 최적화 라는 핵심 기술을 구현하기 위한 툴이었는데, 여기에 기존에 연구되어 왔던 여러 최적화 기법들도 붙이면서 이제는 매우 널리 쓰이는 프레임워크가 된 듯하다.
여기서는 vLLM을 파이썬 패키지로 설치하고, vLLM의 API를 호출하는 간단한 예제 코드를 작성해본 뒤, LLM inference serving에 필요한 API들이 어떤 식으로 구현되어 있는지 코드의 계층적 흐름을 이해해보기로 한다. 이를 바탕으로 시스템, 하드웨어 상의 다양한 계층에 구현되어 있는 최적화 기법들이 코드 상으로는 어디에 해당되는지도 이해해보기로 한다.
Installation: Build from the Source
Pypi를 통해 빌드되어 있는 파이썬 패키지를 받아서 vLLM을 설치할 수도 있지만, 여기서는 코드를 이해하고, 필요한 부분의 경우 일부를 수정하며 살펴볼 생각이므로, github에 공개되어 있는 vLLM의 소스 코드를 받아, 코드를 돌리고자 하는 우리 환경에서 직접 빌드부터 하여 패키지를 설치해보기로 한다.
* vLLM의 소스 코드 github repository
https://github.com/vllm-project/vllm
GitHub - vllm-project/vllm: A high-throughput and memory-efficient inference and serving engine for LLMs
A high-throughput and memory-efficient inference and serving engine for LLMs - vllm-project/vllm
github.com
* Conda 환경 준비 및 Python-only build
vLLM example 코드를 돌리기 위한 conda 환경을 만들어주고, 위 github repository에 들어가서 clone을 해서 소스 코드를 받은 다음, 파이썬 패키지로 설치해준다.
conda create -n vllm python=3.10 # Conda 환경 생성
git clone https://github.com/vllm-project/vllm.git
cd vllm
VLLM_USE_PRECOMPILED=1 pip install --editable . # Python 패키지 빌드하여 설치
https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html 이 링크에 들어가보면, 소스 코드로부터 빌드를 하여 vLLM을 파이썬 패키지로 설치할 수 있는 몇 가지 variation들 및 권장 설치 방법이 나와 있는데, 우리의 경우 cpp나 CUDA kernel 코드까지는 일단 수정하지 않기로 하고, Python-only build의 권장 사항을 따라 VLLM_USE_PRECOMPILED=1 이라는 환경 변수를 설정해주고 설치를 진행하였다. 이 경우, cpp나 CUDA kernel 코드들에 대해서는 우리의 환경에서 직접 compile을 하는 것이 아니라, vLLM이 미리 컴파일해둔 binary를 받게 된다. (물론 cpp와 kernel 부터 직접 빌드할 수도 있으며, 빌드에 약간의 시간이 더 걸린다고 나와 있다)
참고로 torch 같은 라이브러리들을 conda 환경에 따로 미리 받지는 않아도 되는데, vLLM 패키지의 설치에 대한 specification이 적혀 있는 setup.py 를 보면 install_requires 를 통해 vLLM이 필요로 하는 파이썬 패키지에 대한 dependency를 알아서 resolve 하도록 해두었기 때문이다. 구체적으로 어떤 식으로 파이썬 패키지들 간의 dependency가 resolve 되는지는 모르지만, pip install -e 커맨드를 통해 setup.py 스크립트가 실행되기 때문에, 일반 python 스크립트와는 달리 Pypi에서 뭔가 특수한 처리를 해주고 있을 것이다. vLLM에 대한 설치가 끝난 뒤
conda list커맨드를 입력해보면, torch 와 같은 라이브러리들이 함께 설치되어 있는 것을 확인할 수 있다.
* CUDA version issue?
참고로 환경변수가 CUDA-12.1로 설정되어 있는 상태에서는 설치 오류가 났는데, https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html 에서 Pre-build wheels 에 대한 설명을 적어둔 걸 보면, 디폴트로 받게 되는 pre-built binary는 CUDA-12.4에서 빌드된 것이라고 한다. 그래서 환경 변수를 CUDA-12.4로 바꿔줬더니 설치가 잘 되었으며, 설명을 마저 읽어보면 CUDA-11.8, 12.1에 대한 pre-built binary도 설치 링크를 제공하고 있다.
Example code: A Simple Text Generation with LLM
LLM 추론에 있어 가장 직관적이고 간단한 예시는 prompt 문장 하나를 넣어주고, LLM에게 output text를 생성하도록 하는 것이다.
import torch
from vllm import LLM
if __name__ == '__main__':
model_path="meta-llama/Meta-Llama-3.1-8B"
prompt = "Hi, my language model assistant. Infer the rule and give me the best answer. \
Strawberry: red, Apple: red, Lemon: yellow, Grape: purple, Watermelon: "
llm = LLM(model=model_path,
dtype="bfloat16",
)
outputs = llm.generate(prompt)
prompt = outputs[0].prompt
generated_text = outputs[0].outputs[0].text
print(f"> Prompt: {prompt}")
print(f"> Generated text: {generated_text}")
class "LLM": entrypoint class to user / application programmer
코드 이해에서 조금 더 자세히, 단계적으로 살펴보겠지만, LLM이라는 클래스가 user 또는 application level의 개발자에게 있어 entrypoint가 된다. 여기서는 Llama-3.1-8B 모델을 돌리기 위해, 모델을 생성할 때 모델의 path를 argument로 넘겨주었고, 그 외에 아주 많은 argument를 넘겨줄 수 있다. 예를 들어, 자료형에 해당하는 dtype부터 vLLM이 제공하는 각종 서비스에 대한 설정 값을 넘겨줄 수 있도록 되어 있다. 물론 vLLM의 개발자들이 알아서 default 값을 잘 정해두었기 때문에, 이렇게 model path 정도만 필수 argument로 넣어줘도 잘 돌아간다
generate() API
이렇게 모델을 선언해줬으면, 모델에 대해 문장을 만들어달라는 API를 호출할 수 있는데, huggingface 및 다른 프레임워크와 마찬가지로 generate() 라는 함수를 API로 제공한다. 여기도 역시 generate 관련된 설정 값을 줄 수 있는 매우 많output argument가 있겠으나, 입력 문장에 해당하는 prompt 문장을 필수로 넣어주면 문장 생성 예제 코드는 작성해서 돌려볼 수 있다.
output을 보면 클래스 안에 또 클래스.. 아무튼 여러 정보를 계층적으로 담고 있는 것을 볼 수 있는데, 그 중에서 위 코드처럼 text라는 멤버 변수의 값을 읽으면 가장 직관적인 출력인 모델이 생성한 문장을 얻을 수 있다.