Cara kerja penyangga

Bayangkan Anda sedang berobrol dengan seseorang, dan alih-alih mengingat apa yang telah mereka katakan sebelumnya, mereka harus memulai dari awal setiap kali Anda merespons. Hal ini pasti terasa lambat dan tidak efisien, bukan?

Anda dapat memperluas analogi ini ke model transformer. Generasi model autoregresif dapat lambat karena membuat prediksi satu token pada satu waktu. Setiap prediksi baru bergantung pada seluruh konteks sebelumnya.

Untuk memprediksi token ke-1000, model memerlukan informasi dari 999 token sebelumnya. Informasi ini direpresentasikan sebagai perkalian matriks di antara representasi token.

Untuk memprediksi token ke-1001, Anda membutuhkan informasi yang sama dari 999 token sebelumnya ditambah informasi apa pun dari token ke-1000. Ini adalah banyak perkalian matriks yang harus dihitung ulang dan ulang oleh model untuk setiap token!

Cache key-value (KV) menghilangkan inefisiensi ini dengan menyimpan pasangan kv yang diturunkan dari lapisan perhatian token yang telah diproses sebelumnya. Pasangan kv yang disimpan diambil dari cache dan digunakan kembali untuk token berikutnya, sehingga menghindari kebutuhan untuk menghitung ulang.

Caching hanya harus digunakan untuk inference. Penggunaannya selama pelatihan dapat menyebabkan error yang tidak terduga.

Untuk lebih memahami bagaimana dan mengapa caching bekerja, mari kita lihat lebih dekat struktur matriks perhatian.

Matriks perhatian

Perhatian dot-product berskala dihitung seperti yang ditunjukkan di bawah ini untuk batch berukuran b, jumlah kepala perhatian h, panjang urutan sejauh ini T, dan dimensi per kepala perhatian d_head. $$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$$

Matriks query (Q), key (K), dan value (V) adalah proyeksi dari input embeddings dengan bentuk (b, h, T, d_head).

Untuk causal attention, mask mencegah model memperhatikan token di masa depan. Setelah sebuah token diproses, representasinya tidak pernah berubah terhadap token di masa depan, yang berarti $K_{\text{past}}$ dan $V_{\text{past}}$ dapat di-cache dan digunakan kembali untuk menghitung representasi token terakhir. $$\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}])$$

Saat inferensi, Anda hanya perlu query dari token terakhir untuk menghitung representasi $x_t$ yang memprediksi token berikutnya $ t+1 $. Pada setiap langkah, vektor key dan value baru disimpan di cache dan ditambahkan ke key dan value sebelumnya. $$K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$$

Attention dihitung secara independen di setiap lapisan model, dan caching dilakukan berdasarkan lapisan.

Lihat tabel di bawah ini untuk membandingkan bagaimana caching meningkatkan efisiensi.

Kelas Cache

Antarmuka KV cache dasar menerima tensor key dan value untuk token saat ini dan mengembalikan tensor K dan V yang diperbarui. Hal ini dikelola secara internal oleh metode forward model.

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

Saat Anda menggunakan kelas Cache dari Transformers, modul self-attention melakukan beberapa langkah penting untuk mengintegrasikan informasi masa lalu dan saat ini.

1. Modul perhatian menggabungkan pasangan kv saat ini dengan pasangan kv masa lalu yang disimpan dalam cache. Hal ini menghasilkan bobot perhatian dengan bentuk (new_tokens_length, past_kv_length + new_tokens_length). Pasangan kv saat ini dan masa lalu pada dasarnya digabungkan untuk menghitung skor perhatian, sehingga memastikan model menyadari konteks sebelumnya dan input saat ini.

2. Ketika metode forward dipanggil secara iteratif, sangat penting agar bentuk attention mask sesuai dengan panjang gabungan dari pasangan kv masa lalu dan saat ini. Attention mask harus memiliki bentuk (batch_size, past_kv_length + new_tokens_length). Hal ini biasanya ditangani secara internal di generate(), tetapi jika Anda ingin mengimplementasikan loop generasi Anda sendiri dengan Cache, perhatikan hal ini! Attention mask harus menyimpan nilai token masa lalu dan saat ini.

Implementasi penyimpanan cache

Cache disusun sebagai daftar lapisan, di mana setiap lapisan berisi cache kunci dan nilai. Cache kunci dan nilai adalah tensor dengan bentuk [batch_size, num_heads, seq_len, head_dim].

Layers dapat memiliki tipe yang berbeda (misalnya DynamicLayer, StaticLayer, StaticSlidingWindowLayer), yang sebagian besar mengubah cara panjang sekuens ditangani dan cara cache diperbarui.

Yang paling sederhana adalah DynamicLayer yang tumbuh seiring dengan diprosesnya lebih banyak token. Dimensi panjang sekuens (seq_len) bertambah dengan setiap token baru:

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)

Jenis layer lain seperti StaticLayer dan StaticSlidingWindowLayer memiliki panjang sekuens tetap yang diatur saat cache dibuat. Hal ini membuatnya kompatibel dengan torch.compile. Dalam kasus StaticSlidingWindowLayer, token yang sudah ada akan digeser keluar dari cache ketika token baru ditambahkan.

Contoh di bawah ini menunjukkan cara membuat loop generasi dengan DynamicCache. Seperti yang dibahas, attention mask adalah konkatenasi dari nilai token masa lalu dan saat ini.

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