transformers/main/ko/cache_explanation.md

캐싱 작동 방식

누군가와 대화하고 있다고 상상해 보세요. 상대방이 이전에 한 말을 기억하지 못하고, 당신이 응답할 때마다 처음부터 다시 시작해야 한다면 어떨까요? 이는 느리고 비효율적일 것입니다.

이 비유를 transformer 모델에도 확장할 수 있습니다. 자기회귀 모델 생성은 한 번에 하나의 토큰씩 예측을 수행하기 때문에 느릴 수 있습니다. 각 새로운 예측은 이전의 모든 컨텍스트에 의존합니다.

1000번째 토큰을 예측하려면 모델은 이전 999개 토큰의 정보가 필요합니다. 이 정보는 토큰 표현에 걸친 행렬 곱으로 표현됩니다.

1001번째 토큰을 예측하려면 1000번째 토큰의 정보 외에도 이전 999개 토큰의 동일한 정보가 필요합니다. 이는 모델이 각 토큰마다 반복적으로 계산해야 하는 막대한 양의 행렬 곱셈입니다!

key-value (KV) cache는 이전에 처리된 토큰의 attention layer에서 파생된 kv 쌍을 저장하여 이러한 비효율성을 제거합니다. 저장된 kv 쌍은 캐시에서 검색되어 후속 토큰에 재사용되며, 이를 통해 재계산할 필요가 없습니다.

캐싱은 추론에만 사용해야 합니다. 학습 중에 활성화하면 예상치 못한 오류가 발생할 수 있습니다.

캐싱이 어떻게 그리고 왜 작동하는지 더 잘 이해하기 위해, 어텐션 행렬의 구조를 자세히 살펴보겠습니다.

Attention mat

b 크기의 배치, h개의 attention head, 현재까지의 시퀀스 길이 T, attention head당 차원 d_head에 대해 scaled dot-product attention은 아래에 표시된 대로 계산됩니다.

$$\text{Attention}(Q,K,V)=\text{softmax}(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_{\text{head}}}}\times \text{mask})V\backslash text\{Attention\}(Q,\; K,\; V)\; =\; \backslash text\{softmax\}\backslash left(\; \backslash frac\{Q\; K^\backslash top\}\{\backslash sqrt\{d\_\{\backslash text\{head\}\}\}\}\; \backslash times\; \backslash text\{mask\}\; \backslash right)\; V$$

쿼리(Q), 키(K), 값(V) 행렬은 (b, h, T, d_head) 모양의 입력 임베딩에서 투사된 것입니다.

causal attention의 경우, mask는 모델이 미래 토큰에 attend하는 것을 방지합니다. 한 토큰이 처리된 후에는 미래 토큰에 대해 그 표현이 더 이상 변하지 않으므로, $ K_{\text{past}} $와 $ V_{\text{past}} $를 캐시하여 마지막 토큰의 표현을 계산할 때 재사용할 수 있습니다.

$$\text{Attention}(q_t,[\underset{\text{캐시됨}}{\underbrace{k_1,k_2,\dots ,k_{t-1}}},k_t],[\underset{\text{캐시됨}}{\underbrace{v_1,v_2,\dots ,v_{t-1}}},v_t])\backslash text\{Attention\}(q\_t,\; [\backslash underbrace\{k\_1,\; k\_2,\; \backslash dots,\; k\_\{t-1\}\}\_\{\backslash text\{캐시됨\}\},\; k\_\{t\}],\; [\backslash underbrace\{v\_1,\; v\_2,\; \backslash dots,\; v\_\{t-1\}\}\_\{\backslash text\{캐시됨\}\},\; v\_\{t\}])$$

추론 시에는 다음 토큰 $ t+1 $을 예측하는 표현 $ x_t $를 계산하기 위해 마지막 토큰의 query만 필요합니다. 각 단계에서 새로운 key와 value 벡터는 캐시에 저장되고, 과거의 key와 value에 추가됩니다.

$$K_{\text{cache}}\leftarrow \text{concat}(K_{\text{past}},k_t),V_{\text{cache}}\leftarrow \text{concat}(V_{\text{past}},v_t)K\_\{\backslash text\{cache\}\}\; \backslash leftarrow\; \backslash text\{concat\}(K\_\{\backslash text\{past\}\},\; k\_t),\; \backslash quad\; V\_\{\backslash text\{cache\}\}\; \backslash leftarrow\; \backslash text\{concat\}(V\_\{\backslash text\{past\}\},\; v\_t)$$

Attention은 모델의 각 레이어에서 독립적으로 계산되며, 캐싱은 레이어별로 수행됩니다.

캐싱이 효율성을 어떻게 향상시키는지 비교하려면 아래 표를 참조하세요.

캐싱 없이	캐싱을 사용하여
각 단계에서 이전의 모든 K와 V를 다시 계산합니다	각 단계에서 현재 K와 V만 계산
각 단계의 attention 비용은 시퀀스 길이에 대해 2차입니다	각 단계의 attention 비용은 시퀀스 길이에 대해 선형입니다(메모리는 선형으로 증가하지만, 토큰당 계산량은 낮게 유지됩니다).

Cache 클래스

기본적인 KV 캐시 인터페이스는 현재 토큰의 key와 value 텐서를 받아 업데이트된 K와 V 텐서를 반환합니다. 이는 모델의 forward 메서드에 의해 내부적으로 관리됩니다.

new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)

Transformers의 Cache 클래스를 사용할 때, self-attention 모듈은 과거와 현재 정보를 통합하기 위해 몇 가지 핵심 단계를 수행합니다.

1. Attention 모듈은 캐시에 저장된 과거 kv 쌍과 현재 kv 쌍을 연결(concatenate)합니다. 이를 통해 (new_tokens_length, past_kv_length + new_tokens_length) 형태의 attention weights가 생성됩니다. 현재와 과거의 kv 쌍은 기본적으로 결합되어 attention scores를 계산하며, 이를 통해 모델이 이전 컨텍스트와 현재 입력을 인식할 수 있도록 보장합니다.

2. forward 메서드가 반복적으로 호출될 때, 어텐션 마스크의 형태가 과거와 현재 kv 쌍의 결합된 길이와 일치하는 것이 중요합니다. 어텐션 마스크의 형태는 (batch_size, past_kv_length + new_tokens_length)여야 합니다. 이는 일반적으로 generate() 내부에서 처리되지만, Cache를 사용하여 자체 생성 루프를 구현하려는 경우 이 점을 반드시 유의하세요. 어텐션 마스크에는 과거와 현재 토큰 값이 포함되어야 합니다.

캐시 저장 구현

Caches는 레이어의 리스트로 구성되며, 각 레이어에는 key와 value 캐시가 포함됩니다. key와 value 캐시는 [batch_size, num_heads, seq_len, head_dim] 형태의 텐서입니다.

Layer는 서로 다른 유형(예: DynamicLayer, StaticLayer, StaticSlidingWindowLayer)을 가질 수 있으며, 이는 주로 시퀀스 길이를 처리하는 방식과 캐시를 업데이트하는 방식을 변경합니다.

가장 간단한 방법은 토큰이 처리될수록 확장되는 DynamicLayer입니다. 시퀀스 길이 차원(seq_len)은 각 새 토큰마다 증가합니다:

cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)

StaticLayer와 StaticSlidingWindowLayer와 같은 다른 레이어 유형은 캐시가 생성될 때 설정되는 고정된 시퀀스 길이를 가집니다. 이로 인해 torch.compile과 호환됩니다. StaticSlidingWindowLayer의 경우, 새로운 토큰이 추가되면 기존 토큰이 캐시에서 밀려나갑니다.

아래 예제는 DynamicCache를 사용하여 생성 루프를 만드는 방법을 보여줍니다. 앞서 논의했듯이, 어텐션 마스크는 과거와 현재 토큰 값의 연결입니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    # Greedily sample one next token
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    # Prepare inputs for the next generation step by leaving unprocessed tokens, in our case we have only one new token
    # and expanding attn mask for the new token, as explained above
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"