Bộ nhớ cache KV từ đầu trong nanoVLM
KV Cache từ đầu trong nanoVLM
- 4 min read
KV Cache từ đầu trong nanoVLM
Bài viết gốc: https://huggingface.co/blog/kv-cache
TL;DR
Chúng tôi đã triển khai KV Caching từ đầu trong kho lưu trữ nanoVLM của mình (một cơ sở mã nhỏ để huấn luyện Mô hình Ngôn ngữ Tầm nhìn của riêng bạn bằng PyTorch thuần túy). Điều này cho chúng tôi mức tăng tốc 38% trong quá trình tạo. Trong bài đăng trên blog này, chúng tôi đề cập đến KV Caching và tất cả kinh nghiệm của chúng tôi trong khi triển khai nó. Các bài học kinh nghiệm là tổng quát và có thể được áp dụng cho tất cả các thế hệ mô hình ngôn ngữ tự hồi quy. Triển khai từ đầu trên một cơ sở mã nhỏ là một trải nghiệm học tập tuyệt vời, hãy tham gia cùng chúng tôi!
Giới thiệu
Các mô hình ngôn ngữ tự hồi quy tạo văn bản bằng cách lấy mẫu mỗi lần một token. Trong quá trình suy luận, mô hình xử lý một chuỗi đầu vào nhất định, dự đoán token tiếp theo, nối nó vào chuỗi và lặp lại quá trình này cho đến khi đạt được một tiêu chí dừng nào đó:
Việc tạo từng bước này vốn dĩ là tuần tự:
- Để tạo mã thông báo ti+1 t_{i+1} ti+1, mô hình phải xem xét toàn bộ chuỗi từ t0 t_0 t0 đến ti t_i ti. Từ phiên bản đầu tiên trong ví dụ trên ti+1 t_{i+1} ti+1 sẽ là
the, trong khi tất cả các mã thông báo trước đó t0 t_0 t0 đến ti t_i ti sẽ là[What, is, in,]. - Mặc dù các máy biến áp song song bên trong, nhưng mỗi dự đoán mới yêu cầu một đường chuyền chuyển tiếp đầy đủ thông qua tất cả các lớp máy biến áp, điều này làm phát sinh một bộ nhớ/tính toán bậc hai theo độ dài chuỗi.
Sự lặp lại này cũng dẫn đến dư thừa tính toán. Trong bài đăng này, chúng ta sẽ khám phá KV Caching, một kỹ thuật tối ưu hóa giúp giảm thiểu sự kém hiệu quả này.
Mục lục:
- Xem lại Kiến trúc Transformer
- Dư thừa xuất hiện ở đâu
- KV Caching khắc phục nó như thế nào
- KV Caching trong nanoVLM: Từ lý thuyết đến thực hành
- Tóm tắt: Tại sao KV Caching lại quan trọng
Xem lại Kiến trúc Transformer
Trước khi đi sâu vào bộ nhớ đệm, hãy xem lại cách hoạt động của sự chú ý trong các mô hình máy biến áp. Một mô hình ngôn ngữ Transformer bao gồm các lớp xếp chồng lên nhau, mỗi lớp bao gồm:
- Tự chú ý nhiều đầu
- Mạng chuyển tiếp (MLP)
- Kết nối còn sót lại và chuẩn hóa lớp
Để hiểu KV Caching giúp ích ở đâu, chúng ta tập trung vào cơ chế tự chú ý, cụ thể là trong một đầu chú ý duy nhất.
Hãy xem qua một triển khai PyTorch đơn giản để hình dung các tính toán chính.
import torch
input_seq_length = 5
dim_model = 10
input_ids_emb = torch.randn(input_seq_length, dim_model)
W_q = torch.randn(dim_model, dim_model)
W_k = torch.randn(dim_model, dim_model)
W_v = torch.randn(dim_model, dim_model)
Q = input_ids_emb @ W_q
K = input_ids_emb @ W_k
V = input_ids_emb @ W_v
Tính toán Tự chú ý
Đối với một chuỗi các nhúng đầu vào T T T được biểu thị là X∈RT×D X \in \mathbb{R}^{T \times D} X∈RT×D, tự chú ý được tính như sau:
- Q=XWQ Q = XW_Q Q=XWQ, với WQ∈RD×Dq W_Q \in \mathbb{R}^{D \times D_q} WQ∈RD×Dq
- K=XWK K = XW_K K=XWK, với WK∈RD×Dk W_K \in \mathbb{R}^{D \times D_k} WK∈RD×Dk
- V=XWV V = XW_V V=XWV, với WV∈RD×Dv W_V \in \mathbb{R}^{D \times D_v} WV∈RD×Dv
- Mặt nạ nhân quả M M M để ngăn chặn việc truy cập mã thông báo trong tương lai
Đầu ra cuối cùng là:
Attention(X;Q,K,V)=softmax(QK⊤⋅Mdk)V \text{Attention}(X; Q, K, V) = \text{softmax}\left( \frac{QK^\top \cdot M}{\sqrt{d_k}} \right)V <span class=“strut” style=“height
Link bài viết gốc
- Tags:
- Ai
- June 4, 2025
- Huggingface.co