كيف يعمل التخزين المؤقت

تخيّل أنك في محادثة مع شخص ما، وبدلًا من أن يتذكر ما قاله سابقًا، عليه أن يبدأ من الصفر في كل مرة تردّ فيها. أليس هذا بطيئًا وغير فعّال؟

يمكنك توسيع هذه التشبيه إلى نماذج التحويل (Transformer). قد يكون توليد النماذج التوليدية الذاتية (Autoregressive) بطيئًا لأنه يتنبأ برمز واحد في كل مرة. كل تنبؤ جديد يعتمد على كامل السياق السابق.

لتنبؤ الرمز رقم 1000، يحتاج النموذج إلى معلومات من الرموز التسعمئة والتسعة والتسعين السابقة. يتم تمثيل هذه المعلومات على شكل ضربات مصفوفية عبر تمثيلات الرموز.

لتنبؤ بالرمز رقم 1001، تحتاج إلى نفس المعلومات من الرموز التسعمئة والتسعة السابقة بالإضافة إلى أي معلومات من الرمز رقم 1000. وهذا يعني الكثير من عمليات ضرب المصفوفات التي يجب على النموذج حسابها مرة بعد مرة لكل رمز!

تُزيل ذاكرة التخزين المؤقت للمفتاح والقيمة (KV) هذه الكفاءة المنخفضة عن طريق تخزين أزواج kv المشتقة من طبقات الانتباه للرموز التي تمت معالجتها سابقًا. يتم استرجاع أزواج kv المخزنة من ذاكرة التخزين المؤقت وإعادة استخدامها للرموز اللاحقة، مما يتجنب الحاجة إلى إعادة الحساب.

يجب استخدام التخزين المؤقت فقط لـالاستدلال. قد يؤدي تفعيله أثناء التدريب إلى أخطاء غير متوقعة.

لفهم أفضل لكيفية عمل التخزين المؤقت ولماذا يعمل، دعونا ننظر عن كثب في بنية مصفوفات الانتباه.

مصفوفات الانتباه

يُحسب انتباه ضرب النقطة المُقيَّم كما هو موضح أدناه لدُفعة حجمها b, عدد رؤوس الانتباه h, طول التسلسل حتى الآن T, والأبعاد لكل رأس انتباه d_head. $$\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).

في الانتباه السببي، تمنع القناع النموذج من الانتباه إلى الرموز المستقبلية. بمجرد معالجة رمز، لا تتغير تمثيلاته فيما يتعلق بالرموز المستقبلية، مما يعني أن $K_{\text{past}}$ و $V_{\text{past}}$ يمكن تخزينها مؤقتًا وإعادة استخدامها لحساب تمثيل الرمز الأخير. $$\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}])$$

عند الاستدلال، لا تحتاج سوى استعلام الرمز الأخير لحساب التمثيل $x_t$ الذي يتنبأ بالرمز التالي $ t+1 $. في كل خطوة، يتم تخزين متجهات المفاتيح والقيم الجديدة في ذاكرة التخزين المؤقت وإضافتها إلى المفاتيح والقيم السابقة. $$K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$$

يتم حساب الانتباه بشكل مستقل في كل طبقة من طبقات النموذج، وتُنفَّذ التخزين المؤقت على أساس كل طبقة على حدة.

راجع الجدول أدناه لمقارنة كيفية تحسين التخزين المؤقت للكفاءة.

فئة الكاش

تستقبل واجهة تخزين KV الأساسية متجهين للمفتاح والقيمة للرمز الحالي وتعيد متجهي 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)

عند استخدامك لفئة Cache في Transformers، يقوم وحدة الانتباه الذاتي (self-attention) بعدة خطوات حاسمة لدمج المعلومات السابقة والحالية.

1. يقوم وحدة الانتباه بدمج أزواج kv الحالية مع أزواج kv السابقة المخزنة في الكاش. وهذا يُنشئ أوزان انتباه بالشكل (new_tokens_length, past_kv_length + new_tokens_length). تُدمج أزواج kv الحالية والسابقة بشكل أساسي لحساب درجات الانتباه، مما يضمن أن يكون النموذج على دراية بالسياق السابق والمدخل الحالي.

2. عند استدعاء طريقة forward بشكل متكرر، من الضروري أن يتطابق شكل قناع الانتباه مع الطول المجمع للأزواج kv السابقة والحالية. يجب أن يكون شكل قناع الانتباه (batch_size, past_kv_length + new_tokens_length). عادةً ما يتم التعامل مع ذلك داخليًا في generate()، ولكن إذا كنت ترغب في تنفيذ حلقة توليد خاصة بك باستخدام Cache، فاحرص على مراعاة ذلك! يجب أن يحتوي قناع الانتباه على قيم الرموز السابقة والحالية.

تنفيذ تخزين الكاش

تُبنى الكاشات على شكل قائمة من الطبقات، حيث تحتوي كل طبقة على كاش للمفاتيح وكاش للقيم. وتكون كاشات المفاتيح والقيم من نوع tensors بأبعاد [batch_size, num_heads, seq_len, head_dim].

يمكن أن تكون الطبقات من أنواع مختلفة (مثل 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,"