커뮤니티2

Transformer 아키텍처는 현재 LLM 에서 사용하는 de facto의 모델 구조입니다.

처음 Vaswani가 제안했을때에는 Encoder와 Decoder를 함께 사용하는 언어번역 테스크를 위해 고안되었으나, 각 Encoder와 Decoder의 특징에 따라서 다양한 모델로 적용되고 있습니다.

여러분들이 많이 사용하고 계신 GPT의 원류 또한, Vaswani의 Attention is All you needs의 Decoder를 사용한 모델입니다.

그런데 현재 LLM을 통한 생성형 AI가 서비스화가 되면서 다양한 문제를 직면하였습니다.

Decoder를 사용하는 모델이 문장을 생성할때에 "Auto-Regression" 방식을 사용한다는 것입니다. 이는 Prompt를 모델에게 제공하면, 모델은 이 Prompt를 받아서 다음 단어를 뽑고하고, (Prompt + 다음 단어)를 다시 모델에 제공하여 그 다음 단어를 뽑는 방식입니다. 이런 방식은 우선 문장을 생성하는 시간이 오래 소요됩니다. 이러한 방법은 서비스 형태에서는 사용자에게 큰 이슈로 적용될수 있습니다.

해당 이슈를 해결하기 위해서, 문장 생성까지의 시간을 줄이기 위해 고안된 방법은 다양한데, 그중 오늘은 KV-Cache라는 방식에 대해서 알아보도록 하겠습니다. 

우선 KV-Cache에 대해 설명하기 앞서서, Transformer (Decoder) 구조에 대해서 다시한번 살펴뵤고, 어떻게 적용되는지에 대해 설명해보고자 합니다.

Transformer는 아래와 같은 구조를 가지고 있습니다.

왼쪽의 Encoder와  오른쪽의 Decoder로 구성된  Transformer 아키텍쳐 입니다.

포스트의 주제인 KV-Cache를 이해하기 위해서는 Muti-Head Attention 블럭의 내부구조를 이해해 봐야 합니다.

Multi-Head Attention은 다음과 같이 구성되어있습니다.

조금 더 설명하기 쉽게 예시를 통해 Prompt가 어떻게 변환 되는지 보겠습니다.

```python

import torch

from transformers import AutoTokenizer

input_text = "Hello, how are you?"

# Tokenize the input text

tokenizer = AutoTokenizer.from_pretrained('gpt2-medium')

x = tokenizer(input_text, return_tensors="pt")

x = x['input_ids']

print('Shapf of the input tensor:', x.shape)

print('Input tensor:', x)

```

```

Shape of the input tensor: torch.Size([1, 6])

Input tensor: tensor([[15496, 11, 703, 389, 345, 30]])

```

위의 예시는 "Hello, how are you"가 토크나이저를 통해 변환된 결과 입니다.

토크나이저에 학습된 Vocab은 다음과 같습니다.

```

Token: 15496 -> Hello 

Token: 11 -> , 

Token: 703 -> how 

Token: 389 -> are 

Token: 345 -> you 

Token: 30 -> ?

```

* 예시를 통해 아키텍쳐의 Dimension은 32로 축소하였습니다.

```python

embeddings = torch.nn.Embedding(tokenizer.vocab_size, 32)

emb_output = embeddings(x)

print(emb_output.shape)

```

```

torch.Size([1, 6, 32])

```

Batch Size: 1, Sequence_length: 6, Dimension: 32인 결과를 얻을수 있습니다.

* Positional Encoding 또는 Rotary Encoding을 적용하는 부분은 이전 블로그를 참고해주세요.

* 예시는 Positional Encoding을 가정한 결과로 가정합니다.

- Multi-Head Attention 부분을 살펴보면, 아래와 같은 흐름으로 진행 됩니다.

```python

head = 2

head_dim = 32//2

wq = torch.nn.Linear(32, head*head_dim)

wk = torch.nn.Linear(32, head*head_dim)

wv = torch.nn.Linear(32, head*head_dim)

wq_output = wq(emb_output)

wk_output = wk(emb_output)

wv_output = wv(emb_output)

print('Query shape:', wq_output.shape)

print('Key shape:', wk_output.shape)

print('Value shape:', wv_output.shape)

```

```

Query shape: torch.Size([1, 6, 32]) 

Key shape: torch.Size([1, 6, 32]) 

Value shape: torch.Size([1, 6, 32])

```

- Multi-Head

```python

wq_head_output = wq_output.view(1, 6, head, head_dim)

wq_head_output = wq_head_output.transpose(1, 2)

wk_head_output = wk_output.view(1, 6, head, head_dim)

wk_head_output = wk_head_output.transpose(1, 2)

wv_head_output = wv_output.view(1, 6, head, head_dim)

wv_head_output = wv_head_output.transpose(1, 2)

print('Query head shape:', wq_head_output.shape)

print('Key head shape:', wk_head_output.shape)

print('Value head shape:', wv_head_output.shape)

```

```

Query head shape: torch.Size([1, 2, 6, 16]) 

Key head shape: torch.Size([1, 2, 6, 16]) 

Value head shape: torch.Size([1, 2, 6, 16])

```

Batch Size: 1, Heads: 2, Sequence Length: 6, Dimension:16 인 텐서를 Q(query), K(key), V(value) 별로 얻을수 있습니다.

다음은 Scaled-Dot Product Attention을 적용시켜 보겠습니다.

```python

score = torch.matmul(wq_head_output, wk_head_output.transpose(-2, -1))

score = score / (32**0.5) # scale

print('Score shape:', score.shape)

attention = torch.nn.Softmax(dim=-1)(score)

print('Attention shape:', attention.shape)

output = torch.matmul(attention, wv_head_output)

print('Output shape:', output.shape)

```

```

Score shape: torch.Size([1, 2, 6, 6]) 

Attention shape: torch.Size([1, 2, 6, 6]) 

Output shape: torch.Size([1, 2, 6, 16])

```

마지막으로 Concat과 Linear를 통하면 Multi-Head Attention의 내부 구조가 됩니다.

```python

wo = torch.nn.Linear(32, 32) 

output = output.transpose(1, 2).contiguous().view(1, 6, 32)

mha_output = wo(output)

print('MHA output shape:', mha_output.shape)

```

```

MHA output shape: torch.Size([1, 6, 32])

```

Multi-Head Attention 블럭 이후에는 Feed-Forward Layer를 통해 아키텍처의 블럭결과를 받을수 있지만, KV-Cache를 이해하기 위해서는 FFN에 대해서는 생략합니다.

KV-Cache

KV-Cache는 이름 그대로 Attention에서 K와 V에 대한 값을 Cache로 저장하는 방식을 얘기합니다.

이런 방식이 어떻게  문장 생성시 시간을 줄일수 있는지 한번 확인해보겠습니다.

문장 생성을 요청 할 시, Prompt 또는 Text를 모델에 제공하면, 모델은 다음과 같이 반복적인 작업을 수행하며, 다음 문장을 추출합니다.

```

Input: Large language models are recent advances in deep learning

Step 0 input: Large language models are recent advances in deep learning 

Step 1 input: Large language models are recent advances in deep learning, 

Step 2 input: Large language models are recent advances in deep learning, which Step 3 input: Large language models are recent advances in deep learning, which uses 

Step 4 input: Large language models are recent advances in deep learning, which uses deep 

Step 5 input: Large language models are recent advances in deep learning, which uses deep neural 

Step 6 input: Large language models are recent advances in deep learning, which uses deep neural networks 

Step 7 input: Large language models are recent advances in deep learning, which uses deep neural networks to 

Step 8 input: Large language models are recent advances in deep learning, which uses deep neural networks to learn 

Step 9 input: Large language models are recent advances in deep learning, which uses deep neural networks to learn to

```

예를 들면, Step 0은 초기 Input을 받아서 가장 확률이 높은 토큰을 뽑고 해당 토큰은 , (쉼표)가 되었습니다.

Step 1은 Input과 , (쉼표)를 받아서 다음 단어인 which를 추출하였습니다.

이러한 방식으로 문장을 생성하는데, 위의 예시를 보면, 동일한 문장을 반복적으로 사용하는 것을 알수 있습니다.

이러한 식으로 문장을 생성할때  모델은 엄청난 FLOPS (Floating point operations per second)가 필요합니다. Attention Layer에서의 FLOPS의 계산은 아래와 같은 수식을 가지게 됩니다.

```

2 * batch_size * N_Layers * N_head * head_dim * (sequence_length**2)

```

여기서 가장 중요하게 보아야 할 점은, Attention Score를 구할때, Scaled Dot Attention의 Q와 K의 Dot Product로 인해서 (Sequence_length ** 2) 처럼 FLOPS가 Quadratic으로 증가한다는 것입니다.

많이 사용하는 모델들의 파라미터는 아래와 같습니다.

출처: https://medium.com/@plienhar/llm-inference-series-4-kv-caching-a-deeper-look-4ba9a77746c8

Attention만을 보았을때도 엄청난 FLOPS가 필요로 하는것을 볼수 있습니다.

FLOPS의 증가는 연산량의 증가로, 결국에는 긴 문장 (Context)를 생성하는데에는 연산량이 Quadratic (제곱)으로 증가하는 것을 볼 수 있습니다.

그래서, 문장 생성시, FLOPS을 최소화 하기 위해  KV-Cache가 고안되었습니다.

KV-Cache의 큰 틀은, Inference시에, 0-step (Initial Step)이후에, <Inital + step-0-output>을 합친 문장을 Inference하는 1-Step부터 사용됩니다.

예를 들어, 

```

0-Step-Input: Large language models are recent advances in deep learning

0-Step-Output: ,

1-Step-Input: Large language models are recent advances in deep learning,

```

1-Step-Input이 모델에 들어가서 1-Step-Output을 추출할때, KV-Cache가 없을시에는, 1-Step-Input이 그대로 모델의 인풋으로 사용됩니다.

그런데, 0-Step-Input과 1-Step-Input에서 동일한 단어들이 발생하는데, 바로

```

 Large language models are recent advances in deep learning

```

이 문장이 반복적으로 사용된다는 점 입니다.

KV-Cache는 이렇게 반복적인 연산의 결과를 Cache로 저장하는 방식을 채택하여, 불필요한 연산 (FLOPS)를 줄이는 최적화 방법입니다.

KV-Cache의 최적화 방법을 조금더 자세히 봐보면 다음과 같습니다.

0-Step-Input을 통해 생성된 인풋은 다음과 같습니다.

```

x = tokenizer.encode('Large language models are recent advances in deep learning', return_tensors='pt')

print(x.shape)

```

```

torch.Size([1, 9])

```

0-Step의 Attention연산은 기존의 Vanilla Transformer와 동일합니다. 하지만 KV-Cache를 사용하기 위해 Attention Layer 내부에 K_Cache와 V_Cache에 대한 홀더를 가지고 있도록 합니다.

```

cache_k = torch.zeros((1, 50, 2, 16)) # b, max_seq_len, head, head_dim

cache_v = torch.zeros((1, 50, 2, 16)) # b, max_seq_len, head, head_dim

```

- 모델 학습시에도 동일한 레이어를 사용하기에, KV-Cache를 사용하여 Attention Layer를 다시 구성해야 합니다

구현 코드는 아래에서 찾아볼수 있습니다.

https://github.com/meta-llama/llama/blob/main/llama/model.py 

https://github.com/meta-llama/llama/blob/main/llama/generation.py

KV-Cache의 구현에서는 0-Step에 대해서, Attention Layer는 cache_k 와 cache_v에  "transpose" 하기 이전의  wk_head_output 와  wv_head_output을 저장합니다.

- KV-Cache에 대한 대략적인 내용으로 Llama의 구현코드에서 많은 부분을 생략하였습니다.

```python

wq_head_output = wq_output.view(1, 6, head, head_dim)

wk_head_output = wk_output.view(1, 6, head, head_dim)

cache_k = wq_head_output

wv_head_output = wv_output.view(1, 6, head, head_dim)

cache_v = wv_head_output

```

Attention 에 사용하는 K와 V는, cache_k와 cache_v에서 파싱 하도록 합니다.

```python

bsz, seqlen, _ = x.shape

keys = cache_k[:bsz, : start_pos + seqlen]

values = cache_v[:bsz, : start_pos + seqlen]

```

여기서 중요한점은 start_pos라는 파라미터로, 이는 어텐션 레이어의 argument로 받도록 합니다.

```python

def forward(

        self,

        x: torch.Tensor,

        start_pos: int,

        freqs_cis: torch.Tensor,

        mask: Optional[torch.Tensor],

    ):

    ...

    return output

```

KV-Cache를 사용한 Attention Block은 다음과 같은 그림을 가지게 됩니다.

출처: https://medium.com/@plienhar/llm-inference-series-3-kv-caching-unveiled-048152e461c8

0-Step은 이전 정의가 없기때문에 K와 V는 그대로 사용하도록 합니다.

이때 모델내부에 있는 Multi-Head Attention은 해당 K와 V를 각 cache에 가지고 있게 됩니다.

여기서 중요한점은 1-Step의 Input으로 (0-step-input + 0-step-output)을 모델에 사용하지 않고, 0-step-output만을 모델에 사용하는 것이 KV-Cache의 핵심입니다.

 하나의 단어를 통해서 Q, K, V의 인풋으로는 이전 Vanilla Transformer 에서는 Embedding Layer의 아웃풋인 [Batch_size, Sequence_length , Dimension]이었던 [1, 6, 32]가 [1, 1, 32]로 축소되게 됩니다.

이로 인해 FLOPS의 연산이 Quadratic이 아닌 Linear로 증가하게 됩니다.

```

2 * batch_size * N_Layers * N_head * head_dim * sequence_length * 1

```

이로 인해서, 문장 생성시 FLOPS의 최적화로 Vanilla Transformer보다 빠른 Inference를 나타낼 수 있습니다.

하지만, 이로 인한 또 다른 문제가 발생하는데, 저장된 Cache의 크기는 문장의 길이에 따라 엄청난 양의 데이터를 가지고 있게 됩니다. LLM을 사용하는데 있어서 GPU연산을 통해 가속화를 하지만, GPU의 용량은 정해져 있어, 결국에는 GPU 메모리에 Cache를 담을수 없게 됩니다.

이 문제를 해결하기 위해 다양한 방법이 제시되었는데, Llama에서 사용한 방법은 GQA (Grouped Query Attention)을 사용하였습니다. (http://arxiv.org/abs/2305.13245)

Llama의 모델의 SelfAttention 클래스에 있는 (self.n_kv_heads) Argument는 이 GQA를 적용한 부분입니다. 

추가적인 방법은 Quantization 을 통해서 모델의 파라미터의 용량을 줄이는 방법도 많이 사용됩니다.

이상으로 Llama 등 LLM의 서비스화를 위해 연구된 KV-Cache에 대해 소개드렸습니다.

감사합니다.