KV Caching được giải thích- Tối ưu hóa hiệu quả suy luận Transformer
KV Caching được giải thích- Tối ưu hóa hiệu quả suy luận Transformer
- 6 min read
KV Caching Explained: Optimizing Transformer Inference Efficiency
Khi các mô hình AI tạo văn bản, chúng thường lặp đi lặp lại nhiều phép tính giống nhau, điều này có thể làm chậm quá trình. Key-Value caching (KV caching) là một kỹ thuật giúp tăng tốc quá trình này bằng cách ghi nhớ các thông tin quan trọng từ các bước trước. Thay vì tính toán lại mọi thứ từ đầu, mô hình tái sử dụng những gì đã tính toán, giúp việc tạo văn bản nhanh hơn và hiệu quả hơn.
Trong bài viết này, chúng ta sẽ tìm hiểu KV caching một cách dễ hiểu, giải thích tại sao nó hữu ích và cách nó giúp các mô hình AI hoạt động nhanh hơn.
Prerequisites (Yêu cầu tiên quyết)
Để hiểu đầy đủ nội dung, người đọc nên quen thuộc với:
- Kiến trúc Transformer: Hiểu biết về các thành phần như cơ chế attention.
- Mô hình tự hồi quy (Autoregressive Modeling): Hiểu cách các mô hình như GPT tạo ra chuỗi.
- Kiến thức cơ bản về Đại số tuyến tính: Các khái niệm như phép nhân ma trận và chuyển vị, rất cần thiết để hiểu các phép tính attention.
Bạn có thể tìm hiểu bài viết BLOG này để nắm vững các kiến thức cần thiết cho bài viết này.
Standard Inference and the Rise of KV Caching (Suy luận tiêu chuẩn và sự trỗi dậy của KV Caching)
Khi một mô hình tạo văn bản, nó xem xét tất cả các token trước đó để dự đoán token tiếp theo. Thông thường, nó sẽ lặp lại cùng các phép tính cho mỗi token mới, điều này có thể làm chậm quá trình.
How Does KV Caching Work? (KV Caching hoạt động như thế nào?)
Step-by-Step Process (Quy trình từng bước)
- Lần tạo đầu tiên: Khi mô hình xử lý đầu vào đầu tiên, nó sẽ tính toán và lưu trữ các key và value của nó vào bộ nhớ cache.
- Các từ tiếp theo: Với mỗi từ mới, mô hình truy xuất các key và value đã lưu trữ và thêm các key/value mới thay vì bắt đầu lại từ đầu.
- Tính toán Attention hiệu quả: Tính toán attention bằng cách sử dụng các key và value đã lưu trữ cùng với query mới để tính toán đầu ra.
- Cập nhật đầu vào: Thêm token mới được tạo vào đầu vào và quay lại bước 2 cho đến khi quá trình tạo hoàn tất.
Quá trình này được minh họa dưới đây:
Token 1: [K1, V1] ➔ Cache: [K1, V1] Token 2: [K2, V2] ➔ Cache: [K1, K2], [V1, V2] … Token n: [Kn, Vn] ➔ Cache: [K1, K2, …, Kn], [V1, V2, …, Vn]
Trong bảng trên, chúng ta đã sử dụng d_k = 5 để trực quan hóa tốt hơn, lưu ý rằng con số này có thể lớn hơn đáng kể so với những gì chúng ta đã trình bày.
Comparison: KV Caching vs. Standard Inference (So sánh: KV Caching và Suy luận tiêu chuẩn)
Dưới đây là cách KV caching so với các phương pháp tạo thông thường:
| Feature (Đặc điểm) | Standard Inference (Suy luận tiêu chuẩn) | KV Caching |
|---|---|---|
| Computation per Word | Mô hình lặp lại các phép tính giống nhau cho mỗi từ. | Mô hình tái sử dụng các phép tính trước đó để có kết quả nhanh hơn. |
| Memory Usage | Sử dụng ít bộ nhớ hơn ở mỗi bước, nhưng bộ nhớ tăng theo văn bản dài. | Sử dụng thêm bộ nhớ để lưu trữ thông tin quá khứ, nhưng vẫn giữ hiệu quả. |
| Speed | Chậm dần khi văn bản dài hơn do lặp lại công việc. | Vẫn nhanh ngay cả với văn bản dài hơn bằng cách tránh lặp lại công việc. |
| Efficiency | Chi phí tính toán cao và thời gian phản hồi chậm. | Nhanh hơn và hiệu quả hơn vì mô hình ghi nhớ công việc đã thực hiện. |
| Handling Long Texts | Gặp khó khăn với văn bản dài do các phép tính lặp lại. | Hoàn hảo cho văn bản dài vì nó ghi nhớ các bước trước đó. |
KV caching tạo ra sự khác biệt lớn về tốc độ và hiệu quả, đặc biệt đối với văn bản dài. Bằng cách lưu và tái sử dụng các phép tính trước đó, nó loại bỏ nhu cầu bắt đầu lại mỗi lần, giúp nhanh hơn nhiều so với cách tạo văn bản thông thường.
Practical Implementation (Triển khai thực tế)
Đây là một ví dụ đơn giản về việc triển khai KV caching trong PyTorch:
python
Pseudocode cho KV Caching trong PyTorch
class KVCache: def init(self): self.cache = {“key”: None, “value”: None}
def update(self, key, value):
if self.cache["key"] is None:
self.cache["key"] = key
self.cache["value"] = value
else:
self.cache["key"] = torch.cat([self.cache["key"], key], dim=1)
self.cache["value"] = torch.cat([self.cache["value"], value], dim=1)
def get_cache(self):
return self.cache
Khi sử dụng thư viện transformers, hành vi này được kích hoạt mặc định thông qua tham số use_cache. Bạn cũng có thể truy cập nhiều phương pháp caching thông qua tham số cache_implementation. Dưới đây là một đoạn mã tối thiểu:
python from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(‘HuggingFaceTB/SmolLM2-1.7B’) model = AutoModelForCausalLM.from_pretrained(‘HuggingFaceTB/SmolLM2-1.7B’).cuda()
tokens = tokenizer.encode(“The red cat was”, return_tensors=“pt”).cuda() output = model.generate( tokens, max_new_tokens=300, use_cache = True # mặc định là True ) output_text = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
Chúng tôi đã kiểm tra hiệu suất mã trên với/không có kv caching trên GPU T4 và nhận được kết quả sau:
Conclusion (Kết luận)
KV caching là một kỹ thuật đơn giản nhưng mạnh mẽ giúp các mô hình AI tạo văn bản nhanh hơn và hiệu quả hơn. Bằng cách ghi nhớ các phép tính trước đó thay vì lặp lại chúng, nó giảm thời gian và công sức cần thiết để dự đoán các từ mới. Mặc dù nó đòi hỏi thêm bộ nhớ, phương pháp này đặc biệt hữu ích cho các cuộc hội thoại dài, đảm bảo việc tạo ra kết quả nhanh chóng và hiệu quả.
Hiểu về KV caching có thể giúp các nhà phát triển và những người đam mê AI xây dựng các mô hình ngôn ngữ nhanh hơn, thông minh hơn và có khả năng mở rộng hơn cho các ứng dụng trong thế giới thực.
Tôi xin gửi lời cảm ơn chân thành nhất đến Aritra Roy Gosthipaty 🤗 vì sự hỗ trợ vô giá, phản hồi và sự cống hiến của anh ấy trong việc phát triển bài viết blog này.
References & Further Reading (Tài liệu tham khảo & Đọc thêm)
- Transformers KV Caching Explained
- Transformers Key-Value Caching Explained
- Mastering LLM Techniques: Inference Optimization
- Hugging Face Documentation - KV Caching in Transformers
Community (Cộng đồng)
Link bài viết gốc
- Tags:
- Ai
- 30 Tháng 1
- Huggingface.co