LoopVQ: Nén bộ nhớ đệm KV của Transformer lặp lại xuống dưới một bit

Các Transformer lặp lại gọi đi gọi lại cùng một nhóm lớp, có thể bắt kịp các mô hình thông thường lớn gấp nhiều lần với ít tham số hơn, nhưng mỗi bước lặp đều phải lưu trữ bộ nhớ đệm KV riêng, khiến bộ nhớ đệm tăng tuyến tính theo số bước lặp và trở thành nút thắt cổ chai chính về bộ nhớ. Các tác giả nhận thấy bộ nhớ đệm KV có tính dư thừa cao trên trục lặp, trong khi độ nhạy cảm với lượng tử hóa của các bước lặp khác nhau lại rất khác biệt. Do đó, họ đề xuất LoopVQ: gộp các kênh của một vài bước lặp liên tiếp thành một véc-tơ con và chia sẻ cùng một chỉ mục mã hóa (lượng tử hóa véc-tơ trên trục lặp), sau đó tinh chỉnh mã hóa k-means có trọng số Fisher để đầu ra của mô hình sau lượng tử hóa sát với mô hình độ chính xác đầy đủ (trong quá trình huấn luyện, mỗi bước lặp đều đọc bộ nhớ đệm đã lượng tử hóa). Phương pháp này trực giao với các phương pháp dành cho LLM thông thường và có thể kết hợp với phép biến đổi khóa của TaSQ.

Tựa gốc: LoopVQ: Lượng tử hóa bộ nhớ đệm KV dưới 1 bit cho các Transformer lặp lại

Các Transformer lặp lại gọi đi gọi lại cùng một nhóm lớp, có thể bắt kịp các mô hình thông thường lớn gấp nhiều lần với ít tham số hơn, nhưng mỗi bước lặp đều phải lưu trữ bộ nhớ đệm KV riêng, khiến bộ nhớ đệm tăng tuyến tính theo số bước lặp và trở thành nút thắt cổ chai chính về bộ nhớ. Các tác giả nhận thấy bộ nhớ đệm KV có tính dư thừa cao trên trục lặp, trong khi độ nhạy cảm với lượng tử hóa của các bước lặp khác nhau lại rất khác biệt. Do đó, họ đề xuất LoopVQ: gộp các kênh của một vài bước lặp liên tiếp thành một véc-tơ con và chia sẻ cùng một chỉ mục mã hóa (lượng tử hóa véc-tơ trên trục lặp), sau đó tinh chỉnh mã hóa k-means có trọng số Fisher để đầu ra của mô hình sau lượng tử hóa sát với mô hình độ chính xác đầy đủ (trong quá trình huấn luyện, mỗi bước lặp đều đọc bộ nhớ đệm đã lượng tử hóa). Phương pháp này trực giao với các phương pháp dành cho LLM thông thường và có thể kết hợp với phép biến đổi khóa của TaSQ.

Lượt thích 8 · Số sao mã nguồn 3

Lý do

  • Các tác giả tuyên bố rằng với 0.5 bit mỗi phần tử, phương pháp này vượt trội hơn các phương pháp lượng tử hóa véc-tơ bộ nhớ đệm KV hiện có dành cho LLM thông thường trên các tác vụ suy luận chuỗi tư duy dài và truy xuất ngữ cảnh dài của các mô hình Ouro và Huginn; khi kết hợp với TaSQ, ở mức 0.25 bit vẫn giữ được phần lớn độ chính xác BF16.
  • Lợi ích hệ thống đã được chứng minh thực tế: Sau khi triển khai trong SGLang, các tác giả cho thấy độ trễ giải mã giảm gần một nửa với ngữ cảnh 32K, thông lượng phục vụ đỉnh điểm của Ouro 1.4B / 2.6B tăng lần lượt lên 4.3 lần và 6.1 lần so với đường cơ sở BF16.
  • Mã nguồn theo giấy phép Apache-2.0: Quy trình hiệu chuẩn, nhánh SGLang với lưu trữ chỉ mục đóng gói và nhân giải mã hợp nhất, cùng các tập lệnh tái tạo cho từng bảng trong bài báo đều có sẵn trong kho lưu trữ.

Cách sử dụng

  • Chạy ./setup_envs.sh để tạo ba môi trường (loopvq-serve cho engine, loopvq-eval cho đánh giá, loopvq-build cho hiệu chuẩn), export MODELS_DIR trỏ tới Ouro-1.4B, Ouro-2.6B, Ouro-*-Thinking và huginn-0125; để tái tạo bảng tổng quát: MODEL=ouro_1p4b GPUS="0 1 2 3" bash experiments/loopvq_general/build.sh, sau đó MODEL=ouro_1p4b bash experiments/loopvq_general/run_all.sh, cuối cùng python3 experiments/loopvq_general/make_table.py.
  • Khởi động dịch vụ: Trong môi trường loopvq-serve, chạy PREFIX_TOKENS=8 RECENT_TOKENS=64 PREFILL_BACKEND=triton ./scripts/serve_method.sh $MODELS_DIR/Ouro-1.4B vq <目录 mã hóa> --port 30000 --disable-piecewise-cuda-graph, thu được giao diện tương thích OpenAI của SGLang; các token ở đầu và gần nhất được giữ ở định dạng bf16.

Rủi ro

  • Chỉ áp dụng cho các Transformer lặp lại (bài báo sử dụng Ouro 1.4B / 2.6B và Huginn-3.5B), không thể áp dụng trực tiếp cho LLM thông thường; quá trình hiệu chuẩn chỉ sử dụng 16 chuỗi GPQA.
  • Các số liệu về độ trễ và thông lượng đến từ môi trường của các tác giả; lượt thích 8 và số sao mã nguồn 3 (chỉ phản ánh mức độ quan tâm).

Mở · Mã nguồn · Trang bài báo

Hugging Face Daily Papers

Xem tin AI mới nhất