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

1. समस्या विवरण

बड़े पैमाने के Mixture-of-Experts (MoE) आर्किटेक्चर विशेषज्ञों के एक उपसमुच्चय (subset) को सक्रिय करके ट्रिलियन-पैरामीटर वाले Large Language Models (LLMs) के प्रशिक्षण को सक्षम करते हैं। हालाँकि, इन मॉडलों के प्रशिक्षण में विशेषज्ञ समानांतरता (expert parallelism - EP) के लिए आवश्यक all-to-all संचार के कारण गंभीर बाधाएं आती हैं। जैसे-जैसे मॉडल कई GPU नोड्स में स्केल होते हैं, सीमित इंटर-नोड नेटवर्क बैंडविड्थ (जैसे, Infiniलैंड/InfiniBand) संचार विलंबता (latency) को बढ़ा देती है, जो एक प्रशिक्षण चरण के समय का 50% तक उपभोग कर सकती है।

मौजूदा अनुकूलन रणनीतियाँ मुख्य रूप से ऑपरेटर स्तर पर गणना को संचार के साथ ओवरलैप करने (जैसे, FFN गणनाओं को चंकिंग करना) या माइक्रो-बैच पाइपलाइनिंग का उपयोग करने पर ध्यान केंद्रित करती हैं। हालाँकि, इन दृष्टिकोणों में दो मौलिक सीमाएँ हैं:

1. अवशिष्ट स्टॉल (Residual Stalls): अटेंशन और FFN परतों के बीच गणना-संचार अनुपात में अंतर्निहित असंतुलन पूर्ण ओवरलैप को रोकता है। अटेंशन गणना अनुक्रम लंबाई (sequence length) के साथ द्विघात (quadratically) रूप से स्केल होती है, जबकि FFN और संचार रैखिक (linearly) रूप से स्केल होते हैं। यह "कम्युनिकेशन टेल्स" (communication tails) बनाता है जिन्हें छिपाया नहीं जा सकता, विशेष रूप से जैसे-जैसे अनुक्रम लंबाई बढ़ती है।
2. संसाधन अक्षमता (Resource Inefficiency): पूर्ववर्ती विसंयोजित (disaggregated) दृष्टिकोण (अटेंशन और FFN को अलग-अलग हार्डवेयर पर रखना) अक्सर यह मांग करते हैं कि प्रत्येक समूह के पास पूर्ण मॉडल पैरामीटर हों, जिससे बड़े पैमाने पर आउट-ऑफ-मेमोरी (OOM) त्रुटियां होती हैं, या वे प्रत्येक घटक की विशिष्ट गणना-संचार विशेषताओं के अनुकूल संसाधन आवंटन करने में विफल रहते हैं।

2. कार्यप्रणाली: DisagMoE

DisagMoE एक प्रशिक्षण प्रणाली है जिसे एक विसंयोजित आर्किटेक्चर (disaggregated architecture) के माध्यम से मॉडल प्लेसमेंट और शेड्यूलिंग को संयुक्त रूप से अनुकूलित करके all-to-all बाधाओं को कम करने के लिए डिज़ाइन किया गया है।

2.1 विसंयोजित घटक प्लेसमेंट (Disaggregated Component Placement)

पारंपरिक एक्सपर्ट पैरेललिज्म के विपरीत, जहाँ अटेंशन और FFN परतें एक ही GPU समूह में सह-अस्तित्व में होती हैं, DisagMoE मॉडल को घटक प्रकार के आधार पर विभाजित करता है:

• अटेंशन समूह (A-Workers): इन समूहों में कई ट्रांसफार्मर परतों से सेल्फ-अटेंशन परतें शामिल हैं। वे घने (dense) घटकों को रेप्लिकेट करते हैं और डेटा डिस्ट्रीब्यूशन के लिए डेटा पैरेललिज्म (DP) का उपयोग करते हैं।
• FFN समूह (F-Workers): इन समूहों में कई ट्रांसफार्मर परतों (विशेषज्ञों) से FFN परतें शामिल हैं। विशेषज्ञों को समूह के भीतर EP के माध्यम से GPU में विभाजित किया जाता है।
• लेयर असाइनमेंट: परतों को समूहों में इंटरलीव्ड तरीके से सौंपा जाता है (उदाहरण के लिए, ग्रुप 0 लेयर्स 0, $p$, $2p$ को संभालता है; ग्रुप 1 लेयर्स 1, $p+1$, आदि को संभालता है), जो मेमोरी बाधाओं के तहत डिवाइस उपयोगिता को अधिकतम करता है।

2.2 AF-Pipe: ओवरलैपिंग के साथ मल्टी-स्टेज पाइपलाइन

DisagMoE एक AF-Pipe पाइपलाइन शेड्यूल पेश करता है जो संचार को केवल एक साइड इफेक्ट के बजाय एक प्रथम-श्रेणी के स्टेज के रूप में मानता है।

• Many-to-Many (M2N) संचार: एक एकल नोड समूह के भीतर मानक all-to-all एक्सचेंज के बजाय, AF-Pipe अटेंशन और FFN समूहों के बीच एकदिशीय, many-to-many टोकन एक्सचेंज स्थापित करता है।
• स्टेज अलाइनमेंट: पाइपलाइन समूहों के बीच सीमाओं को संरेखित करती है, पारंपरिक "डिस्पैच" और "कंबाइन" ऑपरेशनों को एक एकल एकीकृत संचार चरण में फ्यूज करती है। यह संचार ओवरहेड को लगभग $1/k$ (जहाँ $k$ शीर्ष-$k$ विशेषज्ञ हैं) कम करता है और अनावश्यक डेटा ट्रांसफर को समाप्त करता है।
• एसिंक्रोनस निष्पादन (Asynchronous Execution): सिस्टम प्रत्येक वर्कर के भीतर तीन समन्वित स्ट्रीम्स (फॉरवर्ड, बैकवर्ड, और कम्युनिकेशन) का उपयोग करता है। यह FFN समूह को उपलब्ध डेटा पर फॉरवर्ड पास की गणना करते समय ही ग्रेडिएंट्स (N2M) भेजने की अनुमति देता है, जिससे गैर-ब्लॉकिंग निष्पादन और निरंतर GPU उपयोग सुनिश्चित होता है।

2.3 अनुकूली वर्कर आवंटन (Adaptive Worker Allocation)

अटेंशन (कंप्यूट-बाउंड) और FFN (कम्युनिकेशन-बाउंड) के बीच असंतुलन को संबोधित करने के लिए, DisagMoE एक कंप्यूट-कम्युनिकेशन रूफलाइन मॉडल द्वारा निर्देशित एक अनुकूली एलोकेटर का उपयोग करता है।

• मॉडलिंग: सिस्टम उस टर्निंग पॉइंट ($\hat{I}$) को मॉडल करता है जहाँ एक घटक नेटवर्क-बाउंड से कंप्यूट-बाउंड में बदल जाता है। अनुक्रम लंबाई बढ़ने के साथ अटेंशन परतें FFN परतों की तुलना में कंप्यूट रूफ तक तेजी से पहुँचती हैं।
• अनुकूलन: सिस्टम GPU और नेटवर्क इंटरफेस कार्ड (NIC) विभाजन को एक Mixed Integer Linear Programming (MILP) समस्या के रूप में तैयार करता है। यह गतिशील रूप से FFN समूहों को अधिक नेटवर्क बैंडविड्थ (उनके कम्युनिकेशन रूफ को उठाने के लिए) और अटेंशन समूहों को अधिक कंप्यूट संसाधन आवंटित करता है, जिससे बबल्स को कम करने के लिए पाइपलाइन चरणों को संतुलित किया जाता है।
• खोज रणनीति (Search Strategy): एक एल्गोरिदम रूफलाइन मॉडल का उपयोग करके एक सीड कॉन्फ़िगरेशन उत्पन्न करता है, जिसे फिर अनमॉडल किए गए ओवरहेड्स को संभालने के लिए प्रति-चरण मापे गए समय का उपयोग करके लोकल सर्च के माध्यम से परिष्कृत किया जाता है।

3. प्रमुख योगदान

1. विसंयोजित घटक प्लेसमेंट: एक नवीन आर्किटेक्चर जो अटेंशन और FFN परतों को अलग-अलग वर्कर समूहों में विभाजित करता है, जिससे स्वतंत्र स्केलिंग और संसाधन आवंटन संभव होता है।
2. AF-Pipe शेड्यूल: एक मल्टी-स्टेज पाइपलाइन जो पारंपरिक all-to-all के स्थान पर अंतर-समूह M2N एक्सचेंजों को लाती है, जो व्यवस्थित रूप से अटेंशन और FFN दोनों गणनाओं के साथ संचार को ओवरलैप करती है।
3. अनुकूली संसाधन आवंटन: एक रूफलाइन-निर्देशित तंत्र जो अटेंशन और FFN समूहों के लिए GPU और NIC संसाधनों को संयुक्त रूप से आकार देता है, जो विभिन्न अनुक्रम लंबाई और मॉडल कॉन्फ़िगरेशन केAcross कंप्यूट-कम्युनिकेशन अनुपात को प्रभावी ढंग से संतुलित करता है।

4. प्रयोगात्मक परिणाम

लेखकों ने तीन MoE मॉडलों: DeepSeek-MoE, GPT-OSS-120B, और Qwen3 के साथ, 4K से 32K की अनुक्रम लंबाई के लिए 16 नोड्स (128× H800 GPUs) के क्लस्टर पर DisagMoE का मूल्यांकन किया।

• थ्रूपुट सुधार: DisagMoE ने Megatron-LM (1F1B इंटरलीव्ड), Tutel, Comet, और DualPipe सहित मजबूत बेसलाइन को लगातार पछाड़ दिया।
  • Megatron-1F1B की तुलना में 1.81× स्पीडअप प्राप्त किया।
  • अत्याधुनिक MoE ओवरलैप सिस्टमों की तुलना में 1.34× स्पीडअप प्राप्त किया।
• संचार में कमी: बेसलाइन की तुलना में, DisagMoE ने गैर-ओवरलैप्ड संचार समय को Tutel (88%), Comet (75%), और DualPipe (45%) के मुकाबले काफी कम कर दिया।
• स्केलेबिलिटी: सिस्टम ने विभिन्न top-$k$ मानों और एक्सपर्ट पैरेललिज्म (EP) आकारों में प्रभावशीलता प्रदर्शित की।
• संसाधन संवेदनशीलता: प्रयोगों ने पुष्टि की कि अटेंशन और FFN समूहों के बीच इष्टतम GPU अनुपात अनुक्रम लंबाई के साथ बदलता है (जैसे, लंबी अनुक्रम लंबाई अधिक अटेंशन संसाधनों के पक्ष में होती है), जो अनुकूली आवंटन की आवश्यकता को प्रमाणित करता है।

5. महत्व और दावे

पेपर का दावा है कि DisagMoE वर्तमान MoE प्रशिक्षण की एक मौलिक सीमा को संबोधित करता है: अटेंशन और FFN परतों के बीच संरचनात्मक असंतुलन के कारण संचार को पूरी तरह से ओवरलैप करने में असमर्थता। इन घटकों को विसंयोजित (disaggregating) करके और रूफलाइन मॉडल के आधार पर संसाधनों को अनुकूली रूप से आवंटित करके, सिस्टम संचार विंडो को विस्तारित करता है।

लेखक इस बात पर जोर देते हैं कि उनका दृष्टिकोण विशेष रूप से स्थिर अनुक्रम लंबाई वाले बड़े पैमाने के प्री-ट्रेनिंग वर्कलोड के लिए महत्वपूर्ण है। वे उल्लेख करते हैं कि हालांकि उनका वर्तमान डिज़ाइन स्थिर आकारों (निश्चित अनुक्रम लंबाई और माइक्रो-बैच) को मानता है, लेकिन मूल अंतर्दृष्टि—कि भविष्य के हार्डवेयर और सॉफ्टवेयर को केवल कच्चे बैंडविड्थ को स्केल करने के बजाय घटकों के लिए अनुकूल कंप्यूट-कम्युनिकेशन अनुपात को लक्षित करना चाहिए—MoE प्रशिक्षण को कुशलतापूर्वक स्केल करने के लिए एक नया दिशा प्रदान करती है। सिस्टम Megatron-LM पर कार्यान्वित है और इसे उन मेमोरी बाधाओं के बिना सैकड़ों अरबों पैरामीटर वाले मॉडलों को प्रशिक्षित करने के लिए एक व्यवहार्य समाधान के रूप में दिखाया गया है जो पिछले विसंयोजित प्रशिक्षण प्रयासों में देखी गई थीं।