관심 있는 NLP 논문을 읽어보고 간단히 정리했습니다.
혹시 부족하거나 잘못된 내용이 있다면 댓글 부탁드립니다 🙇♂️
[UC Berkeley, Stanford University]
- 운영체제에서 가상메모리와 페이징 기법에 착안한 PagedAttention을 제안
- 이를 기반으로 하는 vLLM을 개발했는데, (1) KV 캐시 메모리의 낭비가 거의 없고 (2) 불필요한 메모리 사용을 최소화 한다는 특징이 있음
출처 : https://arxiv.org/abs/2309.06180
깃허브 : https://github.com/vllm-project/vllm
1. Introduction
최근 가장 감명 깊게 읽은 PagedAttention에 대해 간단히 리뷰를 해보려고 합니다.
저자들은 이를 바탕으로 vLLM이라는 라이브러리를 만들어 오픈소스로 공개했는데요, 실제로 사용해보면 추론속도가 말도 안되게 빠릅니다.
원래 생성형 모델들의 고질적인 문제가 추론인데, 기존 방식들 혹은 다른 라이브러리 대비 몇 배가 빠른지 가늠이 되지 않는 정도였습니다..
왜 추론(generate)이 오래 걸리느냐, 하는 것에 대해서는 이미 잘 알려져 있습니다.
decoder 기반의 언어모델들은 특정 time step $t$에서 $t-1$까지의 토큰들을 조건으로, 등장할 확률이 가장 높은 토큰 하나를 예측하게 됩니다.
이 과정을 '문장이 끝났음을 알리는 토큰'이 등장할 때까지 반복하게 되니 그 시간이 엄청나게 많이 걸리는 것이죠.
특히나 요즘 모델들을 보면 답변의 길이가 엄청나게 길고..
결국은 답변을 구성하는 토큰의 개수만큼 forward 하는 과정이 반복되고 있다는 뜻입니다.
그런데 생각해보면 우리는 추론을 시작할 때 모델의 답변의 길이가 얼마나 될지 알 수가 없습니다.
그래서 만일을 대비하여 최대한 많은 메모리 공간을 생성 과정에 할당해야 하죠.
예를 들어 최대 출력 길이가 4096 이라고 한다면 메모리를 이에 맞게 할당하는데, 문제는 실제 답변은 100 길이로 생성될 수도 있다는 것입니다.
결국 남는 메모리 공간은 모델의 생성 결과가 등장할 때까지 이러지도 저러지도 못하고 붙들려 있는 상황이 되고, 이것이 곧 추론의 병목이 됩니다.
추론할 때마다 모든 토큰 간 attention 연산을 수행하지 않고, 한 번 연산된 결과는 KV cache에 저장되어 다음 step에도 활용할 수 있도록 하는 테크닉도 널리 사용되고 있습니다만 결국 최종 출력 길이를 임의로 산정하여 불필요한 메모리 공간을 할당해야 한다는 사실에는 변함이 없죠.
본 논문에서는 이를 해결하는 방법론으로 PagedAttention을 제안하며 이를 기반으로 삼는 vLLM 라이브러리를 공개했는데, 그 성능이 압도적으로 뛰어나다고 말씀드릴 수 있겠습니다.
저자는 친절하게도 Autoregressive Generation 방식에 대해서 설명을 해주었는데 여기서 다루기엔 포스팅이 너무 길어질 것 같아서 패스하도록 하겠습니다.
언어 모델의 생성 과정에 대해 좀 더 깊은 이해가 필요하신 분들은 해당 논문 Chapter 2. Background를 참고해보세요.
2. Memory Challenges in LLM Serving
2.1. KV cache
사실 이것도 배경에 해당하는 내용이지만 이를 제외하고 글을 이어나갈 수는 없을 것 같아서 ㅎㅎ..
위에서 간단히 언급한 것처럼 언어 모델은 매 순간마다 한 개의 토큰을 새로 생성하는 과정을 반복합니다.
이때 기존 입력 ($< t$)은 attention 메커니즘의 Key와 Value 역할이고, 새로 생성된 토큰은 Query가 됩니다.
이를 바탕으로 attention 연산을 수행해 새로운 토큰을 추가로 생성하면, 그 토큰이 새로운 Query가 되고 이전의 결과물들은 Key, Value가 됩니다.
이러한 과정을 반복하기 때문에 매번 새로 attention을 하는 것이 아니라, 연산 결과들을 중간에 다 저장해서 다음 step에 활용할 수 있게 되는 것이죠.
그러다가 <eos> 토큰이 생성되는 순간 생성 과정을 마치고 그때까지의 결과를 반환하게 됩니다.
(스타트할 때 의문이 들 수도 있는데.. 정말 순수하게 <bos> 토큰만 존재한다면 self-attention으로 시작하게 되고, 일반적으로는 prompt phase를 거치게 됩니다.)
2.2. Large KV cache
KV cache의 크기는 request의 숫자에 비례하여 빠르게 증가합니다.
논문에서는 13B 사이즈의 OPT 모델을 예시로, 한 개의 토큰을 위한 KV cache 공간이 800KB라고 언급합니다.
이는 key와 value 각각 필요하므로 2를 곱해야 하고, 모델의 hidden state size인 5120에 layer의 개수인 40을 곱한 뒤 FP16 자료형은 2 byte씩 필요하다는 것까지 곱해주기 때문입니다.
즉, 2 (key & value) x 5120 (hidden state size) x 40 (number of layers) x 2 (bytes per FP16) = 819200 bytes = 800KB 가 됩니다.
만약 모델이 최대 2048 토큰까지 생성이 가능하다면, 2048 x 800KB = 1638400KB = 1600MB = 1.6GB 가 됩니다.
한 개의 request가 1.6GB 의 메모리를 차지하게 된다는 뜻이죠..
가용한 메모리를 전부 KV cache에 할당하더라도 몇 개의 request밖에 처리할 수 없는 상황입니다.
게다가 GPU의 발전 과정을 보면, 플래그십에 해당하는 A100과 H100의 vRAM 크기가 동일하게 80GB라는 특징이 있습니다.
FLOPS는 두 배 이상 늘어난 반면 메모리는 그대로 유지되는 상황이고, 이는 곧 메모리의 크기가 심각한 병목이 될 수 밖에 없다는 것을 방증합니다.
2.3. Complex decoding algorithms
LLM 서비스 관점에서 보면, 유저가 어떤 질문을 여러 개 날리는 과정에서 중복되는 텍스트들이 있습니다.
이는 서비스 공급자 설정한 시스템 프롬프트일수도 있고, 상황에 따라 유저가 정한 few-shot 등이 될 수 있겠죠.
그래서 반복적으로 등장하는 Key-Value에 대해 불필요한 연산 과정을 생략할 수 있도록 프롬프트 KV를 재활용하는데, 이는 전체 KV cache의 12% 정도를 차지한다고 합니다.
하지만 autoregressive 생성 과정의 대부분을 차지하는 KV cache는 한 개의 샘플만을 위해 존재하고 버려지게 됩니다.
결국 이것 역시도 불필요한 메모리 점유 상황을 야기하여 비효율적인 GPU 가동률로 이어지게 됩니다.
2.4. Scheduling for unknown input & output lengths
인트로에서도 언급했지만, 가장 문제가 되는 것은 입출력 길이가 얼마나 될지 알 수가 없다는 것입니다.
시퀀스의 최대 길이를 산정해두면 이 길이를 넘어서는 생성하지 않게 되지만, 이에 한참 미치지 못하는 생성 결과가 나타날 수도 있죠.
(위에서 최대 길이가 4K인데 실제로는 100 토큰 생성 후 종료되는 상황을 예로 들었습니다.)
하지만 메모리를 미리 할당해두지 않으면 생성이 제대로 될 수 없을 것이고..
여러 request가 한 번에 들어오는 상황에서는 이러한 상황이 더 심각해질 것입니다.
결국 이는 여러 개 샘플을 한 번에 추론하지 못하도록 만드는 주요한 원인 중 하나가 됩니다.
2.5. Memory Management in Existing Systems
기존의 메모리 관리 형상을 보면 위와 같습니다.
현존하는 딥러닝 프레임워크는 텐서가 연속적인 메모리 공간에 저장되어야 하게끔 설계가 되어 있습니다.
위 이미지를 보시면 크게 세 종류의 메모리 낭비 원인을 알 수 있습니다.
- 미래에 생성될 토큰을 위해 예약된 슬롯, reserved slots
- 최대 길이만큼 공간을 마련해두었으나 사용되지 않고 버려진, internal fragmentation
- memory allocator로부터 발생하는, external fragmentation
reserved slots와 internal fragmentation은 사실 한 덩어리로 봐도 되겠지만..
어쨌든 internal fragmentation은 이후에 생성 토큰으로 차지가 된다고 하더라도, 생성이 완료될 때까지 붙들고 있어야 하는 문제가 해결되는 건 아닙니다.
3. Method
3.1. Paged Attention
논문에서는 이 방식을 운영 체제의 paging에 착안한 것이라고 설명하고 있지만..
사실 저는 이에 대해 잘 알지 못해 패스하겠습니다.
어쨌든 핵심은 PagedAttention은 KV cache를 연속되지 않은 메모리 공간에 저장할 수 있다는 것입니다.
이때 저장되는 단위는 block이고, 각 block은 일정한(고정된) 숫자의 토큰을 포함하게 됩니다.
$$A_{ij} = \frac{\exp(q_i^\top K_j / \sqrt{d})}{\sum_{t=1}^{\lceil i/B \rceil} \exp(q_i^\top K_t / \sqrt{d})}, \quad o_i = \sum_{j=1}^{\lceil i/B \rceil} V_j A_{ij}^\top$$
$$K_j = (k_{{(j-1)B+1}}, \ldots, k_{jB}), \quad V_j = (v_{{(j-1)B+1}}, \ldots, v_{jB}), \quad A_{ij} = (a_{i,{(j-1)B+1}}, \ldots, a_{i,jB})$$
우선 $K_j, V_j, A_{ij},\,$$o_i$ 각각은 블록 단위의 벡터, attention score, output을 나타내고 있습니다.
이때 $K, V, A$의 범위가 위와 같이 정해지는 이유는 j 번째 블록이기 때문입니다.
예를 들어 한 블록에 4개의 숫자를 집어 넣는데 전체 블록의 개수는 5개 (총 20개의 값이었겠죠), 현재는 2번째 블록을 보고 있다고 하겠습니다.
즉, $B=4, j=2$ 입니다.
1부터 20개의 값 중에서 두 번째 블록에 해당하는 값들은 [5, 6, 7, 8] 이 될 것입니다.
이는 $(j-1) * B + 1 = (2-1) * 4 + 1 = 5$에서 출발하여 $j * B = 2 * 4 = 8$에서 끝나는 것으로 표현 가능하다는 걸 알 수 있습니다.
Attention score를 계산하는 것은 기존의 self-attention score를 구하는 것과 동일한데, 이를 블록 단위로 나누었기 때문에 분모의 변수가 $t$라는 점만 다릅니다.
참고로 $[i/B]$의 대괄호는 천장 함수(ceiling function)로 $i$를 $B$로 나눈 값을 올림한 것입니다.
전체 블록의 개수를 의미하게 되고, $t$는 블록의 번호를 의미하는 것으로 볼 수 있습니다.
마지막 순서의 블록은 이전의 것들과 개수가 달라질 수 있습니다. (ex. 4 + 4 + 4 + 1)
위 이미지를 보시면, 'forth'에 대한 Query vector를 Block 0, 1, 2에 산재되어 있는 Key & Value vectors와 attention하여 output을 구하는 모습이 시각화되어 있다는 걸 알 수 있습니다.
3.2. KV Cache Manager
vLLM의 memory manager의 핵심 아이디어는 운영 체제의 virtual memory와 유사하다는데.. 이것도 모르겠으니 넘어가고요 ㅎㅎ..
위에서 블록 단위로 쪼개 놓은 KV cache들을 논리적으로 물리적으로는 비연속적으로 관리할 수 있고(non-contiguous physical memory page), 이를 논리적인 관계로 연결해서 관리하는 것입니다.
KV cache request가 들어오면 블록 단위에 맞게끔 메모리를 할당하고요, 이전과 동일하게 future token을 위한 공간도 미리 예약을 해둡니다.
여기서 남는 공간은 필연적으로 생기는 것이지만 단위가 블록이라서 최소화 될 수 있겠죠.
그리고 block tables 개념을 사용하는데, 이는 각 request의 logical & physical block을 mapping 하는 역할을 수행합니다.
이걸 글로만 보면 잘 이해가 되지 않기 때문에 아래에서 그림과 함께 설명드리도록 하겠습니다.
3.3. Decoding with PagedAttention and vLLM
① vLLM은 prompt를 처리하는데 필요한 블록만을 초기에 할당합니다.
생성 가능한 최대 시퀀스 길이만큼 초기화할 필요가 없다는 것이죠.
위 예시에서 block size는 4이고 프롬프트는 'Four score and seven / years ago our'로 7개의 토큰으로 구성됩니다.
따라서 4 / 3 개로 나뉘어 두 개의 블록이 필요하고, 남은 한 칸은 미래에 생성될 토큰을 위해 reserved 됩니다.
두 개의 블록은 왼쪽 Logical KV blocks에서 보이는 것처럼 Block 0, 1에 할당됩니다.
즉, 논리적으로는 0번 블록이 1번 블록으로 이어진다는 뜻입니다.
하지만 실제로는 오른쪽 Physical KV blocks에서 각각 Block 7, 1에 할당되는 것을 확인할 수 있습니다.
이 정보는 중간의 Block Table에 기록이 되므로, 우리는 어떤 Block이 다른 Block으로 어떻게 이어지는지를 Block Table을 통해 알 수 있게 됩니다.
(0 → 7, 1 → 1)
예를 들어 Logical KV block의 0번이 'seven' 이라는 토큰으로 끝난 뒤 다음에는 1번 block으로 이어져야 하는데, 사실 이 정보는 GPU의 1번 block에 저장되어 있다는 것이죠.
물론 0번 block 자체도 GPU의 7번 block으로부터 정보를 꺼내와야 하긴 하지만요.
② 위 설명처럼 physical block 7 & 1 번으로 PagedAttention 알고리즘을 수행하여 새로운 토큰을 생성합니다.
새로 생성된 토큰은 Logical Block 1번의 마지막 자리 & Physical Block 1번의 마지막 자리를 차지하게 됩니다.
그리고 Block Table에서 # filled 란을 3 → 4로 표시하죠.
이는 해당 블록에서 몇 개의 토큰 자리가 채워졌는지를 표시하는 것입니다.
그래서 처음에는 한 칸이 비워진 3이었던 것이고, 생성된 토큰이 남은 한 자리를 채우면서 4로 변하게 된 것입니다.
③ 두 번째 decoding step에서, 기존의 마지막 logical block의 빈자리가 없기 때문에 새로운 physical block을 할당하게 됩니다.
예시에서는 Physical 3번 block이 할당되었습니다.
Block Table을 보면 Logical 2번 block에 Physical 3번 block에 매핑되어 있는 것이 확인됩니다.
이와 같은 과정을 생성이 중단될 때까지 반복하면 되겠습니다.
그런데 대부분은 이런 decoding 과정을 배치 단위로 실행할 것인데, 두 개의 샘플에 대해 동시에 추론하는 예시를 도식화한 것은 아래와 같습니다.
이때의 주요한 특징은 두 가지인 것 같습니다.
1) 불필요한 공간을 다 할당받지 않는다.
physical block은 필요할 때마다 새로 할당받고 그 관계는 mapping table에 기록하기 때문에 당장 쓰이지도 않을 메모리를 할당하지 않아도 됩니다.
2) 한 개 샘플의 생성이 끝나면 해당 메모리는 해제한다.
어차피 블록 단위로 매 step 추론이 이뤄지고 있기 때문에 한 배치 내 모든 시퀀스에 대한 생성이 종료될 때까지 기다릴 필요가 없습니다.
그냥 메모리를 해제하고 묶여 있던 메모리들을 다른 작업에 사용할 수 있게 되는 것이죠.
결국 이러한 특징들 덕분에 GPU를 최대한 효율적으로 꽉꽉 채워서 활용할 수 있게 됩니다.
3.4. Application to Other Decoding Scenarios
3.4.1. Parallel sampling & Shared prefix
parallel sampling에서는 한 개의 request가 여러 개의 sample을 포함하게 되고, 이때 입력 프롬프트는 공통적으로 사용되므로 공유 대상이 됩니다.
위 그림에서는 입력 프롬프트 중 4개의 토큰을 Physical 7번 블록에 한 번만 저장하여 샘플 A1, A2에서 둘 다 참조하고 있다는 걸 알 수 있습니다.
그리고 입력 프롬프트 중 나머지 3개의 토큰은 1번 블록에 저장됩니다.
여기에서의 핵심은 한 개의 physical block이 여러 개의 샘플에 logically 연결될 수 있다는 것이죠.
이를 관리하는 방법이 reference count 입니다.
참조해야 할 횟수가 몇 번 남았는지를 표시한 것으로 이해할 수 있습니다.
입력 프롬프트가 저장된 7 & 1 번 블록 둘 다 reference count가 2라는 뜻입니다.
논문의 설명을 그대로 따라가보면 이렇습니다.
[예시]
샘플 A1이 생성을 할 때 Block 1의 마지막 칸 자리에 생성을 시도할 것입니다.
그런데 reference count가 2이므로 이 숫자를 2 → 1 로 변경하고, Block 3에 Copy-on-write 한 뒤 생성된 토큰을 마지막 자리에 추가해 줍니다.
이어서 샘플 A2가 생성을 할 때는 reference count가 1이므로 Block 1의 마지막 자리에 생성된 토큰을 저장하게 됩니다.
사실 설명은 sampling 때문에 되어 있긴 하지만, 공통된 입력 프롬프트를 중복하여 메모리 할당하지 않는다는 점은 그냥 Shared prefix라고도 볼 수 있을 것 같습니다.
이를 잘 나타내는 그림은 아래와 같습니다.
프롬프트 엔지니어링의 중요성이 점점 더 커지고 있는 만큼, 최근 입력 프롬프트는 정말 길고 복잡한 사례들이 많습니다.
few-shot을 사용하면 더 그렇고요..
위 예시에서는 실제 task input과 output은 굉장히 짧은 반면 shared prefix는 엄청나게 길다는 것을 알 수 있습니다.
vLLM에서는 LLM service provider에 의해 미리 정의된 일련의 shared prefix가 한 세트의 physical block을 형성하여 저장되고, 이를 반복적으로 참조한다고 볼 수 있겠습니다.
그래서 실제로 여러 데이터를 추론하고자 할 때 초기 로드 시간이 상대적으로 짧지 않습니다.
3.4.2. Beam search
언어 모델의 생성 과정에서 가장 흔히 사용되는 beam search 입니다.
매 스텝마다 가장 그럴싸한(등장할 확률이 높은) k개의 조합을 취하는 방식이죠. (예시에서 k=4)
그림을 보시면 딱 느낄 수 있는 것이, 입력 토큰만을 공유하는 parallel sampling 방식과 달리 다른 후보군의 생성 결과 일부를 공유하는 경우도 있다는 점입니다.
결국 고정된 shared prefix를 제외하고도 생성 과정에서 공통되는 부분을 효율적으로 관리할 필요성이 있다는 뜻이죠.
[예시]
처음에 Block 0으로부터 출발해서 Block 1, 2로 갈라집니다.
결과적으로 Block 0, 1, 3 은 후보 0, 1, 2가 공유하는 Block이 되고 후보 3만 다른 노선을 타게 되었습니다.
그런데 네 번째 단계에서 후보 0과 3은 각각 Block 5와 8에서 최후를 맞이하게 됩니다.
반대로 최종 후보는 1, 2에서 분기한 네 개 Block 9-12 이죠.
그러므로 이후 단계에서 더 이어지지 않는 후보 0과 3에 대한 logical block을 해제해야 합니다.
정확히는 이와 연결된 physical block의 reference count를 줄여야 하는 것이죠.
즉, Block 2, 4, 5, 8의 reference count는 0이 되어 메모리를 해제하고, vLLM은 새로운 Block 9-12에 메모리를 할당하게 되는 것입니다.
이전에는 중간 단계의 토큰들을 중복해서 메모리에 할당하게 됨으로써 메모리 낭비가 심했는데, 블록 단위로 메모리를 관리할 수 있어 이를 개선할 수 있었다고 해석 가능합니다.
3.5. Scheduling and Preemption
vLLM에서는 first-come-first-serve (FCFS) 를 원칙으로 scheduling 한다고 합니다.
쉽게 말하면 선입선출, 먼저 들어온 요청부터 먼저 처리한다는 것이죠.
이와 관련하여 사용되는 주요한 두 개념은 다음과 같습니다.
Swapping & Recomputation
아무래도 Block 단위로 GPU를 꽉꽉 채워 활용하다보면 스케쥴링이 완벽하게 되지 않을 가능성이 있나 봅니다.
물리적인 메모리 공간을 미리 할당하는 것이 아니라 한 블록 내에 더이상 새로운 토큰을 저장할 공간이 없을 때 할당해서 그런 것 같습니다.
그런데 이미 전체 GPU를 사용 중이어서 더이상 할당한 공간이 없는 경우 CPU의 Block을 사용하게 됩니다.
이때 sequence의 일부 KV cache셋을 GPU로부터 해제하여 CPU로 옮긴다고 설명하고 있습니다.
그리고 preempted sequence가 reschedule 되면 KV cache를 단순히 다시 계산한다는데 그 의미는 무엇인지 잘 모르겠습니다.
3.6. Distributed Execution
일반적으로 LLM은 single GPU로 감당하기 힘든 경우가 많아 여러 대의 GPU로 분산 처리해야 할 때가 많습니다.
vLLM은 흔히 사용되는 Megatron-LM 스타일의 tensor model parallelsim 을 효율적으로 지원할 수 있다고 설명하고 있습니다.
정확히는 저도 무슨 내용인지 잘 모르겠고..
흔히 사용되는 All-reduce 방식에 관한 것 같은데 관심이 있는 분들은 관련 자료를 추가적으로 찾아보시는 걸 추천드립니다.
어쨌든 vLLM에서는 KV cache manager를 메인 scheduler로 사용하여 여러 대의 GPU가 존재할 때의 연산을 관리하도록 설계되어 있는 것 같습니다.
저도 vLLM을 single GPU에서만 사용해봤는데 링크를 보시면 분산 환경에서의 추론 및 서빙에 대해 설명하고 있습니다.
4. Implementation
이 챕터에서는 Kernel-level Optimization, Supporting Various Decoding Algorithms에 대한 내용을 다루고 있습니다.
제가 이런 지식에 관해서는 아는 바가 많이 없어서.. 키워드만 정리하고 넘어가도록 하겠습니다 🙏🏻
[Kernel-level Optimization]
- Fused reshape and block write
- Fusing block read and attention
- Fused block copy
[Supporting Various Decoding Algorithms]
- fork
- append
- free
5. Evaluation
논문에서는 여러 실험에 대한 결과를 상세히 보고 및 분석한 내용을 설명하고 있습니다.
저는 이중에서 개인적으로 중요하다고 느꼈던 것 일부만 공유해보도록 하겠습니다.
위 그래프에 나타난 ShareGPT와 Alpaca는 데이터셋의 이름입니다.
OPT 모델이 두 데이터셋에 대해 추론하는 속도를 비교한 것이라고 볼 수 있겠습니다.
서로 다른 색의 선들은 각각 추론용 라이브러리이고 이중 vLLM은 파란색입니다.
x축은 request rate이므로 우측으로 갈수록 더 많은 request를 처리하게 됨을 의미하는데 기준은 vLLM입니다.
y축은 normalized latency로 갚이 높을수록 느리다는 것을 의미하며 이것 역시 기준이 vLLM입니다.
즉, 위 그래프는 vLLM을 기준으로 측정한 request와 latency의 정도가, 다른 라이브러리를 사용했을 때에 비해 얼마나 뛰어난지를 잘 보여주고 있는 것입니다.
예를 들어 (d) 그래프의 경우 vLLM이 초당 30개 가까운 request를 처리하는 동안, Orca는 다섯 개 정도의 request만을 처리할 수 있다는 걸 의미하는 것이죠.
노란색의 Orca (Pow2)는 절반 정도의 처리 속도를 나타내는 것이겠네요.
이 그래프 역시 vLLM의 성능이 Orca의 두 배 이상 뛰어나다는 것을 잘 보여주고 있습니다.
잠깐 언급하자면 ShareGPT를 구성하는 텍스트의 평균 길이가 Alpaca에서보다 훨씬 길기 때문에 값에 차이가 나는 것으로 이해할 수 있습니다.
(8.4x longer input prompts & 5.8x longer outputs)
vLLM의 뛰어난 성능은 한 번에 여러 후보를 생성하게 되는 parallel generation, beam search에서도 발휘될 수 있다는 게 눈에 띕니다.
그리고 이러한 경향은 공통된 prefix가 많고 길수록 효과적이라는 것도 아래 그래프를 통해 확인 가능합니다.
5-shot을 사용하게 되면 공통으로 참조하게 되는 shared prefix가 차지하는 메모리 공간이 많아지는데, 이전에는 활용하지 못했던 자원들을 vLLM에서는 병렬적으로 활용할 수 있게 되어 추론 속도가 훨씬 빨라지는 것으로 이해할 수 있습니다.
마지막으로 ablation study에 해당하는 내용 중 block size에 대한 비교가 있어서 그 결과를 가져왔습니다.
빨간색으로 표시한 부분을 보시면 알 수 있는 것처럼 block size는 8-32 사이의 값이 가장 적절합니다.
사이즈가 커지면 한 번에 할당하는 physical 메모리의 크기가 커지고 반대는 작아지는데 뭐든지 적당히가 중요한가 봅니다.
정확히 서술되어 있지는 않지만 뇌피셜을 펼쳐보자면, block size가 너무 작은 경우 매핑 테이블만 지나치게 커지고 이 테이블을 참조하는 횟수가 많아지면서 병목이 발생할 것 같습니다.
또 새로운 토큰을 저장할 공간이 많지 않아 새로운 메모리를 계속해서 할당해야 되는데 이는 결국 scheduler의 일이 복잡해진다는 뜻이 되겠죠.
반대로 block size가 너무 크다면 이전과 마찬가지로 GPU 내에 제대로 활용되지 못하고 묶여서 비효율을 초래하는 경우가 많아질 것입니다.
그렇기 때문에 상대적으로 텍스트 길이가 짧은 Alpaca 데이터셋의 latency가 빠르게 높아지는 것 같습니다.
저정도 크기의 block이라면 사실상 block 단위로 관리하지 않는 것과 비슷하게 되는 시점이라는 의미로 보입니다. (사이즈가 128 -> 256 될 때)
ShargeGPT의 데이터셋은 훨신 긴 텍스트로 구성되어 저정도 사이즈에서도 block 단위의 관리가 기존 대비 큰 이득이 되는 것 같고요.
오늘은 이렇게 PagedAttention, 그리고 이를 구현한 오픈소스 라이브러리 vLLM에 대해 디테일하게 알아보았습니다.
개인적으로는 지금까지 사용했던 도구들 중에 그 효용성이 가장 임팩트있게 느껴졌는데요.
보통 이정도의 속도 향상이 있다고 하면 성능을 잃게 되는 거 아니냐..는 생각이 드는데 그러지도 않습니다.
실제로 결과물이 그대로 잘 재현되고.. 논리적으로는 저장되는 메모리 공간이 달라지는거지 연산 결과가 달라질 이유가 없어서요.
물론 깊은 하드웨어 지식단으로 내려가면 어떤 비밀이 숨겨져 있을지는 잘 모르겠습니다만 ㅎㅎ
여튼 아직까지도 model.generate 로 고통받고 계신분이 있다면 반드시 꼭 vLLM 써보시길 추천드립니다.
LLM을 위한 다른 추론 라이브러리들도 많고 이보다도 더 빠르다고 알려진 것도 있으니 상황에 맞게 쓰는 게 중요한 것 같습니다.
