Dasar-dasar AI

Cache KV Anda Tidak Memiliki Masalah Bit. Ini Memiliki Masalah Geometri.

mm
Tambahkan Unite.AI ke sumber pilihan Anda di Google

Pada presisi 2-bit yang identik, satu keputusan tentang sumbu mana yang Anda kuantisasi dapat mengayun skor benchmark dari 2,88 hingga 63,53. Kunci dan nilai memerlukan perlakuan yang berbeda — dan alasan ini ada di dalam persamaan perhatian, bukan perangkat keras.

Ambil Llama-2-13B. Kelompokkan cache kunci-nilainya dengan ukuran grup kuantisasi 32 menjadi dua bit, sementara meninggalkan semuanya di tempat — model yang sama, anggaran bit yang sama, ukuran grup yang sama, benchmark yang sama.

Tergantung pada satu elemen keputusan implementasi, hasil akurasi CoQA menghasilkan 2,88 atau 63,53. Skor menggunakan presisi penuh adalah 66,37.

Keputusan ini tidak tentang berapa banyak bit total yang digunakan. Pertanyaannya sederhana: apa sumbu yang Anda pilih untuk mengelompokkan saat menghitung setiap faktor skala? Ketika Anda memutuskan untuk menggunakan saluran sebagai dimensi pengelompokan (kunci) dan token sebagai dimensi pengelompokan (nilai), Anda berakhir di suatu tempat dalam empat poin dari kinerja presisi penuh. Jika Anda membalik salah satu dari pilihan ini, Anda akan mengalami kerugian kualitas. Jika Anda membalik keduanya, model tidak akan berfungsi.

Empat cara untuk menghabiskan 2 bit yang sama pada cache yang sama. Hasil dari ablasi KIVI pada Llama-2-13B dengan ukuran grup 32.

Kuantisasi biasanya dipikirkan sebagai satu dial: 8 bit, 4 bit, 2 bit, dengan biaya akurasi yang halus. Di dalam cache KV, itu tidak seperti itu. Ini tentang memilih sistem koordinat, dan sistem yang berbeda diterapkan pada kunci dan nilai. Bagian ini menjelaskan mengapa. Secara singkat: kesalahan kuantisasi tergantung pada rentang nilai dalam grup; kunci dan nilai memiliki struktur yang sangat berbeda; dan orang sering terjebak karena Anda tidak dapat menghasilkan sumbu yang tepat dari distribusi nilai sama sekali. Anda harus melihat bagaimana kesalahan berubah setelah perhatian mengonsumsinya. Ini memberikan prinsip umum untuk mengompresi aktivasi antara dan alasan yang baik untuk meragukan kesalahan rekonstruksi sebagai proksi untuk kualitas.

Mengapa Cache KV Adalah Tempat Ini Menggigit

Selama fase generasi, transformer menyimpan semua data proyeksi kunci-nilai (KV) token yang telah diproses sebelumnya ke dalam cache sehingga tidak perlu menghitung data tersebut lagi. Cache tersebut tumbuh secara linier dengan panjang konteks dan ukuran batch. Akhirnya, ini akan menghasilkan cache yang tumbuh lebih besar dari model itu sendiri.

Peningkatan pertumbuhan ini dapat dengan mudah diidentifikasi saat melihat konsumsi memori dari berbagai bagian model. Dalam analisis KVQuant dari LLaMA-7B, bobot akun untuk sekitar 98 persen memori pada panjang urutan 512, dengan aktivasi pada 2 persen. Pada konteks 128K, rasio ini berbalik menjadi sekitar 16 persen bobot dan 84 persen cache KV. Ketika kita melihat analisis OPT-175B yang dikutip oleh penulis KIVI, mereka menemukan hasil yang sama. Secara khusus, pada ukuran batch 512 dengan prompt token 512, cache KV mencapai 1,2TB — beberapa kali ukuran bobot model.

Namun, kapasitas hanya setengah dari masalah di sini. GPU harus membaca seluruh cache KV dari memori perangkat untuk setiap token yang dihasilkan. Ini berarti bahwa saat GPU membaca cache KV, inti komputasi diam. Dengan demikian, mengurangi ukuran cache secara keseluruhan meningkatkan headroom pemrosesan yang tersedia dan mengurangi waktu yang dihabiskan menunggu transfer data.

Apa yang Sebenarnya Terdiri dari Kesalahan Kuantisasi

Kuantisasi integer seragam secara matematis sederhana. Untuk sekelompok angka, Anda merekam angka terkecil sebagai titik nol dan kemudian membagi rentang kelompok tersebut dengan jumlah level yang dapat diwakili untuk mendapatkan ukuran langkah. Anda kemudian membulatkan setiap elemen ke langkah terdekat. Dua hasil langsung mengikuti. Pertama, kesalahan per elemen dibatasi oleh setengah langkah. Kedua, ukuran langkah adalah rentang kelompok dibagi dengan 2^B − 1. Pada 2 bit, Anda hanya memiliki 4 level untuk menutupi penyebaran yang ada dalam kelompok tersebut. Jadi, elemen yang 100 kali lebih besar dibandingkan dengan tetangganya tidak hanya berkinerja buruk. Ini mengembangkan ukuran langkah untuk semua elemen lain yang berbagi kelompok yang sama, dan semuanya menjadi lebih kasar bersama. Kelompok adalah unit kerusakan. Memilih sumbu berarti memutuskan elemen mana yang menderita bersama. Mengajukan pertanyaan secara berbeda, ini tidak lagi “berapa banyak bit yang bisa saya aford?” tetapi “di mana nilai ekstrem dan apakah saya bisa mengisolasi mereka?”

Kunci: Outlier Berada di Saluran Tetap

Model bahasa besar berisi aktivasi yang tidak biasa besar dibandingkan dengan sebagian besar aktivasi. Sun dan rekan-rekan mengkatalogkan aktivasi besar ini di seluruh keluarga model: di Mixtral 8x7B, magnitude terbesar dekat 7000 sementara median magnitude fitur sekitar 0,3 — sekitar empat pesanan besar terpisah. Ini sangat jarang; mereka tetap di dimensi yang jarang berubah dengan input, dan mereka tidak tidak sengaja. Mereka bertindak sebagai bias implisit, dan mereka adalah apa yang memfokuskan perhatian pada beberapa token: perilaku sink perhatian. Di cache kunci, struktur ini sangat jelas: saluran tertentu membawa magnitudo yang sangat besar secara konsisten di seluruh token dalam urutan. Kelompokkan sepanjang token, dan setiap kelompok berisi saluran outlier tersebut, sehingga ukuran langkah setiap kelompok ditetapkan oleh outlier, dan semua saluran biasa membayar untuk itu. Kelompokkan sepanjang saluran, dan saluran outlier membentuk kelompok mereka sendiri. Rentang internal mereka besar tetapi terkandung; saluran biasa dibiarkan sendiri. Hasilnya sesuai. Rata-rata di seluruh lapisan dan kepala pada Llama-2-13B, KIVI melaporkan kesalahan rekonstruksi kunci 13,67 di bawah pengelompokan per-token versus 4,55 per-saluran, dan — yang lebih penting — kesalahan skor perhatian 47,00 versus 9,60. Kuantisasi kunci per-token menghasilkan sekitar lima kali kesalahan skor. Skor kemudian setuju dengan metrik yang berarti untuk kunci; kuantisasi saluran unggul dalam kedua hal.

Nilai: Di Mana Intuisi Patah

Cache nilai tidak menunjukkan pola saluran-outlier. Ini tampaknya cukup datar. Di atasnya, dengan argumen rentang, kita bisa mengharapkan bahwa salah satu dari sumbu ini akan menghasilkan kualitas kompresi yang serupa.

Mereka tidak. Tidak peduli bagaimana manajemen kunci diimplementasikan (hasil 2,80 dan 2,88), mengompresi per-saluran nilai runtuhkan model.

Dan ini adalah tangkapan: jika Anda mengukur kerugian ini menggunakan kesalahan rekonstruksi mentah pada tensor asli yang mana setiap nilai dikompresi, kuantisasi nilai per-saluran sebenarnya terlihat sedikit lebih baik, pada 3,73 versus 4,57. Jika Anda memvalidasi kompresi Anda dengan cara yang jelas, Anda akan memilih konfigurasi yang menghancurkan model.

Kesalahan kuantisasi cache nilai pada Llama-2-13B, diukur dengan dua cara. Metrik tensor yang disimpan dan metrik output yang dikonsumsi tidak setuju oleh lebih dari satu pesanan besar.

Resolusi ini adalah cache nilai tidak pernah dibaca langsung. Ini dikonsumsi oleh produk matriks: output perhatian adalah jumlah tertimbang vektor nilai di seluruh token, dengan skor perhatian softmax sebagai bobot. Karena itu, kesalahan yang relevan adalah kesalahan yang diperkenalkan selama proses ini dan bukan dalam tensor itu sendiri. Diukur dalam hal output perhatian, urutan kesalahan terbalik. Kesalahan relatif yang dilaporkan oleh KIVI untuk output perhatian karena kuantisasi vektor nilai per-token adalah 3,55 dibandingkan dengan 49,89 untuk kuantisasi per-saluran — lebih dari empat belas kali lebih tinggi untuk apa yang tampak seperti pilihan yang lebih baik berdasarkan seberapa baik itu dikompresi.

Penjelasannya adalah kekurangan perhatian, yang mereka ukur sebesar 84,3 persen. Sebagian besar informasi yang terkandung dalam output dapat dikaitkan dengan sejumlah kecil token yang sangat penting. Kuantisasi per-token membatasi kesalahan setiap token ke token itu, sehingga kesalahan pada token yang tidak penting dikalikan dengan bobot perhatian yang hampir nol dan secara efektif menghilang. Kuantisasi per-saluran menyebarkan kesalahan setiap token di seluruh skala saluran yang dibagikan, sehingga token yang direpresentasikan dengan buruk mencemari representasi token yang penting. Kekurangan yang membuat perhatian efisien adalah properti yang sama yang membuat kuantisasi per-token aman.

Pelajaran yang dapat dipindahkan adalah lebih luas dari cache KV: ukur kesalahan kompresi di mana tensor dikonsumsi, bukan di mana tensor disimpan. Asumsi implisit yang dibuat oleh kesalahan rekonstruksi adalah bahwa setiap komponen tensor memiliki bobot yang sama saat menyumbang output akhir. Perhatian secara eksplisit tidak. Setiap operasi downstream yang menimbang, menggate, atau mensparse inputnya mematahkan asumsi tersebut. Pembaca yang familiar dengan artikel sebelumnya saya tentang buta metrik evaluasi dalam sistem pengambilan akan mengenali bahwa hasil ini mirip dengan kegagalan yang dijelaskan sebelumnya: metrik yang mudah dihitung yang melaporkan sesuatu yang lain dari yang dimaksud.

Penanaman Embedding Rotary Mempersulit Kunci

Ada beberapa masalah dengan menggunakan Penanaman Posisi Rotary (RoPE). RoPE memutar pasangan saluran berdasarkan posisi relatif setiap token. Pencampuran ini sebagian melarutkan struktur saluran tetap yang membuat kuantisasi kunci per-saluran bekerja pada tempat pertama — saluran outlier diputar ke tetangganya, dan tetangganya mewarisi rentang. Jawaban KVQuant adalah pengurutan: kuantisasi kunci sebelum rotasi diterapkan, dan terapkan RoPE setelah dekuantisasi. Bersama dengan kuantisasi kunci per-saluran, tipe data non-seragam, dan mengisolasi sebagian kecil outlier, ini mendapatkan mereka di bawah 0,1 degradasi perplexity pada 3 bit, dan memungkinkan melayani LLaMA-7B hingga 1 juta token konteks pada satu A100-80GB.

Juga penting untuk memahami tingkat dampak dari RoPE. Penulis makalah “RotateKV” melaporkan peningkatan 145% dalam kesalahan kuantisasi setelah RoPE ditambahkan, dan mencatat bahwa saluran outlier berbeda di seluruh kepala perhatian — yang mengapa menerapkan satu matriks rotasi bersama di mana-mana tidak cukup, dan rotasi adaptif kepala melakukan lebih baik.

Pajak Sistem, dan Mengapa Ini Bukan Detail

Kuantisasi per-token cocok untuk decoding. Setiap token tiba; Anda kuantisasi, tambahkan ke urutan (sepanjang dimensi token), tidak ada yang lain bergerak.

Namun, kuantisasi per-saluran tidak cocok. Karena statistik saluran melintasi token yang belum dihasilkan, Anda tidak dapat menghitung faktor skala saat token masuk. Solusi KIVI adalah untuk menyimpan token terbaru — hingga 128 — dalam presisi penuh dalam buffer residual, dan kuantisasi dalam kelompok sekali cukup banyak telah terkumpul.

Seperti yang terjadi, buffer residual menjadi load-bearing, bukan hanya sesuatu yang insidental. Pada GSM8K dengan Llama-2-7B, skor presisi penuh 13,50. Kuantisasi penuh ke 2 bit dengan sumbu yang benar, skor 5,76. Sumbu yang sama dan bit yang sama, plus buffer residual token terbaru pada presisi penuh, skor 12,74. Jendela geser token terbaru pada presisi penuh akan pulih sebagian besar apa yang hilang karena kuantisasi agresif pada masalah multi-langkah yang sulit — yang akan masuk akal jika kita mempertimbangkan token mana yang sedang dihadapi oleh rantai operasi aritmatika.

Ada manfaat signifikan dari melakukan semua hal ini dengan benar — seperti yang dilaporkan KIVI, 2,6 kali lebih sedikit penggunaan memori puncak untuk Llama-2-7B, memungkinkan ukuran batch hingga 4 kali lebih besar, serta 2,35 hingga 3,47 kali lebih baik throughput pada tugas layanan dunia nyata.

Apa yang Harus Dilakukan dengan Ini

  1. Jangan pernah menggunakan satu kuantisasi untuk keduanya. Gunakan kuantisasi yang berbeda untuk kunci (per-saluran) dan untuk nilai (per-token). Pipa yang menerapkan kuantisasi tunggal ke “cache KV” telah mungkin mengorbankan sebagian besar kualitas yang mungkin saat menggunakan jumlah bit kecil untuk mewakili setiap nilai.
  2. Kuantisasi kunci sebelum RoPE. Ini adalah masalah kebenaran daripada preferensi.
  3. Simpan jendela presisi penuh token terbaru. Meskipun menyimpan jendela seperti itu mengambil sangat sedikit memori dibandingkan dengan seberapa besar cache dapat, ini adalah area yang menghasilkan sebagian besar akurasi untuk tugas yang sulit.
  4. Jangan validasi pada kesalahan rekonstruksi. Validasi selalu berdasarkan pada output perhatian atau berdasarkan kinerja tugas akhir. Metrik penyimpanan tidak hanya berisik — untuk nilai, itu menunjuk ke arah yang salah.
  5. Jangan validasi pada benchmark pilihan berganda konteks pendek. Penulis KIVI sengaja menghindari tugas tertutup seperti MMLU untuk evaluasi ini, karena satu langkah dekoding yang membaca logit output hampir tidak menguji cache sama sekali. Setiap evaluasi yang tidak membangun cache dari waktu ke waktu dan kemudian melakukan generasi dari itu tidak akan pernah dapat mengamati kegagalan yang melekat dalam desain sistem Anda.

Di Mana Pekerjaan Ini Menuju

Meskipun masih ada beberapa hal yang harus dilakukan mengenai sifat geometris masalah ini, banyak peneliti terus mempelajari cara saluran outlier didistribusikan di antara berbagai kepala transformer, dan bagaimana keterbatasan perangkat keras mempengaruhi pengelompokan mana yang paling murah: InnerQ melipatkan normalisasi kunci saluran ke dalam bobot kunci dan pertanyaan selama pengisian awal. Oleh karena itu, tidak ada overhead tambahan yang dikeluarkan pada runtime. Selain itu, InnerQ menyimpan jendela presisi tinggi untuk token terbaru yang dihasilkan dan token sink perhatian. Dengan melakukan ini, InnerQ menghilangkan kesempatan untuk outlier di saluran sink untuk mencemari saluran tetangga.

Yang lain mengusulkan bahwa alih-alih menyimpan cache seluruhnya, kita harus menyimpan hanya cukup informasi untuk dapat rematerialisasi kunci dan/atau nilai(v) sesuai kebutuhan dari representasi cache yang lebih kecil.

Akhirnya, penting untuk diingat bahwa akurasi bukan satu-satunya parameter yang dipengaruhi oleh kuantisasi. Penelitian yang baru-baru ini diterbitkan menunjukkan degradasi pemosisian yang dihasilkan dari kuantisasi cache KV. Selain itu, penelitian ini mendokumentasikan degradasi pemosisian bahkan di lingkungan layanan vLLM produksi yang menggunakan cache FP8 bersama dengan protokol pemulihan tanpa pelatihan yang memulihkan hingga 97% dari apa yang hilang dalam hal pemosisian. Sehingga, sementara konfigurasi mungkin mempertahankan hasil benchmark, itu tidak berarti bahwa itu mempertahankan semua parameter lain yang Anda pedulikan.

Prinsip Umum

Gagasan kuantisasi telah dibingkai sebagai “anggaran presisi”: berapa banyak bit yang bisa saya korbankan? Cache KV menunjukkan bahwa pertanyaan yang lebih berguna adalah struktural. Presisi dialokasikan dalam kelompok; kelompok adalah unit kerusakan, dan sumbu yang Anda kelompokkan menentukan elemen mana yang berbagi nasib. Sumbu yang benar adalah sumbu di mana tensor Anda dikonsumsi, yaitu cara Anda menggunakan tensor Anda DAN BUKAN bagaimana tensor Anda muncul saat disimpan dalam memori. Kunci digunakan melalui komputasi dot-produk terhadap pertanyaan. Saluran yang rusak tunggal akan meracuni semua skor. Nilai dikonsumsi melalui komputasi rata-rata tertimbang yang langka di seluruh token. Oleh karena itu, token yang rusak tunggal hanya ditimbang.

Dua tensor dengan dimensi identik dan dihasilkan oleh dua lapisan berturut-turut diperlakukan secara berbeda. Ini layak untuk mempertanyakan setiap aktivasi yang Anda rencanakan untuk mengompresi: apa operasi yang mengontrak ini, dan apakah pengelompokan saya menghormatinya?

Himanshu Goel adalah peneliti AI/ML yang mengkhususkan diri dalam generasi yang ditingkatkan dengan pengambilan untuk domain dengan risiko tinggi, termasuk alur kerja dokumen biomedis, keuangan, dan peraturan.