앞선 포스트에서는 Large Language Model (LLM) 인퍼런스의 중요성과 왜 우리는 LLM을 효율적으로 활용해야 하는지를 알아보았다. 더불어, LLM을 이용한 문장 생성은 autoregressive generation이며, 해당 생성 과정에서 사용할 수 있는 다양한 디코딩 전략을 소개하였다. 지난 포스트는 아래 링크를 참조하자.
LLM 인퍼런스 훑어보기 (1) - LLM을 이용한 문장 생성
인공 지능과 기계 학습 기술의 발전은 현대 사회에 혁명적인 변화를 가져왔다. 특히, Large Language Model(LLM)과 같은 최신 기술은 자연어 처리 및 이해 분야에서 차별화된 성능을 보이며, 다양한 분
dytis.tistory.com
이번 포스트에서는 LLM의 근간이 되는 Transformer 구조와, encoder, decoder 모델 구조에 대해서 다루고자 한다. 더 나아가 이러한 구조에서 효율적인 인퍼런스를 위하여 도입된 KV Cache라는 개념 또한 알아본다.
Transformer
트랜스포머(Transformer)는 2017년에 구글이 발표한 "Attention is all you need"라는 논문에서 제안된 모델이다. 이 모델은 기존의 seq2seq 구조를 따르면서도 이름처럼 어텐션(Attention) 메커니즘만을 사용하여 구현되었다. RNN을 사용하지 않고도 인코더-디코더 구조를 설계한 이 모델은 번역 성능에서 RNN보다 우수한 결과를 보였다.
기존의 seq2seq 모델은 인코더-디코더 구조로 이루어져 있었으며, 인코더는 입력 시퀀스를 하나의 벡터 표현으로 압축하고, 디코더는 이 벡터 표현을 기반으로 출력 시퀀스를 생성한다. 그러나 이러한 구조는 입력 시퀀스의 정보가 일부 손실될 수 있는 단점이 있었으며, 이를 보완하기 위해 어텐션 메커니즘이 도입되었다. 그렇다면 어텐션 메커니즘은 무엇일까?
Attention
어텐션의 기본 아이디어는 디코더가 출력 단어를 예측하는 각 시점(time step)마다 인코더에서의 전체 입력 문장을 다시 한 번 참고하는 것이다. 그러나 모든 입력 단어를 동일한 비율로 고려하는 것이 아니라, 해당 시점에서 예측해야 할 단어와 관련 있는 입력 단어 부분에 더 집중(attention)하여 보게 된다.
어텐션 함수는 일반적으로 다음과 같이 표현됩니다:
Attention(Q, K, V) = Attention Value
어텐션 함수는 주어진 '쿼리(Query)'에 대한 모든 '키(Key)'와의 유사도를 개별적으로 계산하고, 이러한 유사도를 해당 키와 연결된 각 '값(Value)'에 반영한다. 그런 다음, 이 유사도가 반영된 '값(Value)'을 모두 합산하여 반환한다. 이러한 결과를 Attention Value라고 한다. 일반적으로 Q, K, V는 다음과 같이 정의할 수 있다.
- Q: 모든 시점의 디코더 셀에서의 은닉 상태들
- K: 모든 시점의 인코더 셀의 은닉 상태들
- V: 모든 시점의 인코더 셀의 은닉 상태들
셀프 어텐션은 Q, K, V가 전부 동일한 것을 말한다. 트랜스포머에서는 Encoder Self-attention, Masked Decoder Self-attention, Encoder-Decoder Attention의 세 가지의 어텐션이 사용된다. 당연하게도, 두 self-attention은 자신의 동일한 Q, K, V를 가지지만 세번째 Encoder-Decoder Attention에서는 Query가 디코더의 벡터이고 Key와 Value는 인코더의 벡터이다.
인코더의 셀프 어텐션을 직관적으로 이해해보자면, 입력 문장 내의 단어들끼리 유사도를 구한다라는 의미이다. 이는 아래 figure로 더 쉽게 이해할 수 있다.
Q, K, V 벡터 얻기
앞서 설명한 대로, 셀프 어텐션은 입력 문장의 단어 벡터를 기반으로 작동한다. 그러나 실제로는 셀프 어텐션을 수행하기 위해 인코더의 초기 입력인 단어 벡터들로부터 Q벡터, K벡터, V벡터를 얻는 과정이 필요하다. 이 때, 이러한 Q벡터, K벡터, V벡터는 초기 입력의 차원과는 달리 더 낮은 차원을 가지며, 트랜스포머에서는 초기 입력으로부터 512차원의 각 단어 벡터를 64차원의 Q벡터, K벡터, V벡터로 변환했다.
위의 64라는 값은 트랜스포머의 다른 하이퍼파라미터인 num_heads로 결정된다. 트랜스포머에서는 num_heads 값을 8로 설정하였고, 이로인해 512/8=64차원의 Q벡터, K벡터, V벡터로 변환해야 한다.
예를 들자면, student라는 단어의 512차원 임베딩 벡터에 512 X 64 차원의 크기를 가진 Q 가중치, K 가중치, V 가중치 행렬을 각각 곱하여 Q, K, V 벡터를 얻어낸다. Q, K, V 벡터를 얻었다면 지금부터는 기존의 어텐션 메커니즘과 동일하다. 각 Q벡터는 모든 K벡터에 대해서 어텐션 스코어를 구하고, 어텐션 분포를 구한 뒤에 이를 사용하여 모든 V벡터를 가중합하여 어텐션 값 또는 컨텍스트 벡터를 구한다. 그리고 이를 모든 Q벡터에 대해서 반복한다.
Scaled dot-product Attention
트랜스포머에서는 Scaled dot-product Attention이라는 어텐션 함수를 사용한다. 아래는 "I am a student"라는 문장에 대해서 sdpa를 이해할 수 있는 좋은 자료이다.
위의 그림은 단어 I에 대한 Q벡터가 모든 K벡터에 대해서 어텐션 스코어를 구하는 것을 보여준다. 이러한 과정은 am에 대한 Q벡터, a에 대한 Q벡터, student에 대한 Q벡터에 대해서도 모두 동일한 과정을 거친다. 여기서 attention score는 각각 단어 I가 I, am, a, student와 얼마나 연관되어 있는지를 보여주는 수치다. 트랜스포머에서는 두 벡터의 내적값을 스케일링하는 값으로 K벡터의 차원의 제곱근을 나눈다. 위에서 K벡터를 64차원으로 사용한다고 했으니, 실제 스케일링 값은 1/8이다.
이제 어텐션 스코어에 소프트맥스 함수를 사용하여 어텐션 분포(Attention Distribution)을 구하고, 각 V벡터와 가중합하여 어텐션 값(Attention Value)을 구한다. 이를 단어 I에 대한 어텐션 값 또는 단어 I에 대한 컨텍스트 벡터(context vector)라고도 할 수 있다. am에 대한 Q벡터, a에 대 Q벡터, student에 대한 Q벡터에 대해서도 모두 동일한 과정을 반복하여 각각에 대한 어텐션 값을 구한다.
행렬 연산으로 일괄 처리하기
사실 각 단어에 대한 Q, K, V 벡터를 구하고 스케일드 닷-프로덕트 어텐션을 수행하였던 위의 과정들은 벡터 연산이 아니라 행렬 연산을 사용하면 일괄 계산이 가능하다. 지금까지 벡터 연산으로 설명하였던 이유는 이해를 돕기 위한 과정이고, 실제로는 행렬 연산으로 구현된다.
Multi-head attention
앞의 어텐션에서는 512의 차원을 가진 단어 벡터를 num_heads로 나눈 차원을 가지는 Q, K, V 벡터로 바꾸고 어텐션을 수행하였다. 트랜스포머의 연구진은 한 번의 어텐션을 하는 것보다 여러번의 어텐션을 병렬로 사용하는 것이 더 효과적이라고 판단하였고, 512 차원을 num_heads개로 나누어 64차원을 가지는 Q, K, V에 대해서 num_heads개의 병렬 어텐션을 수행한다. 이 때 각각의 어텐션 값 행렬을 어텐션 헤드라고 부른다.
트랜스포머에는 다른 다양한 기술과 개념들이 많지만, 가장 대표적인 어텐션에 대한 내용은 이정도로 마무리 하자. 정리하자면, LLM을 추론하는 과정에서 어텐션 레이어를 거치기 전 입력 단어들에 대해 Q, K, V 가중치를 곱하여 Q, K, V 벡터를 획득하는 과정이 필요하다. 여기서 다시 생각해보면 셀프 어텐션의 경우에는 K, V 값이 이전 단어 생성에서 이미 구해졌으며, 다음 단어 생성 시에도 사용이 된다는 것을 알 수 있다.
KV Cache
앞서 살펴봤듯이, transformer model의 추론에서 샘플링 과정은 먼저 주어진 프롬프트 또는 컨텍스트를 처리하는 단계가 이루어지며, 그런 다음 추가 토큰을 하나씩 샘플링하면서 (= autoregressive) 트랜스포머는 self-attention을 수행한다. 이 self-attention을 위해 현재 시퀀스에 있는 각 항목의 key-value 값을 필요로 하며, 이러한 벡터들은 KV Cache 또는 Past Cache라고 불리는 행렬로 제공된다. Past Cache는 다음과 같은 형태를 가진다: [batch_size, 2, num_heads, seq_len, features].
토큰을 샘플링할 때마다 이러한 벡터들에 대한 중복계산을 피하는 것이 이것의 목적이다. 이미 계산된 k, v 값이 있다면, 우리는 약간의 저장 비용으로 상당한 양의 계산 비용을 절약할 수 있다.
토큰마다 우리가 새롭게 저장하는 bytes는 2 X 2 X n_layers X n_heads X d_head이다. 가장 앞의 팩터 2는 k, v의 두 벡터의 수를 의미하고, 각 레이어마다의 k, v값을 저장한다. 이 때 k, v 행렬은 n_heads X d_head 행렬이고, 2번째 팩터 2는 bytes의 수를 의미한다. (16-bit 포맷이라고 가정)
이 때 절약할 수 있는 flops는 2 X 2 X n_layer X d_model^2이다. token embedding t_e를 W_k로 곱해야 되며, 여기에 2 X d_model^2 flops가 소요된다. 우리는 k와 v에 대해서 한번씩 계산해야 하므로 2를 곱해주고, 이러한 과정을 n_layer만큼 반복해야 한다.
OPT-30B 모델이 있고, seq_len = 1024, batch_size = 128일 때, KV cache에 소요되는 메모리는 다음과 같다:
batch_size(128) X 2 X 48(n_layers) X 7168 (d_model) X 1024(seq_len) = 180 GB
한 가지 염두해야 할 점은 모델은 결국 2 X 30GB = 60GB이기 때문에 KV cache는 3배의 용량을 차지한다는 점이다. 즉 LLM의 inference에서의 중요한 측면은 이 작업 자체가 memory bound라는 것일 수 있다.
Huggingface의 transformers에서는 use_cache라는 parameter를 제공하여, kv cache를 켜거나 끌 수 있다.
이번 포스트에서는 transformer의 구조를 알아보고, inference의 속도향상을 위해 도입된 kv cache의 개념을 확인했다. 다음 포스트에서는 좀 더 다양한 측면에서의 가속 방법론을 다룰 생각이다.
참고한 글
- https://wikidocs.net/31379
- https://velog.io/@cha-suyeon/%EC%9D%B8%EC%BD%94%EB%8D%94%EC%9D%98-%EC%85%80%ED%94%84-%EC%96%B4%ED%85%90%EC%85%98-%EC%9D%B4%EC%A0%90-Q-K-V%EC%9D%98-%EC%A0%95%EC%9D%98
- https://kipp.ly/transformer-inference-arithmetic/
- https://lilianweng.github.io/posts/2023-01-10-inference-optimization/
- https://www.dipkumar.dev/becoming-the-unbeatable/posts/gpt-kvcache/
'AI > 모델 인퍼런스' 카테고리의 다른 글
| LLM 인퍼런스 훑어보기 (6) - quantization (0) | 2023.11.17 |
|---|---|
| LLM 인퍼런스 훑어보기 (5) - continuous batching (0) | 2023.11.10 |
| LLM 인퍼런스 훑어보기 (4) - kernel fusion (0) | 2023.10.14 |
| LLM 인퍼런스 훑어보기 (3) - KV Cache (deep dive) (0) | 2023.09.23 |
| LLM 인퍼런스 훑어보기 (1) - LLM을 이용한 문장 생성 (0) | 2023.09.14 |