Buckets:
| # 캐싱 작동 방식 | |
| 누군가와 대화하고 있다고 상상해 보세요. 상대방이 이전에 한 말을 기억하지 못하고, 당신이 응답할 때마다 처음부터 다시 시작해야 한다면 어떨까요? 이는 느리고 비효율적일 것입니다. | |
| 이 비유를 transformer 모델에도 확장할 수 있습니다. 자기회귀 모델 생성은 한 번에 하나의 토큰씩 예측을 수행하기 때문에 느릴 수 있습니다. 각 새로운 예측은 이전의 모든 컨텍스트에 의존합니다. | |
| 1000번째 토큰을 예측하려면 모델은 이전 999개 토큰의 정보가 필요합니다. 이 정보는 토큰 표현에 걸친 행렬 곱으로 표현됩니다. | |
| 1001번째 토큰을 예측하려면 1000번째 토큰의 정보 외에도 이전 999개 토큰의 동일한 정보가 필요합니다. 이는 모델이 각 토큰마다 반복적으로 계산해야 하는 막대한 양의 행렬 곱셈입니다! | |
| key-value (KV) cache는 이전에 처리된 토큰의 attention layer에서 파생된 kv 쌍을 저장하여 이러한 비효율성을 제거합니다. 저장된 kv 쌍은 캐시에서 검색되어 후속 토큰에 재사용되며, 이를 통해 재계산할 필요가 없습니다. | |
| > [!WARNING] | |
| > 캐싱은 **추론**에만 사용해야 합니다. 학습 중에 활성화하면 예상치 못한 오류가 발생할 수 있습니다. | |
| 캐싱이 어떻게 그리고 왜 작동하는지 더 잘 이해하기 위해, 어텐션 행렬의 구조를 자세히 살펴보겠습니다. | |
| ## Attention mat | |
| `b` 크기의 배치, `h`개의 attention head, 현재까지의 시퀀스 길이 `T`, attention head당 차원 `d_head`에 대해 **scaled dot-product attention**은 아래에 표시된 대로 계산됩니다. | |
| $$ \text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V $$ | |
| 쿼리(`Q`), 키(`K`), 값(`V`) 행렬은 `(b, h, T, d_head)` 모양의 입력 임베딩에서 투사된 것입니다. | |
| causal attention의 경우, mask는 모델이 미래 토큰에 attend하는 것을 방지합니다. 한 토큰이 처리된 후에는 미래 토큰에 대해 그 표현이 더 이상 변하지 않으므로, $ K_{\text{past}} $와 $ V_{\text{past}} $를 캐시하여 마지막 토큰의 표현을 계산할 때 재사용할 수 있습니다. | |
| $$ \text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{캐시됨}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{캐시됨}}, v_{t}]) $$ | |
| 추론 시에는 다음 토큰 $ t+1 $을 예측하는 표현 $ x_t $를 계산하기 위해 마지막 토큰의 query만 필요합니다. 각 단계에서 새로운 key와 value 벡터는 캐시에 **저장**되고, 과거의 key와 value에 **추가**됩니다. | |
| $$ K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t) $$ | |
| Attention은 모델의 각 레이어에서 독립적으로 계산되며, 캐싱은 레이어별로 수행됩니다. | |
| 캐싱이 효율성을 어떻게 향상시키는지 비교하려면 아래 표를 참조하세요. | |
| | 캐싱 없이 | 캐싱을 사용하여 | | |
| |---|---| | |
| | 각 단계에서 이전의 모든 `K`와 `V`를 다시 계산합니다 | 각 단계에서 현재 `K`와 `V`만 계산 | | |
| | 각 단계의 attention 비용은 시퀀스 길이에 대해 **2차**입니다 | 각 단계의 attention 비용은 시퀀스 길이에 대해 **선형**입니다(메모리는 선형으로 증가하지만, 토큰당 계산량은 낮게 유지됩니다). | | |
| ## Cache 클래스 | |
| 기본적인 KV 캐시 인터페이스는 현재 토큰의 key와 value 텐서를 받아 업데이트된 `K`와 `V` 텐서를 반환합니다. 이는 모델의 `forward` 메서드에 의해 내부적으로 관리됩니다. | |
| ```py | |
| 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](/docs/transformers/main/ko/internal/generation_utils#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()](/docs/transformers/main/ko/main_classes/text_generation#transformers.GenerationMixin.generate) 내부에서 처리되지만, [Cache](/docs/transformers/main/ko/internal/generation_utils#transformers.Cache)를 사용하여 자체 생성 루프를 구현하려는 경우 이 점을 반드시 유의하세요. 어텐션 마스크에는 과거와 현재 토큰 값이 포함되어야 합니다. | |
| ## 캐시 저장 구현 | |
| Caches는 레이어의 리스트로 구성되며, 각 레이어에는 key와 value 캐시가 포함됩니다. key와 value 캐시는 `[batch_size, num_heads, seq_len, head_dim]` 형태의 텐서입니다. | |
| Layer는 서로 다른 유형(예: `DynamicLayer`, `StaticLayer`, `StaticSlidingWindowLayer`)을 가질 수 있으며, 이는 주로 시퀀스 길이를 처리하는 방식과 캐시를 업데이트하는 방식을 변경합니다. | |
| 가장 간단한 방법은 토큰이 처리될수록 확장되는 `DynamicLayer`입니다. 시퀀스 길이 차원(`seq_len`)은 각 새 토큰마다 증가합니다: | |
| ```py | |
| 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](/docs/transformers/main/ko/internal/generation_utils#transformers.DynamicCache)를 사용하여 생성 루프를 만드는 방법을 보여줍니다. 앞서 논의했듯이, 어텐션 마스크는 과거와 현재 토큰 값의 연결입니다. | |
| ```py | |
| 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," | |
| ``` | |
Xet Storage Details
- Size:
- 8.23 kB
- Xet hash:
- 4774414b47db46de86c0a0efc749290968d235a1dc6e1bd3ab8fbae2e7acab05
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.