Book Review : Foundations of Large Language Models (4)

Preliminary : Compute-bound or Memory-bound 

 이번 Chapter에 들어가기 전에, 배경지식으로 알아두면 좋은 게 있다. 바로 Compute-bound와 Memory-bound에 관한 얘기이다. Compute-bound와 Memory-bound에 대해서 얘기하려면, 먼저 Hardware Arithmetic Intensity라는 물리량을 정의해야한다.  Hardware Arithmetic Intensity := (max FLOP/s) / {max memory load(GB)/s}로 정의된다. Hardware의 1초 최대 계산량을 최대 Bandwidth로 나눈 것이다. 즉 1byte memory load를 할 때 FLOP을 몇번을 해야 memory load와 computation이 hardware에서 균형을 이루고 있는지를 나타내는 물리량이다. 만약 현재 program의 Arithmetic Intensity가(program Arithmetic Intensity := {(program FLOP/s)/program's memory load(GB)/s}) Hardware Arithmetic Intensity보다 작다면 이 프로그램은 메모리에 비해서 계산이 놀고있는 것이기 때문에 Memory-bound application이라고 할 수 있고, 현재 program의 Arithmetic Intensity가 Hardware Arithmetic Intensity보다 크다면, 메모리가 계산에 비해서 놀고있는 것이기 때문에, compute-bound application이라고 할 수 있다. 

Engineering technics for LLM Inference : prefill vs decoding

 이제 본격적으로 LLM Inference를 위한 Engineering technic에 대해서 살펴보도록 하자. Inference는 크게 총 두 가지 phase로 이루어진다. 첫 번째는 prefill phase이고 두 번째는 decoding phase이다.

 먼저 prefill phase는 user input x를 바탕으로 KV cache를 채우는 작업이다. 이때 특이한 점은 우리는 모든 토큰을 알고있으므로, 모든 token이 parallel하게 처리될 수 있다는 것이다. (이전 layer의 계산이 모두 끝나야 다음 layer의 계산을 할 수 있긴 하지만, 같은 layer안에서는 완전히 병렬처리가 가능하다) 또한, 한 token의 k,v가 계산에 여러번 쓰이게 되므로, 메모리 한번 load당 computation을 여러번 하게 되고, 보통의 hardware spec에서는 prefill phase에서는 computation을 할 동안 memory load가 놀고있어서 compute-bound상태라고 한다. 

 두 번째로 Decoding phase는 output을 auto-regressive하게 만들어내는 과정이다. 이 과정은 한번 load한 key와 value를 계산에 한번만 쓰는데, 보통의 hardware spec에서는 memory-bound라고 한다. iteration i-1에서 있던 KV cache를 L1 cache같은 곳에 저장해두고 쓰면 memory load가 거의 필요없기때문에 compute-bound일 수도 있지 않냐라고 할 수 있는데, 보통 KV cache는 L1 cache에 비해서 엄청나게 크기때문에 불가능하다. KV cache는 하드웨어 cache를 말하는게 아니고 data structure일 뿐인걸 상기하자. 

 

밑의 표는 prefill과 decode phase를 비교한 표이다. 여기서 말하는 Computational cost는 그냥 execution time이다. 보통 decoding stage의 memory load가 너무 오래 걸려서 전체 decoding stage가 보통 prefill stage보다 오래걸린다. (한 decoding stage는 보통 prefill stage보다는 빠르다. 혹시 헷갈릴까봐 한 prefill stage라는것은 없는게, prefill stage는 병렬로 모든 토큰이 한번에 처리된다는것을 기억하자.) 

Engineering technics for LLM Inferece : Decoding Algorithms(Searching Algorithms)

 Decoding phase에서 원래 보통은 매 itereation마다 그냥 확률이 제일 높은 token을 다음 token으로 고르게 된다. 하지만 전체 sequence의 maximum likelihood 관점에서 봤을 때 그렇게 해서 만든 전체 sequence는 optimal이 아니라 sub optimal solution이다. 전체 sequence의 log-likelihood는 sequence내의 각 token이 나올 확률을 전부 log씌워서 더하면 되는데, 진짜 optimal은 전체 sequence가 나올 확률이 최대가 되어야한다. 하지만 optimal solution을 구하려면 가능한 모든 sequence의 loglikelihood를 전부 구해야하므로, sequence의 길이가 n이라고 했을 때, vocabulary size의 n승 만큼 탐색을 해야하는데 이건 feasible하지 않다. 따라서 어떻게 하면 좋은 sub optimal solution을 구할 수 있을지에 대한 방법들을 지금부터 소개한다.  

 첫 번째는 Greedy Decoding이다. 원래 하듯이 매 step에서 확률이 제일 큰 token으로 고르는 것이다. 

 두 번째 방법으로는 Beam Decoding이 있다. 지금 step까지 sentence log likelihood top-k를 유지시키는 것이다. 즉 k개의 sentence를 유지시키고, 다음 step을 각 sentence마다 진행해서, k*(Vocab Size)의 candidate들 중에서 다시 top-k를 뽑는 것이다. inference를 k번 해야한다는 단점이 있다. 

 세 번째 방법으로는 top-k Sampling이 있다. 매 step마다 top k 중에서 확률적으로 1개를 sampling하는 기법이다. 

 마지막 방법으로는 top-p Sampling이 있다. 매 step마다 확률순으로 vocab을 정렬하고, 가장 확률이 큰 것부터 accumulative sum을 해서 그 sum이 p보다 커지는 순간 그 뒤의 vocab들은 전부 삭제해버린다. 이후 살아남은 vocab들을 가지고 1개를 확률적으로 sampling하는 방법이다. 그림을 참고하면 이해가 잘 될 것이다. 

Engineering technics for LLM Inferece : Decoding with Penalty terms 

 다음으로 소개할 technic은 Decoding with Penalty terms이다. 한 Decoding stage가 끝나고 나면 결국 다음 토큰에 대한 후보 토큰들의 확률이 나오게 되는데, 이 확률을 사후 조정하는 테크닉이다. 예를들면 앞에서 나왔던 단어가 너무 많이 등장하지 않도록 확률 조정을 사후처리로 할 수 있다. 이러면 하나의 글 내에서 어휘의 다양성을 높일 수 있다. 

<Decoding with Penalty terms>

 

Engineering technics for LLM Inferece : Speculative Decoding

 다음으로 소개할 technic은 Speculative Decoding이다. 작은 모델과 큰 모델을 모두 decoding에 이용해서 큰 model만을 활용해서 decoding을 할 때보다 속도를 높이는 방법인데, 방법은 다음과 같다. 먼저, 작은 모델로 k개의 token을 미리 예측한다. 그리고 큰 모델에서는 병렬로 그 k개의 token이 진짜 맞았는지 검증을하면 된다. 예를 들면 작은 모델에서 <start> 다음의 3개의 output이 "i am seokmin"으로 나왔다면, 큰 모델에서는 "i am seokmin"을 그냥 input으로 넣고 병렬로 attention을 수행해서 <start>뒤에 나올 토큰 후보들의 확률분포와 <start> i 다음에 나올 토큰들의 확률분포와 <start> i am 다음에 나올 토큰들의 확률분포를 병렬적으로 계산할 수 있게 된다. best 시나리오는 만약 <start>다음에 올 확률이 최대인 토큰이 i였고, <start> i 다음에 올 확률이 최대인 토큰이 am이었고, <start> i am 다음에 올 확률이 최대인 토큰이 seokmin인 시나리오일 것이다. 이 경우에는 우리는 원래 큰 모델을 3번 inference해야하는 시간을 작은 모델 3번 + 큰 모델 1번 inference하는 시간으로 줄 일 수 있다. 하지만 보통은 i am다음에 seokmin이 확률이 최대인 토큰이 아닐 것이다. 이런 경우에는 맞은 token들 까지만 사용하고, 뒤에 틀린 토큰들은 버려버린 후에 big model로 가장 처음 틀린 토큰을 다시 바로잡고, 그 이후에는 작은 모델로 k개를 예측하면서 과정을 반복하면 된다.   

 

Engineering technics for LLM Inferece : Multi-User Inference Request - batching & scheduling

 LLM을 serving하는 입장에서, 당연히 multi-user로부터 Inference Request가 올 수 있다. 이 경우에 어떻게 대처해야하는지 알아보자. 기본적으로는 Scheduler를 만들어서 여러 request들을 batching하는 작업을 해야한다. 

 스케쥴링을 할 때, 최대한 길이가 비슷한 것 끼리 같은 batch안에 넣으면 된다. 이렇게 하면 쓸데없는 padding을 줄일 수 있기 때문에 매우 효율적이다. 만약 user들 간의 input이 너무 길이가 다르다면, disaggregation of prefilling and decoding방법을 쓰는 것도 좋다 disaggregation of prefilling and decoding 방법을 이용하면 prefilling이랑 decoding을 완전히 분리시키기 때문에,(즉 prefill이 끝나고 메모리 안에 있는 값을 재정렬한다) Masked Attention을 이용해서 prefill에서 여러 다른 user input들을 하나의 input으로 이어붙여서 동시에 prefill을 할 수 있다. 그림을 보면 이해가 잘 될 것이다.

 

 조금 더 Scheduling Algorithm에 대해서 깊게 들어가보자면, 보통 Scheduling에는 새 request를 현재 처리중인 batch에 바로 넣을 수 있는지 없는지에 따라 두 가지 알고리즘이 존재한다. 바로 한 batch에 대한 연산이 끝나기 전까지는 batch에 새로운 request를 넣을 수 없다는 Request-level Scheduling과 한 batch 에 대한 연산중에 batch안에 새로운 request를 넣을 수 있다는 Iteration-level Scheduling이다. 밑의 그림의 예시를 참고하자

Iteration-level Scheduling의 가장 대표적인 방법으로는 Countinous Batching이란 방법이 있다. 한 iteration안에서 일어날 수 있는 작업을 prefilling전체 혹은 decoding 한 stage로 정의하고, iteration이 끝날 때 마다 batch를 조정할 수 있는 방법이다. 특이한 건 한 batch 안에서 prefilling과 decoding 한 stage가 동시에 일어날 수 있다고 가정하는 것이다. 자세히는 책에서 소개하진 않지만, compute-bound인 prefill stage와 memory-bound인 decoding stage가 서로를 보완한다고 한다. 

<Continous batching>

 

Countinous Batching은 좋아보이지만, 하나의 문제점이 있다. 보통 prefill stage는 한 decoding stage에 비해서 훨씬 오래걸리기 때문에, 만약 한 batch 안에 prefill stage와 한 decoding stage가 같이 있으면, decoding phase에 있는 데이터는 너무 많이 pending을 해야한다. 이를 해결하기 위해서 등장한 개념이 바로 Chunked Prefill이다. prefill을 chunk단위로 나눠서 하자는 것이다. prefill-stage를 제대로 이해하고 있는 사람이라면, chunked prefill의 동작 방식은 바로 감을 잡을 수 있을 것이다. 

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

 

다음 chapter에서는 user가 inference를 할 때 더 좋은 output을 낼 수 있도록 하는 prompt engineering에 대해서 소개할 것이다.