AI 모델 추론을 위한 VLLM이란 무엇인가? (feat. by Paged Attention)

1. LLM inference의 특징

 

chatGPT 이후로 LLM에 대한 연구가 많이 활성화되었고, 그에 따라 최근에는 LLM 서빙을 다루는 경우가 많아지고 있다.
 
최근 LLM모델이 매우 커짐에 따라 serving비용이 매우 증가하고 있다. LLM이 커짐에 따라 parameter도 많아져 비싸진 token 생성 과정을 생성이 끝날때까지 반복해서 하나의 token밖에 생성하지 못한다. 따라서 이러한 생성 과정은 GPU의 연산 능력을 제대로 활용하지 못하고 throughput을 저하시키는 memory-bound과정이다.
 
Throughput을 향상시키기 위해 많은 request를 하나의 batch로 묶어서 할 수 있지만, 이를 위해서는 GPU memory를 효율적으로 관리해야 한다.

 
위의 그림은 13B의 파라미터를 가진 모델의 inference를 위해 A100(40GB)을 사용할 때, 할당되는 GPU memory의 비율을 나타낸 것이다. Parameter는 inference하는 동안 static하게 남아있는 부분이므로 어떻게 하지 못한다(fp16 등을 사용하더라도 크기는 줄지만 크기 자체는 고정). 하지만, KV Cache부분은 request에 따라 dynamic하게 활용되므로 이 부분을 잘 활용하여야 한다.

 
 

 

2. 기존 KV cache 방법과 문제점

 
 
vLLM을 개발한 팀은 아래의 논문을 통해 PagedAttention을 먼저 소개했고, 직접 개발하여 vLLM 라이브러리를 출시했습니다.
- 논문 제목: Efficient Memory Management for Large Language Model Serving with PagedAttention
- 링크: https://arxiv.org/pdf/2309.06180.pdf
 

2.1. 문제점 1 : internal, external memory fragmentation

 

기존에 사용되던 LLM serving system들은 이 부분이 효율적이지 못하다. 왜냐하면 기존의 방식들은 request에 대한 KV cache를 물리적으로 연속된 하나의 memory 공간에 저장하기 때문이다.
 
생성형 모델에서는 토큰을 생성하면서 매 request마다 할당하고 해제하여 사용하는 메모리가 많이 생기는데, 현재의 플랫폼들은  메모리를 활용하면서 생성 길이에 따른 메모리를 미리 할당 받아놓고 사용하지 않는 경우도 있고, 새로운 메모리를 할당할 때 비워놓는 공간도 생기면서 매우 비효율적으로 메모리를 사용한다.

 

게다가 memory의 공간을 대부분 생성가능한 최대 길이로 할당한다. KV cache는 모델이 생성하는 새로운 토큰에 따라 매우 길어질 수도 있고 매우 짧아질 수도 있기에 예측불가능하여 미리 최대 길이로 할당하는 것이다.

 

 
 
 

예측 불가능하기 때문에 사전에 물리적으로 인접한 최대 길이(e.g., 2048 tokens)의 공간을 할당한다. 여기서 발생하는 fragmentation은 1) reserved, 2) internal fragmentation, 3) external fragmentaion으로 나뉜다.

1. internal fragmentation(2048이 할당되었지만 10만 사용되면 2038은 사용되지 않음: 최대 길이에 도달하기 전에 생성이 완료되어, k,v cache가 이루어지지 않은 공간 (2048이 할당되었지만 10만 사용되면 2038은 사용되지 않음)
2. external fragmentaion: 각 request 마다 사전에 할당된 공간의 크기가 다르기 때문에 생기는 할당 공간 끼리의 비어있는 공간
3. reserved: 실제로 사용되는 공간이다. 하지만 왜 비효율적이라고 하는가? 만약 생성될 길이가 1000이라고 한다면 첫번째 iteration에서는 1개의 공간만 필요하고 999개는 사용되지 않고 있으므로 비효율적임

 

위 그림은 현재 사용되는 시스템에서 GPU memory가 어떻게 사용되고 있는 지를 보여주는데 20.4% - 38.2%밖에 사용하고 있지 않는 것을 알 수 있다.
 
실제 메모리를 프로파일링 해보았을 때, Key-Value 어텐션을 위해서 사용하는 메모리가 30%가 넘는 것으로 확인되었다고 하는데 이 부분을 논문에서는 Key-Value Cache라고 하고 있습니다.

 

 
 

위의 이미지에서 Prefill 단계는 흔히 우리가 생각하는 promt에 대한 부분을 계산하는 것이고, Decode 단계가 생성하는 단계인데 Decode 단계에서는 한 토큰씩만 생성할 수 있는데 이전의 Key-Value Attention을 항상 처음부터 새로 계산하는 부분에 대해서 Caching을 해두겠다는 것이 PagedAttention의 main contribution입니다.

 
 
 

2.2. 문제점 2 : memory를 공유할 수 있는 기회를 활용할 수 없다.

 

decoding 방법 중에는 하나의 request에 대하여 다수의 ouput을 생성하는 방법들(parallel sampling, beam search)이 있고, 실제로 이러한 방법을 많이 이용한다. 이러한 생성 방식 중에는 KV cache를 어느 정도 공유할 수 있는 경우가 많을 텐데(prompt, 일부 생성 token 등), 현재 시스템에서는 KV cache가 분리된 연속된 공간에 존재하기 때문에, KV cache의 공유가 불가능하다.

 
 

3. 해결방안: Paged Attention

 

위와 같은 기존 방식의 두 가지 한계를 해결하기 위해 고안한 attention 방식이다. 기존의 operating system에서 memory fragmentation and sharing: virtual memory with paging 방식에서 영감을 받아서 만든 방식이다.

• KV cache를 여러 Phycial KV block으로 분산
• 각 Phycial KV block은 일정 개수의 토큰의 key, value 저장 (일정 개수는 조절 가능 - 추후 설명)
• Phycial KV block은 물리적으로 인접할 필요가 없음
• Phycial KV block들이 어디에 위치하는지 블록의 순서(주소)와 블록에 몇 개의 값이 저장되어 있는지 기록한 Block Table 존재

위와 같은 방식을 활용하여 필요할 때만 dynamic하게 memory를 할당하여 쓰므로 interenal fragmentaion, reserved의 낭비를 줄일 수 있고, 모든 블록이 같은 크기를 가져서 external fragmentaion도 줄일 수 있다. 또 block단위로 저장되기 때문에 메모리를 공유하는 것도 가능하다.

 
 

 
그래서 vLLM은

• PagedAttention 기반의 high-throughput distributed LLM 서빙 엔진
• PagedAttention을 통해 KV cache 메모리의 낭비가 0에 가까움
• PagedAttentions과 같이 디자인된, block-level memory management와 preemptive request scheduling을 사용
• GPT, OPT, LLaMA와 같은 널리 사용되는 LLM을 지원
• 하나의 GPU에 담을 수 없는 크기도 지원
• 정확도의 손해 없이 sota system 보다 2-4X의 serving throughput 달성
• 특히 모델이 클수록, seq length가 길수록, decoding방법이 복잡할 수록 좋은 성능 향상

 
 

그래서 PagedAttention은 위의 이미지와 같이 메모리를 블록으로 나누어 놓고 매핑 테이블을 관리하는 방식으로 설계되어 낭비되는 메모리도 줄이고, 캐싱을 통해 다시 계산하지 않아도 되는 효과를 얻을 수 있었다고 한다. 캐싱은 GPU 메모리에서도 활용하지만, CPU 메모리도 같이 활용하여 캐싱을 많이 할 수 있도록 실제로 개발되었다고 한다.
 
그래서 저자들은 PagedAttention을 통해서 메모리를 효율적으로 사용하기 때문에, batch 사이즈를 더 크게 사용할 수 있고 최종적으로 같은 환경에서 throughput을 많이 올릴 수 있었다고 한다.

 

 

 

특히, 위의 이미지에서 prompt 길이를 1-shot으로 가져가는 것 보다 5-shot으로 길게 가져가게 되면, 캐싱을 많이 활용하기 때문에 훨씬 더 큰 효과를 발휘한다고 한다.

 

PagedAttention은 메모리 사용을 효율적으로 해주기 때문에, batch를 키워서 throughput은 높였지만 실제 단일 inference 속도가 올라간 것은 아니다.

 

 

 

4. vLLM 소개

 

vLLM은 PagedAttention의 저자들이 직접 구현하여 발표한 라이브러리이다.

 

vLLM은 PagedAttention 기법을 활용하여 문장 생성 속도를 비약적으로 높인 방법론이다. 이뿐만 아니라 실제 서빙을 위해 많은 요소들이 포함되어 있다. 예를 들어 멀티 클러스터 환경에서 안정적인 서빙을 하기 위한 Ray Cluster를 사용하거나 큰 모델과 데이터를 병렬로 처리할 수 있도록 Megatron LM의 Parallelism을 차용하고 있다.

 

 

 

위 다이아그램은 vLLM의 전체적인 구조이다.

 

LLMEngine에서 분산 처리를 위한 워커, PagedAttention의 블록을 관리하는 블록 매니저, KV 캐시 등을 관리하는 컴포넌트들을 생성하고, 매 요청 시 스케줄러를 통하여 요청된 프롬프트들의 생성 순서를 바꿔준다.

 

LM에서 문장을 생성하기 위해서는 마지막 토큰이 나올 때까지 모델에 반복적으로 포워딩 해야 하는데, 스케줄러에 의해 메모리 및 우선순위를 따져 효율적인 방법으로 GPU 유틸리티를 달성하고 있다. 메모리의 효율성을 더욱 높이기 위해 PagedAttention 기법을 사용하며, GPU 메모리가 부족할 경우 CPU 메모리에 스왑하는 방식으로 중간 계산 과정들에 대해서도 안정적으로 관리하고 있다.

 

 

 

 

HF는 HuggingFace, TGI는 HuggingFace Text Generation Inference인데, vLLM을 활용하면 훨씬 더 높은 throughput을 얻을 수 있었다고 한다.
 
 
 

출처: http://arxiv.org/pdf/2309.06180, 스캐터랩, 석현2 블로그

 

뜨리스땅

 

 

 

 

 

https://tristanchoi.tistory.com/655

NVIDIA Triton 란 무엇인가?