EasyBalance: Cross-Layer Load Balancing in Distributed MoE Inference
EasyBalance एक क्रॉस-लेयर लोड बैलेंसिंग रणनीति है जो वितरित Mixture-of-Experts (MoE) इन्फरेंस के लिए है, जो विशेषज्ञ प्रतिकृति (expert replication), माइग्रेशन, या विशेषज्ञ-डिवाइस मैपिंग में संशोधन की आवश्यकता के बिना, परतों के बीच वर्कलोड को लालची (greedily) रूप से शेड्यूलिंग और विलंबित करके विषम रूटिंग वितरण के कारण होने वाली GPU आइडलिंग को कम करती है।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक विशाल, उच्च-गति वाली लाइब्रेरी चला रहे हैं जहाँ हजारों छात्र (टोकन) को एक विशाल विश्वकोश से उत्तर खोजने की आवश्यकता है। यह विश्वकोश किसी एक व्यक्ति द्वारा नहीं लिखा गया है; यह एक "मिश्रण विशेषज्ञों का" (Mixture of Experts - MoE) है, जिसका अर्थ है कि इसके अंदर सैकड़ों अलग-अलग विशेषज्ञ लेखक मौजूद हैं। जब कोई छात्र प्रश्न पूछता है, तो एक स्मार्ट लाइब्रेरियन (राउटर) जल्दी से निर्णय लेता है कि कौन से कुछ लेखक उस प्रश्न का उत्तर देने के लिए सबसे उपयुक्त हैं। इसे सुपर फास्ट बनाने के लिए, लाइब्रेरी इन लेखकों को कई अलग-अलग कंप्यूटरों (डिवाइसों) में विभाजित करती है जो समानांतर (parallel) में काम करते हैं।
लेकिन, एक समस्या है: सभी प्रश्न एक जैसे नहीं होते। कभी-कभी, छात्रों की एक बड़ी भीड़ ऐसे प्रश्न पूछती है जिनका उत्तर केवल एक विशिष्ट लेखक ही जानता है। उस लेखक का कंप्यूटर भर जाता है और धीमा हो जाता है, जबकि अन्य कंप्यूटर, जिनके लेखक कम व्यस्त हैं, बस हाथ पर हाथ धरे बैठे रहते हैं और उस धीमे चलने वाले के खत्म करने का इंतज़ार करते हैं। यह "इंतज़ार का खेल" बहुत सारी ऊर्जा और समय बर्बाद करता है। वर्षों तक, इसका समाधान अधिक लेखकों को काम पर रखना या भार को संतुलित करने के लिए उन्हें इधर-उधर ले जाना था, लेकिन इसमें बहुत अधिक मेमोरी लगती है और इसे तुरंत करना कठिन है।
यहाँ आता है EasyBalance, एक चतुर नई रणनीति जो बिना किसी नए व्यक्ति को काम पर रखे या फर्नीचर बदले इस "इंतज़ार के खेल" को हल करती है। लेखकों को ठीक करने के बजाय, EasyBalance यह बदल देता है कि छात्र अपने प्रश्न कब पूछते हैं। यह महसूस करता है कि हालांकि लाइब्रेरी आमतौर पर एक समय में एक लेयर (परत) के अनुसार प्रश्नों को प्रोसेस करती है, वास्तव में यह विभिन्न लेयर्स के छात्रों को एक साथ काम करने की अनुमति दे सकती है। इन समूहों को आपस में मिलाकर, व्यस्त लेखकों को राहत मिलती है क्योंकि एक समूह के "भारी" प्रश्न दूसरे समूह के "हल्के" प्रश्नों के साथ संतुलित हो जाते हैं। परिणाम यह है कि कंप्यूटर व्यस्त रहते हैं, प्रतीक्षा समय गायब हो जाता है, और पूरी लाइब्रेरी बहुत तेज़ी से चलती है।
समस्या: "सबसे धीमे का इंतज़ार करें" का नियम
AI की दुनिया में, विशेष रूप से इन "मिश्रण विशेषज्ञों के" (Mixture of Experts) मॉडलों के साथ, सिस्टम अविश्वसनीय रूप से कुशल होने के लिए डिज़ाइन किया गया है। यह प्रत्येक डेटा के लिए केवल कुछ ही "विशेषज्ञों" को सक्रिय करता है। लेकिन एक वितरित सेटअप (distributed setup) में—जहाँ ये विशेषज्ञ कई ग्राफिक्स कार्डों (GPUs) पर फैले हुए हैं—चीजें गड़बड़ हो जाती हैं।
एक रिले रेस की कल्पना करें जहाँ बैटन (baton) अगले धावक को तभी दिया जाता है जब सभी अपना भाग पूरा कर लेते हैं। यदि एक धावक एक भारी बैकपैक (एक "हॉट" विशेषज्ञ जिसके पास बहुत अधिक टोकन हैं) ले जा रहा है, तो वह पूरी टीम को धीमा कर देता है। अन्य धावक, जिन्होंने अपना हल्का भार पूरा कर लिया है, खड़े होकर इंतज़ार करते हैं। पेपर के शब्दों में, इसे लोड इम्बैलेंस (load imbalance) कहा जाता है। सिस्टम सबसे धीमे डिवाइस द्वारा बाधित होता है, जिससे बाकी सब खाली बैठे रह जाते हैं।
इसे ठीक करने के पिछले प्रयासों में एक्सपर्ट रेप्लिकेशन (व्यस्त लेखकों की अतिरिक्त प्रतियां बनाना) या एक्सपर्ट माइग्रेशन (लेखकों को अलग-अलग कंप्यूटरों पर ले जाना) शामिल था। हालाँकि ये काम करते हैं, लेकिन इनके बड़े नुकसान हैं: ये बहुत अधिक मेमोरी खाते हैं, संचार में अतिरिक्त समय लेते हैं, और ये कठोर (rigid) हैं। यदि आप लाइब्रेरी को मिलने वाले प्रश्नों के प्रकार को बदलते हैं, तो पुराना प्लान पूरी तरह से विफल हो सकता है।
समाधान: EasyBalance
इस पेपर के लेखक, यिज़े वू और उनके सहयोगियों ने EasyBalance नामक एक नया दृष्टिकोण प्रस्तावित किया है। उनका बड़ा विचार लेखकों को कंप्यूटरों पर मैप करने के तरीके को ठीक करने के बजाय, काम के शेड्यूल को ठीक करना है।
वे दो प्रमुख अंतर्दion (insights) पर भरोसा करते हैं:
- क्रॉस-लेयर रेडंडेंसी (Cross-Layer Redundancy): भले ही मॉडल के एक विशिष्ट लेयर के पास विशेषज्ञों का एक विशिष्ट सेट होता है, मॉडल के अन्य लेयर्स के विशेषज्ञ पहले से ही कंप्यूटर की मेमोरी में मौजूद होते हैं, जाने के लिए तैयार। वे वर्तमान कार्य के लिए "स्वाभाविक रूप से रेडंडेंट" हैं। आपको नए संस्करणों को काम पर रखने की आवश्यकता नहीं है; आपको बस उन्हीं का उपयोग करने की आवश्यकता है जो आपके पास पहले से हैं।
- वर्कलोड कॉम्बिनेशन (Workload Combination): पेपर सुझाव देता है कि आप मॉडल की विभिन्न लेयर्स से आने वाले माइक्रो-बैच (प्रश्नों के छोटे समूह) को एक साथ चला सकते हैं। भले ही मॉडल आमतौर पर चीजों को एक समय में एक लेयर के हिसाब से प्रोसेस करता है, गणित दिखाता है कि इन समूहों को मिलाना सुरक्षित है। वास्तव में, यह अक्सर बेहतर होता है। यदि ग्रुप A का कंप्यूटर 1 पर भारी भार है, लेकिन ग्रुप B का कंप्यूटर 2 पर भारी भार है, तो उन्हें एक साथ चलाने से कुल भार संतुलित हो जाता है। "सबसे खराब" स्थिति (जहाँ दोनों समूह एक ही कंप्यूटर पर टकराते हैं) सांख्यिकीय रूप से दुर्लभ है, खासकर जब आप अधिक कंप्यूटर जोड़ते हैं।
यह कैसे काम करता है: "स्मार्ट शेड्यूलर"
EasyBalance एक स्मार्ट ट्रैफिक कंट्रोलर की तरह काम करता है। सभी छात्रों को लेयर 1 से भेजने के बजाय, फिर लेयर 2 से, यह वर्तमान भीड़ को देखता है। यह एक साथ चलाने के लिए विभिन्न लेयर्स से छात्रों का एक मिश्रण चुनता है।
- यह तुरंत निष्पादित (execute) करने के लिए माइक्रो-बैच के एक उपसमूह (subset) का चयन करता है।
- यह उन अन्य समूहों को टाल देता है (इंतज़ार करवाता है) जो बाधा (bottleneck) उत्पन्न कर सकते हैं।
- यह यह सब बिना यह बदले कि कौन सा विशेषज्ञ किस कंप्यूटर पर रहता है, करता है।
इसका मतलब है कि सिस्टम बिना हार्डवेयर या मेमोरी को पुनर्गठित किए किसी भी नए प्रकार के कार्य के लिए तुरंत अनुकूलित हो सकता है। यह एक रेस्टोरेंट किचन की तरह है जो बर्गर और सलाद को एक साथ पकाने का निर्णय लेता है क्योंकि ग्रिल व्यस्त है लेकिन सलाद स्टेशन खाली है, बजाय इसके कि ग्रिल के सब कुछ खत्म करने का इंतज़ार किया जाए।
परिणाम: तेज़ और कम बर्बादी
शोधकर्ताओं ने Qwen3-30B और Moonlight-16B सहित कई बड़े मॉडलों पर 8 GPUs पर EasyBalance का परीक्षण किया। उन्होंने LongBench नामक एक बेंचमार्क का उपयोग किया, जिसमें रीडिंग कॉम्प्रिहेन्शन और कोड जनरेशन जैसे कई प्रकार के कार्य शामिल हैं।
निष्कर्ष सुसंगत और प्रभावशाली थे:
- कम आइडलिंग (Reduced Idling): EasyBalance ने GPU "अंडर-यूटिलाइजेशन" (वह समय जब कंप्यूटर कुछ नहीं कर रहा होता) को मुख्य रूप से 40% से अधिक कम कर दिया। कई मामलों में, खाली समय लगभग 0.35 (35% बर्बादी) से गिरकर लगभग 0.2 (20% बर्बादी) हो गया।
- तेज़ इन्फरेंस (Faster Inference): क्योंकि कंप्यूटर एक-दूसरे का इंतज़ार नहीं कर रहे थे, इसलिए उत्तर प्राप्त करने का कुल समय (एंड-टू-एंड लेटेंसी) काफी कम हो गया।
- लचीलापन (Flexibility): पिछले तरीकों के विपरीत, जिन्हें सेटअप करने के लिए विशेषज्ञों को पहले से जानने की आवश्यकता थी, EasyBalance सभी 13 परीक्षण किए गए कार्यों (ट्रिविया से लेकर कोड तक) में उतना ही अच्छा काम करता है।
पेपर ने विभिन्न "शेड्यूलिंग" रणनीतियों (कैसे सिस्टम तय करता है कि समूहों को कैसे मिलाया जाए) का भी पता लगाया। उन्होंने पाया कि MaxUtil नामक एक रणनीति (जो हर GPU के उपयोग को अधिकतम करने की कोशिश करती है) सबसे अच्छा काम करती है, लेकिन CumUtil (बैच को एक-एक करके जोड़ना यदि वे मदद करते हैं) जैसी सरल, तेज़ रणनीतियाँ भी कुछ न करने से बहुत बेहतर हैं।
यह क्यों महत्वपूर्ण है
सबसे रोमांचक बात यह है कि EasyBalance के लिए किसी अतिरिक्त मेमोरी या जटिल पुनर्गठन की आवश्यकता नहीं है। यह मौजूदा सेटअप के साथ काम करता है। जैसे-जैसे AI मॉडल बड़े और अधिक जटिल होते जा रहे हैं, कुछ कंप्यूटरों के काम करने के दौरान दूसरों के खाली बैठने की समस्या और भी बढ़ जाएगी। यह पेपर बताता है कि विशेषज्ञों को रखने के स्थान के बजाय, हम काम को कब चलाते हैं, इसके बारे में स्मार्ट होकर, इन विशाल AI सिस्टम को काफी अधिक कुशल बना सकते हैं।
लेखक बताते हैं कि हालांकि उनकी विधि अत्यधिक प्रभावी है, यह इस सांख्यिकीय संभावना पर निर्भर करती है कि भारी लोड हमेशा एक ही समय में एक ही कंप्यूटर पर नहीं पड़ेगा। विभिन्न मॉडलों और कार्यों में उनके परीक्षणों में, इस रणनीति ने लगातार इन्फरेंस को तेज किया, जिससे यह साबित हुआ कि कभी-कभी बाधा को हल करने का सबसे अच्छा तरीका श्रमिकों को स्थानांतरित करने के बजाय, उन्हें फिनिश लाइन तक एक-दूसरे की मदद करने देना है।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।