लेखकहरू:
(१) बेन अथिवरातकुन, AWS AI ल्याब्स;
(२) सुजन कुमार गोनुगोन्डला, AWS AI ल्याब्स;
(३) संजय कृष्ण गौडा, AWS AI ल्याब्स;
(4) Haifeng Qian, AWS AI प्रयोगशाला;
(५) संजय कृष्ण गौडा, AWS AI ल्याब्स;
(६) हान्टियन डिङ, AWS एआई ल्याब्स;
(७) क्विङ सन, AWS एआई ल्याब्स;
(८) जुन वाङ, AWS AI ल्याब्स;
(9) Jiacheng Guo, AWS AI ल्याब्स;
(१० लिआंगफु चेन, AWS AI ल्याब्स;
(११) परमिंदर भाटिया, जीई हेल्थकेयर (एडब्ल्यूएसमा काम गरेको);
(12) रमेश नल्लापति, Amazon AGI (AWS मा गरिएको काम);
(१३) सुदिप्ता सेनगुप्ता, AWS AI ल्याब्स;
(१४) बिंग जियाङ, गोल्डम्यान साक्स (AWS मा गरिएको काम)।
लिङ्कहरूको तालिका
३.१. नोटेशन र ३.२. भाषा मोडेल अनुमान
३.३. बहु-प्रश्न, बहु-शीर्षक र सामान्यीकृत बहु-प्रश्न ध्यान
४. सन्दर्भ-सचेत विभाजित ध्यान र ४.१. प्रेरणा
४.२. सूत्रीकरण र ४.३. मेमोरी IO जटिलता
५.१. बहु-मुख्य, बहु-क्वेरी, र बहु-समूह ध्यानको क्षमताहरूको तुलना गर्ने
५.२. क्षमताहरूको विलम्बता-समतुल्य मोडेलहरू
G. सट्टा डिकोडिङ र द्रुत डिकोडिङ प्रविधिहरूसँग अनुकूलता
२. सम्बन्धित काम
साहित्यमा, अनुमान विलम्बता र/वा विलम्बता सुधार गर्ने धेरै तरिकाहरू छन्। क्वान्टाइजेसनले int8, int4, र fp8 जस्ता कम-बिटविड्थ प्रतिनिधित्वहरू प्रयोग गरेर मेमोरी प्रयोग घटाउँछ (Wei et al., 2023; Yao et al., 2022; Dettmers et al., 2022; Frantar et al., 2022; Kuzmin et al., 2022; Xiao et al., 2022)। मोडेल प्यारामिटरहरूमा मात्र लागू गर्दा क्वान्टाइजेसनले घट्दो परिणामहरू प्रदान गर्दछ जस्तै लामो अनुक्रम लम्बाइ र ठूला ब्याच आकारहरू जहाँ मेमोरी पहुँच र डट-उत्पादन ध्यानसँग सम्बन्धित कम्प्युटले समग्र अनुमान विलम्बतामा प्रभुत्व जमाउँछ।
स्पर्स एटेन्सन (बेलटागी एट अल।, २०२०; चाइल्ड एट अल।, २०१९; जहीर एट अल।, २०२०) लाई लामो सन्दर्भहरूको लागि ध्यानको जटिलता कम गर्ने र छिटो अनुमान गर्ने तरिकाको रूपमा व्यापक रूपमा अध्ययन गरिएको छ। पोप एट अल। (२०२२) ले विलम्बता र मोडेल FLOPs उपयोगमा पारेटो सीमा प्राप्त गर्न TPUs (सामूहिक einsum) को लागि अनुकूलित बहु-आयामी विभाजन प्रविधिहरू प्रयोग गरेर ठूला भाषा मोडेलहरूको जेनेरेटिभ इन्फरेन्स दक्षताको अनुसन्धान गर्छन्। पेपरले यो पनि देखाउँछ कि बहु-क्वेरी एटेन्सनले उच्च ब्याच साइज अन्तर्गत दक्षतामा जोड दिँदै ३२x ठूलो सन्दर्भ लम्बाइसम्म स्केल गर्न अनुमति दिन्छ। पृष्ठबद्ध ध्यान (क्वोन एट अल।, २०२३) ले KV क्यासको मेमोरी व्यवस्थापनलाई ब्लकहरूमा विभाजन गरेर र म्यापिङ उद्देश्यका लागि ब्लक तालिका प्रयोग गरेर बढाउँछ। यो दृष्टिकोणले गतिशील कार्यभार परिवर्तनहरूलाई प्रभावकारी रूपमा समायोजन गर्दछ र धेरै आउटपुट अनुक्रमहरूमा प्रम्प्टको KV क्यासको साझेदारी मार्फत मेमोरी भण्डारण आवश्यकताहरू कम गर्दछ। यद्यपि, यसले KV क्यासको मेमोरी रिड घटाउँदैन।
सट्टा डिकोडिङ, र यसको भेरियन्टहरूले धेरै अनुक्रमिक टोकनहरू प्रस्ताव गर्न सानो ड्राफ्ट मोडेल प्रयोग गर्दछ, जुन मुख्य मोडेलद्वारा समानान्तर रूपमा प्रशोधन गरिन्छ त्यस्ता टोकनहरू स्वीकार वा अस्वीकार गर्न (चेन एट अल।, २०२३; लेवियाथन एट अल।, २०२२; ली एट अल।, २०२४; काई एट अल।, २०२४; फु एट अल।, २०२३)। मुख्य विचार भनेको प्रत्येक चरणमा धेरै टोकनहरूको डिकोडिङ सक्षम गर्नु हो, जसले गर्दा मुख्य मोडेलको मेमोरी IO प्रयोगहरू परिशोधन हुन्छन्। यद्यपि, डिकोडिङको विलम्बता अझै पनि ठूला सन्दर्भ आकारहरूमा KV क्यास I/O ब्यान्डविथद्वारा हावी हुनेछ, जहाँ विभाजित ध्यानले डिकोडिङ गतिलाई अझ बढाउन सक्छ। छोटकरीमा, वृद्धिशील डिकोडिङले मोडेल लोडिङको परिशोधित मेमोरी IO कम गर्नमा केन्द्रित हुन्छ जबकि बहु-क्वेरी र विभाजित ध्यानले KV क्यासको मेमोरी IO कम गर्छ।
३. पृष्ठभूमि
३.१. नोटेशन
हामी सम्पूर्ण पेपरमा निम्न संकेतन प्रयोग गर्छौं।
३.२. भाषा मोडेल अनुमान
भाषा मोडेलका लागि धेरै अनुमान परिदृश्यहरू छन्, जसमा ब्याच अनुमान र एकल-सन्दर्भ ब्याच नमूना समावेश छ (चित्र १)। ब्याच अनुमानले त्यो केसलाई जनाउँछ जहाँ हामी ब्याचमा धेरै इनपुटहरू सँगै प्रशोधन गर्छौं, र प्रत्येक ब्याच अनुक्रमणिकाको लागि स्वतन्त्र रूपमा पछिल्ला टोकनहरू उत्पन्न गर्छौं। ब्याच आकार १ भएको अवस्थामा, यो एकल-सन्दर्भ अनुमानमा घट्छ। अर्को परिदृश्य एकल-सन्दर्भ ब्याच नमूना हो जहाँ हामी एकल सन्दर्भमा आधारित धेरै अनुक्रमहरू उत्पन्न गर्छौं, जहाँ ब्याच अनुमान केस बीचको भिन्नता यो हो कि KV क्यास प्राप्त गर्न एकल सन्दर्भको लागि प्रिफिल मात्र गर्न आवश्यक छ, त्यसपछि अन्य ब्याच सूचकांकहरूमा प्रसारण गरिन्छ।
चित्र १ ले भाषा मोडेल अनुमानका दुई चरणहरूलाई पनि चित्रण गर्दछ: (क) सन्दर्भ एन्कोडिङ वा प्रिफिलिङ र (ख) वृद्धिशील डिकोडिङ। सन्दर्भ एन्कोडिङले एकल फर्वार्ड पासलाई जनाउँछ जसले सन्दर्भमा सबै टोकन स्थितिहरूको लागि कुञ्जी र मान टेन्सरहरू गणना गर्दछ। एक पटक कुञ्जी र मान टेन्सरहरू गणना गरिसकेपछि, हामी वृद्धिशील डिकोडिङ चरणको समयमा ध्यान संयन्त्रको लागि प्रयोग गर्न यी कुञ्जी र मान टेन्सरहरूलाई क्यास गर्छौं, जसले क्रमिक रूपमा एक पटकमा एउटा टोकन उत्पन्न गर्दछ [2]।
सन्दर्भ एन्कोडिङ चरणको समयमा, मेमोरी इनपुट/आउटपुट (IO) अपरेशनहरूको सापेक्षमा फ्लोटिंग पोइन्ट अपरेशनहरूको संख्या उच्च हुन्छ, कम्प्युट-बाउन्ड शासनसँग मेल खान्छ जहाँ विलम्बता FLOPs द्वारा प्रभावित हुन्छ। यद्यपि, वृद्धिशील डिकोडिङको समयमा जहाँ हामी एकल क्वेरी टोकनमा ध्यान दिन्छौं, यो मेमोरी-बाउन्ड शासनमा पर्दछ जहाँ प्रति मेमोरी पहुँच गणनाको संख्या लगभग 1-देखि-1 हुन्छ (विवरणहरूको लागि परिशिष्ट D.1 हेर्नुहोस्)। मेमोरी IO ले उच्च ब्यान्डविथ मेमोरी (HBM) (Jia et al., 2018) बाट द्रुत अन-चिप SRAM सम्म पढ्ने र लेख्ने अपरेशनहरूलाई जनाउँछ जहाँ वास्तविक गणना हुन्छ। वृद्धिशील डिकोडिङको मेमोरी IO मा दुई घटकहरू हुन्छन्: (1) मोडेल प्यारामिटर लोडिङ र (2) KV क्यास लोडिङ। घटक (1) सन्दर्भ लम्बाइ m वा ब्याच आकार b को पर्वाह नगरी स्थिर हुन्छ जहाँ घटक (2) m र b दुवैमा निर्भर गर्दछ र यदि m वा b उच्च छ भने समग्र मेमोरी IO मा हावी हुन्छ, जुन अनुमानको लागि महत्त्वपूर्ण बाधा बन्न सक्छ। हाम्रो काम मुख्यतया घटक (2) घटाउने कुरामा केन्द्रित छ।
यो पेपर CC BY 4.0 DEED लाइसेन्स अन्तर्गत arxiv मा उपलब्ध छ।