KV ქეშის კვანტიზაცია: მეხსიერების რევოლუცია დიდ ენობრივ მოდელებში

ოდესმე გიცდიათ თქვენი ენობრივი მოდელით გრძელი ტექსტის გენერირება, მხოლოდ იმისთვის, რომ მეხსიერების შემაწუხებელი შეზღუდვების გამო დაბრკოლებას წაწყდომოდით? რამდენადაც ენობრივი მოდელები აგრძელებენ ზრდას ზომითა და შესაძლებლობებით, გრძელი გენერაციების მხარდაჭერამ შესაძლოა მეხსიერების დიდი ნაწილი მოიხმაროს. ეს ხშირი პრობლემაა, განსაკუთრებით მაშინ, როდესაც შეზღუდულ რესურსებთან გაქვთ საქმე. სწორედ აქ ერთვება kv ქეშის კვანტიზაცია სიტუაციის გადასარჩენად. მაშ, რა არის ზუსტად kv ქეშის კვანტიზაცია? თუ ეს ტერმინი თქვენთვის უცხოა, არ იდარდოთ! მოდით, დავყოთ ის ორ ნაწილად: kv ქეში და კვანტიზაცია. Key-value ქეში, ანუ kv ქეში, საჭიროა ავტორეგრესიულ მოდელებში გენერაციის ოპტიმიზაციისთვის, სადაც მოდელი ტექსტს ტოკენ-ტოკენით წინასწარმეტყველებს. ეს პროცესი შეიძლება ნელი იყოს, რადგან მოდელს შეუძლია მხოლოდ ერთი ტოკენის გენერირება ერთდროულად, და ყოველი ახალი წინასწარმეტყველება დამოკიდებულია წინა კონტექსტზე. ეს ნიშნავს, რომ გენერაციაში 1000-ე ტოკენის წინასწარმეტყველებისთვის, გჭირდებათ ინფორმაცია წინა 999 ტოკენიდან, რომელიც მოდის ამ ტოკენების წარმოდგენების მატრიცული გამრავლების სახით. მაგრამ 1001-ე ტოკენის წინასწარმეტყველებისთვის, ასევე გჭირდებათ იგივე ინფორმაცია პირველი 999 ტოკენიდან, პლუს დამატებითი ინფორმაცია 1000-ე ტოკენიდან. სწორედ აქ გამოიყენება key-value ქეში თანმიმდევრული გენერაციის პროცესის ოპტიმიზაციისთვის წინა გამოთვლების შენახვით, რათა ისინი ხელახლა იქნას გამოყენებული მომდევნო ტოკენებში, ამიტომ მათი ხელახლა გამოთვლა საჭირო არ არის. უფრო კონკრეტულად, key-value ქეში მოქმედებს როგორც მეხსიერების ბანკი ავტორეგრესიული გენერაციული მოდელებისთვის, სადაც მოდელი ინახავს key-value წყვილებს, რომლებიც მიღებულია self-attention ფენებიდან ადრე დამუშავებული ტოკენებისთვის. ტრანსფორმერების არქიტექტურაში, self-attention ფენები ითვლიან attention ქულებს queries-ის keys-ზე გამრავლებით, რის შედეგადაც მიიღება value ვექტორების შეწონილი ჯამები. ამ ინფორმაციის შენახვით, მოდელს შეუძლია თავიდან აიცილოს ზედმეტი გამოთვლები და სანაცვლოდ ქეშიდან მოიძიოს წინა ტოკენების key-ები და value-ები. K+1-ე ტოკენისთვის attention ქულების გამოთვლისას, ჩვენ არ გვჭირდება წინა key-ების და value-ების ხელახლა გამოთვლა, არამედ ვიღებთ მათ ქეშიდან და ვამატებთ მიმდინარე ვექტორს. ეს, როგორც წესი, იწვევს ტექსტის უფრო სწრაფ და ეფექტურ გენერაციას. გადავიდეთ მეორე ტერმინზე: კვანტიზაცია არის რიცხვითი მნიშვნელობების სიზუსტის შემცირება მეხსიერების დასაზოგად. კვანტიზაციის დროს, თითოეული რიცხვითი მნიშვნელობა მრგვალდება ან იჭრება, რათა მოერგოს შემცირებული სიზუსტის ფორმატს, რამაც შესაძლოა გამოიწვიოს ინფორმაციის დაკარგვა. თუმცა, კვანტიზაციის პარამეტრებისა და ტექნიკების ფრთხილად შერჩევას შეუძლია მინიმუმამდე დაიყვანოს ეს დანაკარგი, ხოლო მაინც მიაღწიოს დამაკმაყოფილებელ შესრულებას. არსებობს სხვადასხვა კვანტიზაციის მეთოდი. მიუხედავად იმისა, რომ kv ქეში აჩქარებს ავტორეგრესიულ გენერაციას, ის შეიძლება გახდეს მეხსიერების შეზღუდვა ხანგრძლივი კონტექსტის ან დიდი batch ზომის დროს. მოდით, შევაფასოთ, რამდენი მეხსიერება დაგვჭირდება kv ქეშის შესანახად 10000 ტოკენის სიგრძის შეყვანისთვის 7B Llama-2 მოდელისთვის. ერთი ტოკენის kv ქეშის შესანახად საჭირო მეხსიერება არის დაახლოებით 2 * 2 * num_layers * num_key_value_heads * head_dim, სადაც პირველი 2 აღნიშნავს key-ებს და value-ებს, ხოლო მეორე 2 არის ბაიტების რაოდენობა, რომელიც გვჭირდება (იმ ვარაუდით, რომ მოდელი ჩატვირთულია float16-ში). ასე რომ, თუ გვაქვს 10000 ტოკენის სიგრძის კონტექსტი, დაგვჭირდება 2 * 2 * 32 * 32 * 128 * 10000 ≈ 5GB მეხსიერება მხოლოდ წინა key-value ქეშის შესანახად, რაც თითქმის მესამედია იმ მეხსიერებისა, რომელიც საჭიროა მოდელის პარამეტრების ნახევრად სიზუსტით შესანახად. ამიტომ, kv ქეშის უფრო კომპაქტურ ფორმაში შეკუმშვით, შეგვიძლია დიდი რაოდენობით მეხსიერების დაზოგვა და უფრო ხანგრძლივი კონტექსტის გენერაცია ჩვეულებრივ GPU-ებზე. ჩვენს ექსპერიმენტებში, ჩვენ შევძელით მნიშვნელოვნად შეგვემცირებინა მეხსიერების მოხმარება ხარისხის ზედმეტად გაუარესების გარეშე kv ქეშის დაბალი სიზუსტის ფორმატებში კვანტიზაციით. ამ ახალი კვანტიზაციის ფუნქციით, ჩვენ ახლა შეგვიძლია მხარი დავუჭიროთ უფრო ხანგრძლივ გენერაციებს მეხსიერების ამოწურვის გარეშე, რაც ნიშნავს, რომ თქვენ შეგიძლიათ გააფართოვოთ თქვენი მოდელის კონტექსტის სიგრძე მეხსიერების შეზღუდვაზე ფიქრის გარეშე. Key-value ქეშის კვანტიზაცია Transformers-ში დიდწილად შთაგონებული იყო ნაშრომით KIVI: A Tuning-Free Asymmetric 2bit Quantization for kv Cache. ნაშრომმა წარმოადგინა 2-ბიტიანი ასიმეტრიული კვანტიზაცია დიდი ენობრივი მოდელებისთვის ხარისხის დეგრადაციის გარეშე. KIVI ახდენს key ქეშის კვანტიზაციას არხის მიხედვით (per-channel) და value ქეშის კვანტიზაციას ტოკენის მიხედვით (per-token), რადგან მათ აჩვენეს, რომ LLM-ებისთვის key-ებს აქვთ გამონაკლისების უფრო მაღალი მაგნიტუდები ზოგიერთ არხში, ხოლო value-ები ასეთ ნიმუშს არ ავლენენ. ამიტომ, კვანტიზებულ და ორიგინალ სიზუსტეს შორის ფარდობითი შეცდომა გაცილებით მცირეა, როდესაც key-ები კვანტიზებულია არხის მიხედვით და value-ები ტოკენის მიხედვით. მეთოდში, რომელიც ჩვენ ინტეგრირებული გვაქვს...