Cách hoạt động của bộ nhớ đệm

Hãy tưởng tượng bạn đang trò chuyện với ai đó, và thay vì nhớ những gì họ đã nói trước đó, họ lại phải bắt đầu lại từ đầu mỗi khi bạn phản hồi. Điều này sẽ chậm chạp và kém hiệu quả, đúng không?

Bạn có thể mở rộng phép so sánh này cho các mô hình transformer. Việc tạo ra nội dung của mô hình tự hồi quy có thể chậm vì nó đưa ra dự đoán từng token một. Mỗi dự đoán mới đều phụ thuộc vào toàn bộ ngữ cảnh trước đó.

Để dự đoán token thứ 1000, mô hình cần thông tin từ 999 token trước đó. Thông tin này được biểu diễn dưới dạng các phép nhân ma trận trên các biểu diễn token.

Để dự đoán token thứ 1001, bạn cần cùng một lượng thông tin từ 999 token trước đó, cộng thêm bất kỳ thông tin nào từ token thứ 1000. Đây là một lượng lớn phép nhân ma trận mà mô hình phải tính toán lặp đi lặp lại cho từng token!

Bộ nhớ đệm khóa-giá trị (KV cache) loại bỏ sự kém hiệu quả này bằng cách lưu trữ các cặp kv được suy ra từ các lớp chú ý của các token đã được xử lý trước đó. Các cặp kv đã lưu sẽ được truy xuất từ bộ nhớ đệm và tái sử dụng cho các token tiếp theo, tránh việc phải tính toán lại.

Caching chỉ nên được sử dụng cho inference. Nó có thể gây ra các lỗi không mong muốn nếu được bật trong quá trình huấn luyện.

Để hiểu rõ hơn về cách thức và lý do hoạt động của bộ nhớ đệm, hãy cùng xem xét kỹ hơn cấu trúc của các ma trận chú ý.

Ma trận chú ý

Cơ chế scaled dot-product attention được tính toán như bên dưới cho một batch có kích thước b, số đầu attention h, độ dài chuỗi cho đến hiện tại T, và số chiều cho mỗi đầu attention d_head. $\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$

Các ma trận truy vấn (Q), khóa (K) và giá trị (V) là các phép chiếu từ các nhúng đầu vào có hình dạng (b, h, T, d_head).

Đối với causal attention, mask ngăn mô hình chú ý đến các token trong tương lai. Một khi một token được xử lý, biểu diễn của nó không bao giờ thay đổi so với các token trong tương lai, điều này có nghĩa là $K_{\text{past}}$ và $V_{\text{past}}$ có thể được lưu cache và tái sử dụng để tính toán biểu diễn của token cuối cùng. $\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}])$

Khi suy luận, bạn chỉ cần query của token cuối cùng để tính toán biểu diễn $x_t$ dự đoán token tiếp theo $ t+1 $. Ở mỗi bước, các vector key và value mới được lưu trữ trong cache và thêm vào các key và value trong quá khứ. $K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$

Attention được tính toán độc lập ở mỗi lớp của mô hình, và việc lưu cache được thực hiện theo từng lớp.

Tham khảo bảng dưới đây để so sánh cách caching cải thiện hiệu suất.

không có bộ nhớ đệm	với bộ nhớ đệm
cho mỗi bước, tính lại tất cả `K` và `V` trước đó	cho mỗi bước, chỉ tính toán `K` và `V` hiện tại
chi phí attention cho mỗi bước là bậc hai theo độ dài chuỗi	chi phí attention cho mỗi bước là tuyến tính theo độ dài chuỗi (bộ nhớ tăng tuyến tính, nhưng tính toán/token vẫn thấp)

Lớp Cache

Giao diện KV cache cơ bản nhận tensor key và value cho token hiện tại và trả về các tensor K và V đã cập nhật. Điều này được quản lý bên trong bởi phương thức forward của mô hình.

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

Khi bạn sử dụng lớp Cache của Transformers, module self-attention thực hiện một số bước quan trọng để tích hợp thông tin quá khứ và hiện tại.

Module chú ý ghép nối các cặp kv hiện tại với các cặp kv trước đó được lưu trong bộ nhớ đệm. Điều này tạo ra các trọng số chú ý có hình dạng (new_tokens_length, past_kv_length + new_tokens_length). Các cặp kv hiện tại và trước đó về cơ bản được kết hợp để tính toán điểm chú ý, đảm bảo mô hình nhận thức được ngữ cảnh trước đó và đầu vào hiện tại.
Khi phương thức forward được gọi lặp lại, điều quan trọng là hình dạng của attention mask phải khớp với tổng độ dài của các cặp kv trong quá khứ và hiện tại. Attention mask nên có hình dạng (batch_size, past_kv_length + new_tokens_length). Điều này thường được xử lý bên trong generate(), nhưng nếu bạn muốn tự triển khai vòng lặp sinh với Cache, hãy lưu ý điều này! Attention mask nên chứa các giá trị token trong quá khứ và hiện tại.

Thực hiện lưu trữ bộ nhớ đệm

Caches được cấu trúc dưới dạng một danh sách các lớp, trong đó mỗi lớp chứa một bộ nhớ đệm key và value. Các bộ nhớ đệm key và value là các tensor có hình dạng [batch_size, num_heads, seq_len, head_dim].

Các lớp có thể thuộc các kiểu khác nhau (ví dụ: DynamicLayer, StaticLayer, StaticSlidingWindowLayer), chủ yếu thay đổi cách xử lý độ dài chuỗi và cách cập nhật bộ nhớ đệm.

Cách đơn giản nhất là một DynamicLayer tăng kích thước khi có thêm token được xử lý. Chiều dài chuỗi (seq_len) tăng lên với mỗi token mới:

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)

Các loại layer khác như StaticLayer và StaticSlidingWindowLayer có độ dài chuỗi cố định được thiết lập khi cache được tạo. Điều này giúp chúng tương thích với torch.compile. Trong trường hợp của StaticSlidingWindowLayer, các token hiện có sẽ bị dịch chuyển ra khỏi cache khi một token mới được thêm vào.

Ví dụ dưới đây minh họa cách tạo vòng lặp sinh với DynamicCache. Như đã thảo luận, attention mask là sự nối tiếp của các giá trị token trong quá khứ và hiện tại.

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

Update on GitHub