तकनीकी सारांश: PagedWeight

समस्या विवरण

Mixture-of-Experts (MoE) Large Language Models (LLMs) प्रति टोकन केवल विशेषज्ञों के एक उपसमुच्चय (subset) को सक्रिय करके उच्च दक्षता और सटीकता प्रदान करते हैं। हालांकि, सर्विंग परिदृश्यों में, MoE मॉडल एक गंभीर मेमोरी तनाव का सामना करते हैं: GPU को भारी मॉडल वेट्स (weights) और बढ़ते हुए Key-Value (KV) कैश दोनों को स्टोर करना होता है। जबकि PagedAttention जैसी तकनीकें KV कैश विखंडन (fragmentation) को प्रबंधित करती हैं, वे MoE वेट्स के महत्वपूर्ण मेमोरी फुटप्रिंट को संबोधित नहीं करती हैं, जो GPU मेमोरी के 60% से अधिक हिस्से पर कब्जा कर सकते हैं।

मौजूदा क्वांटिज़ेशन विधियाँ काफी हद तक स्थिर (static) हैं, जो इन्फरेंस शुरू होने से पहले ही प्रिसिजन (precision) को निर्धारित कर देती हैं। हालांकि कुछ रनटाइम प्रिसिजन-अनुकूलन विधियाँ मौजूद हैं, वे GPU मेमोरी को वेट्स और KV कैश के बीच गतिशील रूप से पुनर्वितरित करने के लिए प्रिसिजन परिवर्तनों का उपयोग नहीं करती हैं। फलस्वरूप, ऐसा कोई सिस्टम उपलब्ध नहीं है जो कम महत्वपूर्ण विशेषज्ञ वेट्स से मेमोरी को गतिशील रूप से मुक्त कर सके ताकि KV कैश के विस्तार के साथ कार्य की सटीकता से समझौता किए बिना या इन्फरेंस को बाधित किए बिना इसे किया जा सके।

कार्यप्रणाली: PagedWeight

लेखक PagedWeight का प्रस्ताव करते हैं, जो MoE LLM सर्विंग के लिए एक नवीन प्रबंधन प्रणाली है जो विशेषज्ञ-वेट प्रिसिजन और KV कैश आकार के बीच संतुलन बनाने के लिए रनटाइम पर मॉडल वेट्स को गतिशील रूप से क्वांटाइज़ करती है। यह सिस्टम विशेषज्ञ वेट्स को पज्ड (paged) KV कैश ब्लॉक्स की तरह मानता है, जिससे डायनेमिक ऑफलोडिंग और रीलोडिंग संभव होती है।

मुख्य घटक

1. Paged Weight प्रतिनिधित्व (Representation):

   • PagedWeight प्रति-विशेषज्ञ लीनियर-ब्लॉक (लेयर, विशेषज्ञ और मॉड्यूल प्रकार: gate_up या down द्वारा पहचाने गए) के स्तर पर कार्य करता है।
   • यह Any-Precision LLM (APL) तकनीक का उपयोग करता है, जो वेट्स को संबंधित लुकअप टेबल्स (LUTs) के साथ एक ओवरलेड बिट-प्लेन फॉर्मेट में प्रदर्शित करता है।
   • एक Weight Page Table प्रत्येक लीनियर-ब्लॉक के लिए वांछित बिटविड्थ ($d_i$) और प्रतिबद्ध (committed) बिटविड्थ ($q_i$) को ट्रैक करता है। पेज GPU-रेसिडेंट, CPU-रेसिडेंट, या ट्रांसफर अवस्थाओं में हो सकते हैं।

2. क्वालिटी-अवेयर रनटाइम प्लानर (Quality-Aware Runtime Planner):
   प्लानर सटीकता हानि को न्यूनतम करने के लिए यह निर्धारित करता है कि किन वेट्स को क्वांटाइज़ (बिटविड्थ कम) किया जाना चाहिए, जो तीन कारकों पर आधारित है:

   • ऑफलाइन ग्लोबल सेंसिटिविटी: क्वांटिज़ेशन के प्रति प्रत्येक लीनियर-ब्लॉक की अंतर्निहित संवेदनशीलता का अनुमान लगाने के लिए हेसियन-वेटेड (Hessian-weighted) स्कोर का उपयोग करता है।
   • ऑनलाइन रूटिंग सांख्यिकी (Online Routing Statistics): विशेषज्ञों के चयन की आवृत्ति और भार ("रूटिंग मास") को ट्रैक करता है। बार-बार रूट किए जाने वाले ("हॉट") विशेषज्ञों को उच्च बिटविड्थ फ्लोर और डैमेज मल्टीप्लायर के साथ सुरक्षित रखा जाता है।
   • प्रॉम्प्ट रेसिडुअल्स (Prompt Residuals): लीनियर रिग्रेशन हेड्स के माध्यम से अनुमानित एक प्रॉम्प्ट-विशिष्ट सुधार शब्द। यह वर्तमान इनपुट सीक्वेंस के आधार पर ग्लोबल सेंसिटिविटी अनुमान को समायोजित करता है, यह स्वीकार करते हुए कि क्वांटिज़ेशन का प्रभाव प्रॉम्प्ट के अनुसार भिन्न होता है।

   प्लानर संभावित बिटविड्थ कटौती के लिए एक "डैमेज" स्कोर की गणना करता है। जब KV कैश दबाव एक मेमोरी लक्ष्य को ट्रिगर करता है, तो प्लानर सबसे कम-डैमेज वाली कटौती (उच्चतम गुणवत्ता प्रति बचा हुआ बाइट) का लालची (greedy) चयन करता है ताकि लक्ष्य प्राप्त किया जा सके।

3. एसिंक्रोनस निष्पादन और कर्नेल ऑप्टिमाइज़ेशन:

   • एसिंक्रोनस पेज मूवमेंट: लेटेंसी को छिपाने के लिए, सिस्टम वेट पेजों को CPU RAM में ऑफलोड करता है और उन्हें एसिंक्रोनस रूप से रीलोड करता है। बिटविड्थ कमिटमेंट्स को केवल सुरक्षित सीमाओं (जैसे, एक प्लान के पूरी तरह से निष्पादित होने के बाद) पर अपडेट किया जाता है, जिससे यह सुनिश्चित होता है कि इन्फरेंस बाधित न हो।
   • फ्यूज्ड मिक्स्ड-प्रिसिजन कर्नेल: एक कस्टम CUDA कर्नेल Any-Precision बिट-प्लेन और LUTs को सीधे पढ़ता है, जो एकल फ्यूज्ड ऑपरेशन के भीतर प्रत्येक लीनियर-ब्लॉक के लिए अलग-अलग बिटविड्थ का समर्थन करता है ताकि ओवरहेड को कम किया जा सके।

मुख्य योगदान

1. सिस्टम डिज़ाइन: ऑनलाइन Any-Precision MoE वेट्स के लिए एक पज्ड-वेट सिस्टम का प्रस्ताव, जो कमिटेड/वांछित बिटविड्थ, पेज स्टेट्स और मेमोरी खपत को गतिशील रूप से प्रबंधित करता है।
2. प्लानिंग एल्गोरिदम: एक क्वालिटी-अवेयर रनटाइम प्लानर जो लो-डैमेज वेट रिडक्शन रणनीतियों को चुनने के लिए ऑफलाइन सेंसिटिविटी, ऑनलाइन रूटिंग सांख्यिकी और प्रॉम्प्ट रेसिडुअल्स को संश्लेषित करता है।
3. निष्पादन रणनीति: लेटेंसी को छिपाने के लिए न्यूनतम थ्रूपुट लॉस के साथ एसिंक्रोनस पेज मूवमेंट का कार्यान्वयन, जो एक फ्यूज्ड मिक्स्ड-प्रिसिजन MoE कर्नेल के साथ जुड़ा हुआ है।
4. अनुभवजन्य सत्यापन (Empirical Validation): तीन MoE मॉडल्स (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B) में भाषा मॉडलिंग, रीजनिंग और लॉन्ग-कॉन्टेक्स्ट कार्यों पर व्यापक मूल्यांकन, जो मौजूदा बेसलाइन्स की तुलना में बेहतर प्रदर्शन प्रदर्शित करता है।

परिणाम

लेखकों ने भाषा मॉडलिंग, रीजनिंग और लॉन्ग-कॉन्टेक्स्ट कार्यों में FP16, Any-Precision LLM (APL), DP-LLM, और MxMoE के विरुद्ध PagedWeight का मूल्यांकन किया।

• क्वालिटी-मेमोरी ट्रेडऑफ़: PagedWeight बेसलाइन्स की तुलना में 72.0% GPU मेमोरी बचत और 1.94× थ्रूपुट सुधार के साथ FP16-तुल्य सटीकता प्राप्त करता है।
• निश्चित बजट पर क्वालिटी: समान मेमोरी बजट पर, PagedWeight क्वांटिज़ेशन विधियों की तुलना में कार्य गुणवत्ता में 39.3% तक सुधार करता है, जिसमें थ्रूपुट का नुकसान अधिकतम 4.1% है।
• लॉन्ग-कॉन्टेक्स्ट प्रदर्शन: कड़े मेमोरी बजट (जैसे, 10 GB) के तहत, PagedWeight FP16 के तुलनीय लॉन्ग-कॉन्टेक्स्ट क्वालिटी (Passage Retrieval, NarrativeQA) बनाए रखता है, जबकि स्टैटिक क्वांटिज़ेशन बेसलाइन्स (जैसे, APL) महत्वपूर्ण गिरावट का सामना करते हैं।
• थ्रूपुट: PagedWeight यूनिफॉर्म क्वांटिज़ेशन बेसलाइन्स के थ्रूपुट से मेल खाता है, जिसमें बैच साइज 4 पर सबसे बड़ा देखा गया ड्रॉप 4.1% है।
• एब्लेशन (Ablation): रूटिंग सांख्यिकी, प्रॉम्प्ट रेसिडुअल्स, या पेज मूवमेंट जैसे घटकों को हटाने से प्रदर्शन में लगातार गिरावट आती है, जो पूर्ण सिस्टम डिज़ाइन की आवश्यकता की पुष्टि करती है।

महत्व

पेपर का दावा है कि PagedWeight प्रभावी रूप से मॉडल कार्य सटीकता, मेमोरी खपत और थ्रूपुट/लेटेंसी के बीच जटिल ट्रेडऑफ़ को नेविगेट करता है। विशेषज्ञ वेट्स को स्थिर संपत्तियों के बजाय पज्ड, प्रबंधनीय इकाइयों के रूप में मानकर, PagedWeight कार्य सटीकता और वेट्स एवं KV कैश के बीच GPU मेमोरी के आवंटन के बीच एक अनछुए ट्रेडऑफ़ स्पेस को खोलता है। यह सिस्टम प्रदर्शित करता है कि डायनेमिक, क्वालिटी-अवेयर क्वांटिज़ेशन, स्टैटिक क्वांटिज़ेशन या उन विधियों की तुलना में एक व्यवहार्य और बेहतर रणनीति है जो रनटाइम मेमोरी दबाव के अनुकूल नहीं होती हैं।