भाग 2 Transformer वास्तुकला के विकास को वैनिला Transformers से आधुनिक LLMs तक ट्रैक करता है।

यह लेख उन सभी के लिए है जो बुनियादी ध्यान तंत्र को समझते हैं और देखना चाहते हैं कि वार्तालाप बढ़ने पर एक बड़े भाषा मॉडल की सेवा क्यों महंगी हो जाती है। आप सीखेंगे कि कैसे KV Cache निरर्थक गणना को हटाता है, इसकी स्मृति लागत का अनुमान कैसे लगाया जाए, और इस लागत ने शोधकर्ताओं को MQA, GQA और MLA की ओर धकेल दिया।

क्वेरी, कुंजी और मान वास्तव में क्या करते हैं

ध्यान को पढ़ने का एक सरल तरीका खोज प्रणाली के रूप में है। क्वेरी (Q) प्रश्न है। कुंजी (K) आपको बताती है कि प्रत्येक टोकन क्या दर्शाता है, इसलिए मॉडल यह न्याय कर सकता है कि वह टोकन क्वेस्ट से प्रासंगिक है या नहीं। वैल्यू (V) वास्तव में पुनर्प्राप्त की गई जानकारी है जब कोई टोकन संबंधित हो जाता है।

श्रम का यह विभाजन बताता है कि आरओपीई में स्थिति संबंधी जानकारी मुख्य रूप से क्यू और के पर क्यों लागू होती हैः ध्यान तय करता है कि किस टोकन को क्यू-के संबंध के माध्यम से बातचीत करनी चाहिए। KV Cache, हालांकि, पूरी तरह से एक अलग समस्या हल करती है।

ऑटोरेग्रेसिव पीढ़ी में पुनरावृत्ति समस्या

अनुमान के दौरान, LLM एक समय में एक टोकन उत्पन्न करता है। प्रत्येक नए चरण पर, नव निर्मित टोकन को सभी पिछले टोकनों का ध्यान रखना चाहिए। मान लीजिए कि मॉडल पहले से ही उत्पादित किया गया हैः → नाम → है → नाइमिश. अगली टोकन पैदा करते समय, इसे सभी पूर्ववर्ती टोकन के कुंजी और मूल्य प्रतिनिधित्व की आवश्यकता होती है।

कैशिंग के बिना, बढ़ते क्रम मॉडल को पहले से संसाधित किए गए टोकन के लिए K और V की पुनः गणना करने के लिए मजबूर करते हैं। चरण एक में, यह "द" के लिए के 1 और वी 1 की गणना करता है। कदम दो पर, यह फिर से K1 और V 1 का हिसाब लगाता है, फिर k2 और v 2 जोड़ता है। चरण तीन पर, इसे k1 और v 1 की पुनर्गणना करनी होती है, k2 तथा v 2 की पुनर्गणना करना होता है, इसके बाद k3 और v 3 को जोड़ा जाता है। the और name के लिए कुंजी और मानों की गणना पहले ही कर दी गई है, लेकिन उन्हें फिर से गणना किया जा रहा है।

अतीत की गणना एक बार करें, फिर इसे पुनः उपयोग करें

फिक्स यह है कि पहले से गणना की गई K और V को डिस्पोज करने के बजाय मेमोरी में संग्रहीत किया जाए। जब अगला टोकन आता है, तो मॉडल केवल नए के और वी की गणना करता है और उन्हें मौजूदा कैश में जोड़ता है। वैचारिक रूप से: चरण एक के बाद कैश [K1, V1] रखता है; चरण दो के बाद इसे [K1 , K2] और [V1 , V2] रखती है; कदम तीन के बाद यह [के 1, K2 , K3] और [वी 1 , वी 2 , वी 3 ] रखता है ।

मॉडल अतीत की पुनर्गणना करने के बजाय पहले से गणना किए गए प्रतिनिधित्वों का पुनः उपयोग करता है। यह ऑटोरेग्रेसिव अनुमान के दौरान रिडंडेंट कंप्यूटेशन को काफी कम कर देता है और LLM पीढ़ी को कुशल बनाने के लिए सबसे महत्वपूर्ण तकनीकों में से एक है। लेकिन अनुकूलन एक नया सवाल पेश करता हैः यदि हम प्रत्येक पिछले टोकन के लिए K और V संग्रहीत करते रहते हैं, तो संदर्भ बढ़ने पर कैश कितनी GPU मेमोरी खपत करेगा?

क्यू, के और वी की गणना कैसे होती है

प्रशिक्षण के दौरान, मॉडल अपने मापदंडों को सीखता है, जिसमें प्रक्षेपण मैट्रिक्स भी शामिल हैं जो Q, K और V उत्पन्न करते हैं। एक बार प्रशिक्षण पूरा हो जाने पर, ये भार अनुमान के दौरान तय किए जाते हैं, और इनपुट टोकन के लिए क्यू, के और वी की गणना करने के लिए उन प्रशिक्षित वजन का उपयोग किया जाता है।

एक चार टोकन इनपुट लें: The, name, is, naimish. सरलता के लिए, मान लीजिए कि प्रत्येक टोकन को 3-आयामी एम्बेडिंग द्वारा दर्शाया गया है, जो 4 × 3 आकार का एक इनपुट मैट्रिक्स X देता है, जहां 4 टोकन की संख्या है और 3 एम्बेड आयाम है। ट्रांसफार्मर में तीन सीखे गए प्रक्षेपण मैट्रिसेस होते हैंः क्वेरी के लिए Wq, कुंजी के लिए डब्ल्यूके और मूल्य के लिए वी। इस सरलीकृत उदाहरण में, डब्लूक्यू पंक्तियों [0.12, 1.3, 8.4], 2.4, 5.7, 2.9 और 1.2, 7.4, 1.7] के साथ एक 3 × 3 मैट्रिस है; Wk और Wv समान रूप से परिभाषित किए जाते हैं। ये मैट्रिक प्रशिक्षण के दौरान सीखा गया था और निष्कर्षण के दौरान तय रहते थे।

फिर हम गणना करते हैं Q = XWq, K=XWk, और V= XWv. 4 × 3 पर एक्स के साथ और प्रत्येक प्रक्षेपण 3 × 3, क्यू, के, और वी में से प्रत्येक 4 × 3. एक वास्तविक मल्टी-हेड ट्रांसफार्मर में आते हैं। ये अनुमान ध्यान सिरों में विभाजित होते हैं, लेकिन मूल विचार अपरिवर्तित रहता है।

Q, K, V से ध्यान स्कोर तक

यह तय करने के लिए कि कौन से टोकन को अधिक ध्यान दिया जाना चाहिए, मॉडल डॉट उत्पाद का उपयोग करके प्रत्येक क्वेरी और हर कुंजी के बीच समानता की गणना करता हैः Attention Scores = QKT. फिर यह स्कोर को प्रमुख आयाम के वर्गमूल द्वारा बढ़ाता है, जिससे QKT / √dk मिलता है, और सॉफ्टमैक्स लागू होता है ताकि ध्यान वजन प्राप्त किया जा सके। अंत में, भार वैल्यू मैट्रिक्स से गुणा किए जाते हैं। पूरा स्केलेड डॉट-उत्पाद ध्यान सूत्र Attention ((Q, K, V) = softmax (QKT/√dk) वी है।

प्रवाह मायने रखता हैः क्यू × केटी मापता है कि प्रत्येक टोकन कितना प्रासंगिक है, सॉफ्टमैक्स इसे ध्यान भार में बदल देता है, और वी द्वारा गुणा करने से संदर्भ या भाग लिया गया प्रतिनिधित्व उत्पन्न होता है। यह ठीक वही है जहां KV Cache भुगतान करता है। नए जनित टोकन की क्वेरी को सभी पिछले टोकनों की कुंजी के साथ बातचीत करनी चाहिए, लेकिन उन पिछली कुंजियों की गणना पहले ही कर दी गई है। उन्हें फिर से गिनने के बजाय, मॉडल कैश से उनका पुनः उपयोग करता है. केवल नए टोकन के K और V की गणना और संलग्न करने की आवश्यकता होती है।

स्मृति बिल

KV Cache एक प्रमुख समस्या को हल करता है, बार-बार पिछले टोकन के कुंजियों और मूल्यों की पुनः गणना करना, लेकिन यह दूसरा बनाता हैः उन सभी कुंजी और मूल्य अब GPU मेमोरी में बैठते हैं, और कैश संदर्भ के साथ बढ़ता है। दस उत्पन्न टोकन का मतलब K और V के दस सेट होते हैं। एक हजार टोकन एक हजार सेटों का मतलब होता है।

7B-पैरामीटर मॉडल के लिए 32 परतों, 32 ध्यान सिर, हेड आयाम d = 128, 1 के बैच आकार के साथ FP16 में चल रहा है, KV Cache लगभग GB को 4K टोकन पर खपत करता है, लगभग 8 32K टोकन और 128K टोकन के आसपास लगभग 32 GB. 128K टोकन में, कैश अकेले वीआरएएम के बारे में 32 GB की आवश्यकता हो सकती है, इससे पहले कि मॉडल वजन और अन्य रनटाइम आवश्यकताओं के लिए लेखांकन किया जाए। यही कारण है कि संदर्भ लंबाई बढ़ने पर KV Cache एक गंभीर मेमोरी बाधा बन जाती है। संक्षेप में, यह स्मृति के लिए गणना का व्यापार करती है।

क्यों कैश स्केल जिस तरह से यह करता है

L परतों के साथ एक ट्रांसफार्मर पर विचार करें, H_KV कुंजी-मूल्य सिर, हेड आयाम d और संदर्भ लंबाई T. प्रत्येक परत और हर टोकन के लिए, K और V दोनों को संग्रहीत किया जाना चाहिए। बैच आकार, डेटा प्रकार और कार्यान्वयन विवरण की अनदेखी करते हुए, कैश किए गए मूल्यों की अनुमानित संख्या 2 × L × T × H_ KV × d_head के आनुपातिक है, जहां 2 के और वी दोनों को स्टोर करने से आता है।

प्रैक्टिकल रीडिंग: KV Cache साइज स्केल के साथ संदर्भ लंबाई × KV हेड्स की संख्या × हेड आयाम × परतों की संख्या। जैसे-जैसे संदर्भ लंबाई बढ़ती है, कैश बहुत बड़ा हो सकता है। यह भी ध्यान दें कि कैश मेमोरी सीधे कुंजी-मूल्य वाले सिरों की मात्रा से आनुपातिक होती है, जिसका अर्थ है कि K / V प्रमुखों को कम करने से सीधे निष्कर्ष स्मृति में कमी आती है। मल्टी-Head Attention (एमएचए) कई केवी प्रमुखों का उपयोग करता है, समूहीकृत -Query Attention (जीक्यूए) कम और बहु-Query Attention (MQA) एक का उपयोग करते हैं।

क्यों मल्टी-Head Attention महंगा मामला है

मानक एमएचए में, प्रत्येक ध्यान सिर का अपना क्वेरी, कुंजी और मूल्य प्रक्षेपण होता हैः हेड 1 → क्यू 1, के 1, वी 1; हेड 2 → क्यू 2, के 2, वी 2; और इसी तरह हेड एच → क्यूएच, ख, व्ही के माध्यम से। ऑटोरेग्रेसिव अनुमान के दौरान, मॉडल को इसलिए हर पिछले टोकन और हर ट्रांसफार्मर परत के लिए K1, V1 से Kh, Vh तक कैश करना होगा।

क्वेरी वेक्टर समस्या नहीं हैं. नए टोकन की क्वे्री को प्रत्येक डिकोडिंग चरण में गणना और उपयोग किया जाता है। प्रमुख मेमोरी लागत सभी टोकन के लिए ऐतिहासिक K और V को संग्रहीत करने से आती है। यह सवाल उठाता है कि ध्यान डिजाइन को फिर से आकार दिया गया हैः क्या हमें वास्तव में हर ध्यान सिर के लिए एक अलग कुंजी और मूल्य कैश की आवश्यकता है?

अगला क्या है

KV Cache कुशल LLM अनुमान के लिए सबसे सरल लेकिन सबसे महत्वपूर्ण अनुकूलन में से एक है। यह मेमोरी समस्या के साथ गणना की समस्या को बदल देता है, और वह स्मृति समस्या संदर्भ लंबाई, परतों की संख्या और KV हेड गिनती बढ़ने पर बदतर हो जाती है। MHA के साथ यह विशेष रूप से तीव्र हो जाता है, जहां प्रत्येक ध्यान सिर अपने स्वयं के कुंजी और मूल्य प्रतिनिधित्व बनाए रखता है।

अनुवर्ती प्रश्न, क्या प्रत्येक क्वेरी हेड को वास्तव में अपने स्वयं के K और V की आवश्यकता होती है, ने अधिक स्मृति-कुशल वास्तुकला का नेतृत्व किया: MQA, फिर GQA, और अंततः एमएलए (मल्टी-Head Latent Attention) जैसे अधिक परिष्कृत दृष्टिकोण। Reducing KV-कैश मेमोरी कुशल ध्यान का केवल एक हिस्सा है, हालांकि। वास्तविक गणना और ध्यान में स्मृति आंदोलन भी एक प्रमुख बाधा बन गया, जिसने Flash Attention जैसी तकनीकों और अन्य कुशल-ध्यान तंत्रों को जन्म दिया।

विकास रातोंरात ट्रांसफार्मर को बदलने के बारे में नहीं है. शोधकर्ताओं ने पूछा कि क्या ध्यान तेज, सस्ता और अधिक स्मृति-कुशल बनते हुए मजबूत रह सकता है। यह अगले भाग में पालन करने का मुद्दा है।