Önbellekleme nasıl çalışır

Birisiyle sohbet ettiğinizi ve her yanıt verdiğinizde, o kişinin daha önce söylediklerini hatırlamak yerine her seferinde sıfırdan başlamak zorunda kaldığını hayal edin. Bu, yavaş ve verimsiz olurdu, değil mi?

Bu analojiyi transformer modellerine de genişletebilirsiniz. Otoregresif model üretimi, tahmini token token yaptığında yavaş olabilir. Her yeni tahmin, önceki tüm bağlama bağlıdır.

1000. jetonu tahmin etmek için model, önceki 999 jetondan bilgiye ihtiyaç duyar. Bu bilgi, jeton temsilleri üzerinden matris çarpımları olarak temsil edilir.

1001. tokeni tahmin etmek için, 1000. tokenin bilgilerine ek olarak, önceki 999 tokenin aynı bilgisine ihtiyacınız var. Bir modelin her token için tekrar tekrar hesaplaması gereken çok fazla matris çarpımı var!

Bir anahtar-değer (KV) önbelleği, daha önce işlenmiş jetonların dikkat katmanlarından türetilen kv çiftlerini saklayarak bu verimsizliği ortadan kaldırır. Saklanan kv çiftleri önbellekten alınır ve sonraki jetonlar için yeniden kullanılır; böylece yeniden hesaplama gereksinimi ortadan kalkar.

Önbellekleme yalnızca inference için kullanılmalıdır. Eğitim sırasında etkinleştirilirse beklenmeyen hatalara neden olabilir.

Önbellekleme işleminin nasıl ve neden çalıştığını daha iyi anlamak için dikkat matrislerinin yapısına daha yakından bakalım.

Dikkat matrisleri

Ölçekli nokta çarpımı dikkati şu şekilde hesaplanır: boyutu b, dikkat başı sayısı h, şu ana kadarki dizi uzunluğu T, ve dikkat başı başına boyut d_head. $$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$$

Sorgu (Q), anahtar (K) ve değer (V) matrisleri, (b, h, T, d_head) boyutundaki giriş gömme vektörlerinden yapılan projeksiyonlardır.

Nedensel dikkat (causal attention) için maske, modelin gelecekteki jetonlara dikkat etmesini engeller. Bir jeton işlendikten sonra, gelecekteki jetonlara göre temsili asla değişmez; bu da $K_{\text{past}}$ ve $V_{\text{past}}$ değerlerinin önbelleğe alınarak son jetonun temsilini hesaplamak için yeniden kullanılabilmesi anlamına gelir. $$\text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{cached}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{cached}}, v_{t}])$$

Çıkarım sırasında, temsilin hesaplanması için yalnızca son token'ın sorgusuna ihtiyaç duyulur $x_t$ ve bir sonraki token'ı $ t+1 $ tahmin eder. Her adımda, yeni anahtar ve değer vektörleri önbelleğe kaydedilir ve eklenir geçmiş anahtar ve değerlerin sonuna. $$K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$$

Attention, modelin her katmanında bağımsız olarak hesaplanır ve önbellekleme katman bazında yapılır.

Önbelleklemenin verimliliği nasıl iyileştirdiğini karşılaştırmak için aşağıdaki tabloya bakın.

Önbellek sınıfı

Temel bir KV önbellek arayüzü, mevcut token için bir anahtar ve değer tensörünü alır ve güncellenmiş K ile V tensörlerini döndürür. Bu, bir modelin forward metodu tarafından dahili olarak yönetilir.

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'ın Cache sınıfını kullandığınızda, öz-dikkat modülü geçmiş ve mevcut bilgileri birleştirmek için birkaç kritik adım gerçekleştirir.

1. Dikkat modülü, mevcut kv çiftlerini önbellekte saklanan geçmiş kv çiftleriyle birleştirir. Bu işlem, (new_tokens_length, past_kv_length + new_tokens_length) biçiminde dikkat ağırlıkları oluşturur. Mevcut ve geçmiş kv çiftleri, dikkat skorlarının hesaplanmasında temelde birleştirilir; bu da modelin önceki bağlamdan ve mevcut girdiden haberdar olmasını sağlar.

2. forward metodu iteratif olarak çağrıldığında, dikkat maskesinin (attention mask) şeklinin geçmiş ve mevcut kv çiftlerinin birleşik uzunluğuyla eşleşmesi hayati önem taşır. Dikkat maskesinin şekli (batch_size, past_kv_length + new_tokens_length) olmalıdır. Bu durum genellikle generate() içinde dahili olarak ele alınır, ancak Cache ile kendi üretim döngünüzü (generation loop) uygulamak istiyorsanız bunu göz önünde bulundurun! Dikkat maskesi, geçmiş ve mevcut token değerlerini içermelidir.

Önbellek depolama uygulaması

Önbellekler, bir katman listesinin yapısında düzenlenmiştir; her katman bir anahtar ve değer önbelleği içerir. Anahtar ve değer önbellekleri, [batch_size, num_heads, seq_len, head_dim] biçimine sahip tensörlerdir.

Katmanlar farklı türlerde olabilir (ör. DynamicLayer, StaticLayer, StaticSlidingWindowLayer); bu durum, büyük ölçüde dizinin uzunluğunun nasıl ele alındığını ve önbelleğin nasıl güncellendiğini değiştirir.

En basiti, işlenen token sayısı arttıkça büyüyen bir DynamicLayer'dir. Dizi uzunluğu boyutu (seq_len) her yeni token ile artar:

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 ve StaticSlidingWindowLayer gibi diğer katman türleri, önbellek oluşturulurken belirlenen sabit bir dizi uzunluğuna sahiptir. Bu, onları torch.compile ile uyumlu hale getirir. StaticSlidingWindowLayer durumunda, yeni bir token eklendiğinde mevcut tokenlar önbellekten dışarı kaydırılır.

Aşağıdaki örnek, DynamicCache ile bir üretim döngüsü oluşturmayı göstermektedir. Tartıldığı üzere, dikkat maskesi geçmiş ve mevcut jeton değerlerinin birleştirilmesidir.

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,"