कुशल और पोर्टेबल मिक्सचर-ऑफ़-एक्सपर्ट्स संचार

पोर्टेबल मिक्सचर-ऑफ़-एक्सपर्ट्स (MoE) संचार का अवलोकन, जो GPU समानांतरता (parallelism) के अनुकूलन और बड़े पैमाने के AI मॉडल में विलंबता (latency) को कम करने पर केंद्रित है

लेखकPerplexity Research

हम मिक्सचर-ऑफ़-एक्सपर्ट्स (MoE) संचार के लिए एक उच्च-प्रदर्शन वाली, पोर्टेबल, ओपन-सोर्स लाइब्रेरी प्रस्तुत करते हैं, जो मानक ऑल-टू-ऑल संचार प्रिमिटिव्स की तुलना में 10 गुना तेज़ प्रदर्शन प्राप्त करती है।

हमारा कार्यान्वयन कई प्रमुख तकनीकी नवाचारों को शामिल करता है जो बेहतर MoE संचार दक्षता प्रदान करते हैं:

GPU-आरंभित संचार (IBGDA): प्रत्यक्ष GPU-टू-NIC संचार का समर्थन करता है, जो CPU की संलिप्तता को दरकिनार कर विलंबता को काफी कम करता है

संचार और गणना का ओवरलैप: अलग-अलग सेंड और रिसीव चरणों के साथ स्प्लिट कर्नल आर्किटेक्चर, नेटवर्क ट्रांसफर के साथ-साथ गणना को समवर्ती रूप से जारी रखने में सक्षम बनाता है

सबसे तेज़ सिंगल-नोड प्रदर्शन: सिंगल-नोड कॉन्फ़िगरेशन पर पहले के सबसे तेज़ कार्यान्वयन की तुलना में 2.5 गुना कम विलंबता

कुशल और पोर्टेबल मल्टी-नोड प्रदर्शन: हमारा कार्यान्वयन मानक ऑल-टू-ऑल संचार की तुलना में 10 गुना तेज़ गति प्राप्त करता है। यद्यपि यह अत्यधिक विशिष्ट कार्यान्वयनों की तुलना में लगभग 2 गुना धीमा है, हमारा दृष्टिकोण NVSHMEM संस्करणों और नेटवर्क वातावरणों (NVLink, CX-7, और EFA) के बीच बेहतर पोर्टेबिलिटी प्रदान करता है

यह लाइब्रेरी पूर्णतः ओपन-सोर्स है और https://github.com/ppl-ai/pplx-kernels पर उपलब्ध है।

इस लेख में, हम बड़े पैमाने के MoE मॉडल में एक्सपर्ट समानांतरता की चुनौतियों का पता लगाते हैं और वितरित GPU वातावरण में कुशल टोकन डिस्पैच और संयोजन के लिए अपने दृष्टिकोण का वर्णन करते हैं।

परिचय

मिक्सचर-ऑफ़-एक्सपर्ट्स (MoE) मॉडल, जैसे कि DeepSeek R1 और Mixtral 8x7B, प्रत्येक टोकन के लिए सक्रिय किए गए वेट्स (weights) की संख्या को सीमित करके घने (dense) मॉडलों में सुधार करते हैं। प्रत्येक डिकोडर लेयर के मल्टी-लेयर परसेप्ट्रॉन (MLP) में लेयर्स की संख्या बढ़ाकर या लीनियर लेयर्स का आकार बढ़ाकर उच्च पैरामीटर संख्या प्राप्त करने के बजाय, MoE मॉडल पारंपरिक MLP को कई एक्सपर्ट्स और एक राउटर के साथ प्रतिस्थापित करते हैं। उदाहरण के लिए, DeepSeek R1 के 671B पैरामीटर्स में से, केवल 37B पैरामीटर्स ही इन्फरेंस (inference) के दौरान किसी दिए गए टोकन के साथ गुणा किए जाते हैं। यह एक टोकन को डिकोड करने के लिए आवश्यक गणना की मात्रा को काफी कम कर देता है और घने मॉडल की तुलना में विलंबता में सुधार करता है।

घने मॉडलों की तुलना में, MoE मॉडल इन्फरेंस के लिए कुछ अतिरिक्त चुनौतियाँ प्रस्तुत करते हैं। जबकि एक्सपर्ट्स स्वयं छोटे MLP लेयर्स होते हैं, प्रत्येक डिकोडर लेयर में एक राउटर शामिल होता है जो यह तय करता है कि टोकन किन एक्सपर्ट्स को भेजा जाएगा, जहाँ प्रत्येक टोकन को कई एक्सपर्ट्स को भेजा जाता है। राउटर आमतौर पर एक छोटी लीनियर लेयर होती है जो प्रायिकता वितरण (probability distribution) उत्पन्न करती है। आमतौर पर शीर्ष-K स्कोर वाले एक्सपर्ट्स को चुना जाता है और अंतिम सक्रियण (activation) की गणना भारित औसत के रूप में की जाती है, जिसमें राउटर के प्रायिकता वितरण से प्राप्त भार (weight) के साथ एक्सपर्ट आउटपुट को गुणा करके जोड़ा जाता है।

वितरित MoE प्रणालियों में विलंबता को कम करने के लिए, कई डिवाइसों में समानांतरता का लाभ उठाया जा सकता है, लेकिन यह संचार संबंधी चुनौतियाँ उत्पन्न करता है। Mixtral 8x7B या Llama-70B जैसे मॉडल एक ही नोड के 8 डिवाइसों के भीतर फिट हो जाते हैं, जो तेज़ और कम विलंबता वाले NVLink इंटरकनेक्शन (900Gbps तक) से लाभान्वित होते हैं। हालाँकि, बड़े मॉडलों को InfiniBand (जो 400Gbps पर चरम पर होता है) का उपयोग करके कई नोड्स में एक्सपर्ट्स को शार्ड करने की आवश्यकता होती है, जो अतिरिक्त विलंबता चुनौतियाँ पेश करता है।

मिश्रण-विशेषज्ञ परत में टोकन रूटिंग, विशेषज्ञ गणना, और आउटपुट संग्रहण।

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

शार्डिंग और समानांतरता

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

दो GPUs में विशेषज्ञ समानांतरता।

एक्सपर्ट समानांतरता (EP), जिसे पहले चित्र में दर्शाया गया है, केवल एक्सपर्ट गणना को समानांतर करती है। अलग-अलग एक्सपर्ट्स को अलग-अलग डिवाइसों को सौंपा जाता है, जो उनके वेट्स को रखते हैं। रूटिंग के बाद, टोकन संबंधित डिवाइस पर भेजे जाते हैं, और परिणाम बाद में एकत्र और संचित किए जाते हैं। रूटिंग की जटिलता मॉडल के अन्य भागों में समानांतरता की डिग्री पर निर्भर करती है: अन्य लेयर्स को दोहराना रूटिंग की आवश्यकता को पूरी तरह से समाप्त कर सकता है, क्योंकि प्रत्येक रैंक स्थानीय रूप से प्रतिकृति (replicated) रूटिंग तालिका से टोकन का चयन कर सकता है। हालाँकि, यदि केवल एक रैंक रूटिंग चलाता है, तो टोकन, इंडेक्स और वेट्स को उनके संबंधित एक्सपर्ट्स तक भेजने के लिए ब्रॉडकास्ट की आवश्यकता होती है। अंत में, एक ऑल-गेदर या ऑल-टू-ऑल ब्रॉडकास्ट आउटपुट टोकन को उस रैंक के साथ सिंक्रनाइज़ करता है जो मॉडल का निष्पादन जारी रखता है। ऐसा कार्यान्वयन अपेक्षाकृत सरल है क्योंकि टॉर्च (torch) आवश्यक प्रिमिटिव्स को पहले ही प्रदर्शित करता है, हालाँकि एक्सपर्ट्स से टोकन एकत्र करने वाले गेदर ऑपरेशन के साथ एक्यूमुलेशन (accumulation) को फ्यूज करने से कुछ लाभ प्राप्त हो सकते हैं।

दो GPUs में टेंसर समानांतरता के साथ विशेषज्ञ समानांतरता का संयोजन।

एक्सपर्ट-ओनली समानांतरता आदर्श रूप से स्केल नहीं होती है, क्योंकि जब मॉडल अटेंशन, नॉर्म और सैंपलिंग जैसे नॉन-एक्सपर्ट लेयर्स को चला रहा होता है तो क्लस्टर में नोड्स निष्क्रिय हो सकते हैं। हालाँकि, इन लेयर्स की गणना, मुख्य रूप से अटेंशन, भी टेंसर समानांतरता (TP) से लाभान्वित हो सकती है। अधिकांश मॉडल मल्टी-हेड अटेंशन पर निर्भर करते हैं, जिसका अर्थ है कि अटेंशन हेड्स और उनके संबंधित Q, K और V प्रोजेक्शन को भी डिवाइसों में शार्ड किया जा सकता है, जो विभिन्न लेयर्स के बीच स्लाइस को दोहराते (replicating) या एकत्र (gathering) करते हैं। यदि अटेंशन सभी डिवाइसों में फैली हुई है, तो एक ऑल-गेदर सक्रियणों को सिंक्रनाइज़ कर सकता है, जिससे रूटिंग को दोहराया जा सकता है, जिसके लिए एक कुशल कार्यान्वयन के लिए एक्सपर्ट समानांतरता मामले के समान सिंक्रनाइज़ेशन प्रिमिटिव्स की आवश्यकता होती है। हालाँकि, इस स्तर पर समानांतरता की सीमाएँ हैं, क्योंकि अटेंशन हेड्स की संख्या को एक निश्चित सीमा से नीचे कम करने से प्रतिफल कम (diminishing returns) प्राप्त होंगे।

दो GPUs में अलग-अलग अनुरोधों के लिए डेटा समानांतरता के साथ विशेषज्ञ समानांतरता का संयोजन।

सभी डिवाइसों का सर्वोत्तम उपयोग करने और 128 या 256 GPU तक एक्सपर्ट समानांतरता की एक बहुत उच्च डिग्री का समर्थन करने के लिए, डेटा समानांतरता (DP) की आवश्यकता होती है। इस योजना के तहत, जिन डिवाइसों में एक मॉडल विभाजित होता है, उन्हें समवर्ती रूप से अनुरोधों को संभालने के लिए समूहीकृत किया जाता है, जो अटेंशन की गणना करते हैं और उनके स्थानीय समूह में शार्ड किए गए KV कैश को बनाए रखते हैं। इन समानांतर समूहों के कई उदाहरण एक्सपर्ट मूल्यांकन पर सहयोग करते हैं, जिनमें से प्रत्येक एक्सपर्ट वेट्स का एक अलग सबसेट होस्ट करता है। अटेंशन हेड्स की संख्या के आधार पर, एक समूह आमतौर पर एक पूरे नोड के आकार तक फैल सकता है, क्योंकि इंट्रा-नोड संचार तेज़ होता है। उदाहरण के लिए, ऊपर दिए गए चित्र में, एक DP रैंक स्वतंत्र रूप से दो अनुरोधों, A और B की सेवा करता है, जो अनुरोधों के लिए अटेंशन, नॉर्म और किसी भी अन्य बुककीपिंग को संभालता है। दूसरा DP रैंक एक अलग अनुरोध C को संसाधित करता है। हालाँकि, पहला नोड आधे एक्सपर्ट्स को होस्ट करता है, जबकि दूसरा नोड बाकी आधे को, इसलिए रूटिंग के बाद A और B से टोकन दूसरे नोड पर भेजे जा सकते हैं और इसके विपरीत। यह एक विरल (sparse) संचार समस्या की ओर ले जाता है: प्रत्येक डिवाइस किसी अन्य गंतव्य रैंक पर टोकन की एक अलग संख्या भेज सकता है। टॉर्च के मौजूदा प्रिमिटिव्स, मुख्य रूप से all_to_all, विशेष रूप से उपयुक्त नहीं हैं, क्योंकि उन्हें पैडिंग या GPU-टू-CPU सिंक्रनाइज़ेशन और मेटाडेटा ब्रॉडकास्ट के किसी रूप की आवश्यकता हो सकती है। संचार को प्रभावी ढंग से लागू करने के लिए, रूटिंग जानकारी को संसाधित करने और CPU हस्तक्षेप के बिना डिवाइसों से संचार शुरू करने के लिए कस्टम कर्नल्स की आवश्यकता होती है। रूटिंग के बाद, एक डिस्पैच कर्नल को उन रैंकों पर टोकन भेजने चाहिए जहाँ उन्हें रूट किया गया था, जबकि कंबाइन पक्ष पर वर्तमान DP समूह में अनुरोधों से संबंधित सक्रियणों को एकत्र किया जाना चाहिए। इसके अतिरिक्त, ब्रॉडकास्ट को कम करने और बाद के लेयर्स के लिए सभी रैंकों पर टोकन को सही ढंग से दोहराने के लिए, एक DP समूह के भीतर काम को संतुलित किया जाना चाहिए जिसमें कई डिवाइस हो सकते हैं।

NVSHMEM

NVSHMEM एक NVIDIA-विशिष्ट OpenSHMEM कार्यान्वयन है, जो पोर्टेबल इंटर-डिवाइस संचार सुविधाएं प्रदान करता है जो अंतर्निहित हार्डवेयर की जटिलता को दूर करता है। API डिवाइस-टू-डिवाइस रीड और राइट को व्यक्त कर सकता है, जिन्हें व्यक्तिगत ट्रांसपोर्ट लेयर्स के प्रिमिटिव्स पर मैप किया जाता है। NVLink और RDMA दोनों समर्थित हैं, जो पोर्टेबिलिटी की एक डिग्री प्रदान करते हैं और NVSHMEM कर्नल्स को एक नोड के भीतर या कई नोड्स के क्लस्टर में संचालित करने की अनुमति देते हैं। हमारे कर्नल्स में, हम डेटा स्थानांतरित करने और भाग लेने वाले डिवाइसों को सिंक्रनाइज़ करने के लिए रिमोट मेमोरी राइट्स और सिग्नलिंग ऑपरेशंस का उपयोग करते हैं।

NVSHMEM ऑपरेशंस सममित मेमोरी (symmetric memory) की अवधारणा के इर्द-गिर्द बने हैं: वे उन बफ़र्स पर संचालित होते हैं जिन्हें इंटर-डिवाइस संचार में भाग लेने वाले सभी डिवाइसों पर आवंटित किया गया है। एक स्थानीय डिवाइस ऑपरेशन के गंतव्य को निर्दिष्ट करने के लिए टारगेट रैंक के इंडेक्स के साथ अपने स्वयं के स्थानीय बफ़र के भीतर प्राप्त पते का उपयोग कर सकता है। नीचे दिया गया चित्र इस अवधारणा को स्पष्ट करता है: दोनों GPU समान आकार के सममित बफ़र्स आवंटित करते हैं, जो उनके src पॉइंटर्स को बनाए रखते हैं। पहली रैंक दूसरे को 3 पूर्णांक (integers) भेजना चाहती है, उन्हें बफ़र की शुरुआत में रखती है: nvshmem_int_put_nbi स्थानीय बफ़र से शुरुआती पता प्राप्त करता है, और टारगेट डिवाइस को निर्दिष्ट करता है। दूसरी रैंक अपने स्वयं के बफ़र से एक ऑफ़सेट प्राप्त करती है, और एक तत्व को पहले डिवाइस पर भेजती है, जिसे एक से ऑफ़सेट किया जाता है। जबकि गंतव्य पते हमेशा nvshmem_alloc का उपयोग करके आवंटित सममित बफ़र्स होने चाहिए, स्रोत बफ़र्स डिवाइस मेमोरी के मनमाने क्षेत्र हो सकते हैं, बशर्ते वे NVSHMEM के साथ पहले से पंजीकृत हों।

जहाँ NVSHMEM प्रिमिटिव्स की एक विस्तृत श्रृंखला प्रदान करता है, हमारे कर्नल केवल 3 फंक्शन्स पर निर्भर करते हैं, उन पर सभी सिंक्रनाइज़ेशन और फेंसिंग का निर्माण करते हैं।

NVSHMEM का उपयोग करके दो GPUs के बीच सममित मेमोरी और रिमोट लेखन।

nvshmemx_putmem_signal_nbi_warp: डेटा के एक ब्लॉक को एक डिवाइस से दूसरे डिवाइस में स्थानांतरित करता है, साथ ही रिमोट डिवाइस पर एक फ्लैग सेट करता है। ऑपरेशन या तो 64-बिट लोकेशन को सेट (NVSHMEM_SIGNAL_SET) करता है या इंक्रीमेंट (NVSHMEM_SIGNAL_ADD) करता है। मेमोरी का पूरा ब्लॉक स्थानांतरित हो जाने के बाद फ्लैग को अपडेट किया जाता है। यदि रिमोट डिवाइस ने फ्लैग में कोई बदलाव देखा है, तो वह सुरक्षित रूप से अपनी मेमोरी में बफ़र तक पहुँच सकता है। यह फंक्शन डेटा ट्रांसफर और सिंक्रनाइज़ेशन को कोलिसिंग (coalescing) के लिए उपयोगी है।

nvshmemx_signal_op एक एकल मेमोरी लोकेशन पर संचालित होता है, आमतौर पर एक 64-बिट फ्लैग, जिसे परमाणु रूप से (atomically) सेट या इंक्रीमेंट करता है। यह मेटाडेटा भेजने और डिवाइसों को सिंक्रनाइज़ करने में उपयोगी है।

nvshmem_uint64_wait_until का उपयोग सिग्नल के प्राप्त अंत पर, एक फ्लैग को पोल (poll) करने के लिए किया जाता है जब तक कि रिमोट उसे अपडेट न कर दे।

nvshmem_quiet या nvshmem_fence के माध्यम से ऑपरेशंस के बीच किसी भी प्रकार के क्रम को सुनिश्चित करना महंगा है, इसलिए कर्नल्स का हमारा कार्यान्वयन किसी अन्य NVSHMEM प्रिमिटिव्स के उपयोग से बचता है और व्यक्तिगत सेंड-रिसीव जोड़ों के अलावा किसी भी निर्भरता को कम करता है। डेटा भेजते समय, हम हमेशा फंक्शन्स के नॉन-ब्लॉकिंग संस्करण का उपयोग करते हैं, बिना डेटा के स्थानीय रैंक से बाहर भेजे जाने की प्रतीक्षा किए। कर्नल्स के समूह में अन्य, निहित सिंक्रनाइज़ेशन यह सुनिश्चित करता है कि ये बफ़र्स विनाशकारी रूप से अधिलेखित (overwritten) न हों जबकि NIC को उन तक पहुँच की आवश्यकता हो।

NVSHMEM GPUDirect RDMA का उपयोग करता है ताकि NIC CPU को शामिल किए बिना डेटा विनिमय के लिए सीधे GPU मेमोरी को पढ़ और लिख सके। इसके अलावा, ConnectX NICs पर, NVSHEMEM GPU-आरंभित संचार (जिसे Infiniband GPUDirect Async, या IBGDA भी कहा जाता है) का समर्थन करता है, जो GPU को सीधे NIC पर RDMA ऑपरेशंस पोस्ट करने और सीधे NIC के कंप्लीशन क्यू (completion queue) को पोल करने की अनुमति देता है। उन प्लेटफार्मों पर जो GPU-आरंभित संचार का समर्थन नहीं करते हैं, NVSHMEM होस्ट CPU पर संचार शुरू करने और GPU की ओर से कंप्लीशन को पोल करने के लिए एक "प्रॉक्सी" थ्रेड उत्पन्न करता है। NVSHMEM प्रोग्राम पोर्टेबल है, चाहे GPU-आरंभित संचार समर्थित हो या नहीं। हालाँकि, GPU-आरंभित संचार विलंबता को काफी कम कर देता है क्योंकि यह CPU के चक्कर को पूरी तरह से दरकिनार कर देता है।

पोर्टेबल कर्नल्स

हम डिस्पैच और कंबाइन कर्नल्स की एक जोड़ी के माध्यम से MoE संचार को लागू करते हैं। डिस्पैच कर्नल्स प्रत्येक रैंक पर टोकन और रूटिंग जानकारी को पढ़ने, और उन्हें उपयुक्त एक्सपर्ट्स तक भेजने के लिए जिम्मेदार हैं। कंबाइन कर्नल्स एक्सपर्ट्स द्वारा उत्पादित सक्रियणों को एकत्र करते हैं और उन्हें उनके स्रोत रैंक पर वापस भेजते हैं, साथ ही राउटर द्वारा गणना किए गए वेट्स के आधार पर चयनित एक्सपर्ट्स से भारित औसत की गणना करते हैं। कर्नल्स को आगे एक सेंड और रिसीव घटक में विभाजित किया गया है, ताकि डेटा ट्रांसफर को गणना के साथ ओवरलैप किया जा सके। सेंड कर्नल्स सभी नॉन-ब्लॉकिंग और नॉन-सिंक्रनाइज़िंग हैं: वे केवल रिमोट्स को सभी राइट्स डिस्पैच करते हैं। दूसरी ओर, कंबाइन कर्नल्स केवल मेमोरी से पढ़ते हैं, और तब तक प्रतीक्षा करते हैं जब तक कि सभी आवश्यक डेटा स्थानांतरित नहीं हो जाता। NICs को काम डिस्पैच करने के बाद, जब डेटा तार पर एसिंक्रोनस रूप से स्थानांतरित किया जाता है, तो GPU स्थानीय रूप से अन्य उपयोगी काम कर सकते हैं, जैसे कि साझा एक्सपर्ट्स को लागू करना या अटेंशन की गणना करना।

कर्नल्स की प्रत्येक जोड़ी के अपने सममित मेमोरी बफ़र्स होते हैं, जिसमें प्रत्येक रैंक अन्य सभी रैंकों से डेटा प्राप्त करने के लिए निजी स्टोरेज आवंटित करता है। एक DP समूह से डिस्पैच किए जा सकने वाले टोकन की अधिकतम संख्या रैंकों के बीच तय की गई है, जो उस ऊपरी सीमा को भी निर्धारित करती है जिसे प्रत्येक रैंक प्रत्येक DP रैंक से प्रत्येक स्थानीय एक्सपर्ट के लिए प्राप्त कर सकता है। यह सेंडर रैंक को गंतव्य रैंक पर लिखने के लिए एक अद्वितीय पता प्राप्त करने की अनुमति देता है, बिना उनके बीच किसी सिंक्रनाइज़ेशन की आवश्यकता के। डेटा प्राप्त होने के बाद, कर्नल्स रिसीविंग कर्नल्स द्वारा आवश्यक इष्टतम प्रारूप में मेमोरी में उन्हें लेआउट करने के लिए टोकन को इधर-उधर शफल करते हैं। जबकि बफ़र्स का एक बड़ा आयाम होता है, उन्हें मॉडल की सभी अनुक्रमिक लेयर्स में फिर से उपयोग किया जाता है।

डिस्पैच और संयोजन ऑपरेशनों के लिए अलग-अलग भेजने और प्राप्त करने वाले कर्नेल।

वैश्विक सिंक्रनाइज़ेशन का एकमात्र रूप, जैसा कि ऊपर चित्र में दर्शाया गया है, कंबाइन-रिसीव कर्नल में लागू किया गया है। एक बार जब डेटा सेंड से डिस्पैच हो जाता है, तो एक रैंक वापस आता है और निष्पादन जारी रखता है जब वह अन्य रैंकों से आवश्यक सभी टोकन प्राप्त करता है और संसाधित होते ही कंबाइन सेंड कर्नल्स में प्रवेश करता है। कंबाइन-रिसीव में बैरियर यह सुनिश्चित करता है कि कोई भी रैंक आगे न बढ़ सके और डिस्पैच-सेंड शुरू न कर सके जब कोई अन्य रैंक अभी भी डेटा प्राप्त करने की प्रतीक्षा कर रहा हो, जिससे संभावित रूप से विनाशकारी ओवरलैपिंग हो सकती है। सरलता के लिए सिंक्रनाइज़ेशन कंबाइन कर्नल में किया जाता है, एक अनुकूलन के रूप में वही बैरियर डिस्पैच रिसीव और कंबाइन सेंड कर्नल्स के बीच लागू किया जा सकता है, जिससे GPU अन्य कर्नल्स चला रहा हो तब NICs को ट्रांसफर को व्यवस्थित करने के लिए अधिक समय देकर विलंबता को और कम किया जा सकता है।

दोनों कर्नल्स डिवाइसों के सभी उपलब्ध SMs में विभाजित हैं: जबकि डिस्पैच सेंड और कंबाइन रिसीव कर्नल्स को प्रति-रैंक अधिकतम टोकन संख्या (max_m) के बीच समानांतर होना चाहिए, डिस्पैच रिसीव और कंबाइन सेंड कर्नल्स को प्रत्येक रैंक और स्थानीय एक्सपर्ट से उतने टोकन संभालने होंगे, अधिकतम max_m * (num_experts // EP) * (EP // TP) के लिए।

डिस्पैच

डिस्पैच कर्नल उन पीयर्स (peers) तक टोकन और ब्लॉक स्केलिंग कारकों को डिस्पैच करने के लिए जिम्मेदार है जो उन एक्सपर्ट्स को होस्ट करते हैं जिन्हें टोकन रूट किए गए थे। सेंडर पक्ष दो कार्यों को समानांतर करने के लिए वार्प स्पेशलाइजेशन (warp specialization) पर निर्भर करता है: यह निर्धारित करने के लिए रूटिंग जानकारी को एकत्रित करना कि प्रत्येक रैंक से कितने टोकन भेजे जाते हैं और रिमोट तक डिस्पैच करने के लिए टोकन को संदेशों में पैक करना। रिसीवर पक्ष पहले सभी टोकन काउंट्स के प्राप्त होने की प्रतीक्षा करता है, फिर यह साझा मेमोरी से पेलोड को पढ़ता है और उन्हें उन टेंसर्स (tensors) में अनपैक करता है जो डाउनस्ट्रीम ग्रुपेड GEMM कर्नल्स को पास किए जाएंगे। यह कंबाइन कर्नल्स के साथ साझा किए गए बफ़र्स में भी जानकारी संग्रहीत करता है, यह इंगित करने के लिए कि प्रत्येक टोकन कहाँ वापस भेजा जाना चाहिए। यह जानकारी आवश्यक है क्योंकि रिसीव कर्नल टोकन को गैर-नियतात्मक क्रम में सन्निहित बफ़र्स (contiguous buffers) में शफल करता है।

स्रोत और गंतव्य रैंक के बीच टोकनों का डिस्पैच करना और उनके आगमन का संकेत देना।

सेंडर भाग में, प्रत्येक ब्लॉक पर 32 थ्रेड्स (threads) का एक एकल वार्प रूटिंग जानकारी पढ़ने और यह गिनने के लिए जिम्मेदार है कि प्रत्येक गंतव्य एक्सपर्ट को इस रैंक से कितने टोकन प्राप्त होने की उम्मीद है, nvshmem_signal_op का उपयोग करके काउंट प्लस वन भेजता है। काउंट को एक से बढ़ाया जाता है, क्योंकि रिमोट एंड पर शून्य से नॉन-शून्य में परिवर्तन एक स्रोत रैंक से काउंट्स की प्राप्ति का संकेत देता है। समानांतर में, शेष वार्प्स सभी ब्लॉकों में समानांतर में सममित मेमोरी में टोकन कॉपी करने के लिए सहयोग करते हैं, सक्रियणों, स्केलिंग कारकों और उनके इंडेक्स को स्थानीय रैंक पर एक सन्निहित टुकड़े में पैक करते हैं। कंबाइन सेंडर के लिए इंडेक्स आवश्यक है ताकि वह पता निर्धारित किया जा सके जहाँ टोकन लिखा जाएगा। अगला, वार्प समूहों के बीच एक बैरियर के माध्यम से यह सुनिश्चित करने के बाद कि सभी डेटा कॉपी हो गया है, वार्प्स फिर से स्वतंत्र रूप से काम करते हैं, प्रत्येक समान बफ़र को समानांतर में एक अलग एक्सपर्ट को भेजता है। टोकन nvshmemx_putmem_signal_nbi_warp का उपयोग करके भेजे जाते हैं, जो रिमोट डिवाइस पर स्थानीय रैंक से भेजे गए टोकन की संख्या को भी परमाणु रूप से बढ़ाता है। एक DP समूह के भीतर, चूँकि प्रत्येक रैंक के पास भेजे जाने वाले टोकन की एक प्रतिकृति होती है, डिस्पैच समान रूप से संतुलित होता है, जिसमें प्रत्येक डिवाइस टोकन का एक सबसेट भेजता है।

रिसीव एंड पर, सभी कर्नल्स पहले सेंडर रैंक के साथ सिंक्रनाइज़ होते हैं, जिसके लिए कुल टोकन काउंट्स के भेजे जाने की प्रतीक्षा की जाती है। बाद में, वे सभी परमाणु रूप से बढ़ाए गए सेंड टोकन काउंट्स के कुल काउंट्स तक व्यवस्थित होने की प्रतीक्षा करते हैं, जो यह संकेत देता है कि स्रोत रैंक से सभी पेलोड भी भेज दिए गए हैं, putmem कॉल के सिमेंटिक्स के लिए धन्यवाद। कर्नल्स सभी ब्लॉकों और थ्रेड्स (threads) में ऑपरेशन को समानांतर करते हुए nvshmem_uint64_wait_until का उपयोग करके काउंट्स पर पोल (poll) करते हैं। इसके बाद, एक क्रॉस-ब्लॉक बैरियर यह सुनिश्चित करता है कि कोई भी ब्लॉक बफ़र्स को तब तक न पढ़े जब तक कि सभी डेटा सही ढंग से प्राप्त न हो जाए। ब्लॉकों में फैला हुआ और परमाणु रूप से बढ़ाए गए काउंटर के माध्यम से सिंक्रनाइज़, टोकन को सममित मेमोरी से नियमित सन्निहित बफ़र्स में कॉपी किया जाता है जो एक्सपर्ट्स को लागू करने वाले डाउनस्ट्रीम कर्नल्स के इनपुट बन जाते हैं। स्रोत रैंक, एक्सपर्ट इंडेक्स और टोकन इंडेक्स अलग से संग्रहीत किए जाते हैं, जो ठीक उसी स्थान को इंगित करते हैं जहाँ कंबाइन कर्नल को सक्रियणों को भेजना होता है। भले ही एक DP समूह के भीतर से टोकन अलग-अलग डिवाइसों से भेजे जाते हैं, वे सभी संबंधित एक्सपर्ट को पास किए जाने के लिए एक साथ समूहीकृत किए जाते हैं।

कंबाइन

कंबाइन कर्नल्स फिर से सेंड और रिसीवर हिस्सों में विभाजित हैं: सेंडर्स अन-क्वांटाइज़्ड 16-बिट सक्रियणों को रिमोट डिवाइसों पर कॉपी करते हैं, रिसीव कर्नल्स डेटा के भेजे जाने की प्रतीक्षा करते हैं और स्थानीय रूप से सभी एक्सपर्ट योगदानों के भारित औसत की गणना करते हैं। इसके अतिरिक्त, वे एक बैरियर के रूप में भी कार्य करते हैं, डिस्पैच-कंबाइन अनुक्रम को सिंक्रनाइज़ करने के लिए: प्रत्येक रैंक सेंड कर्नल्स में प्रवेश पर प्रत्येक पीयर पर एक फ्लैग सेट करता है, जिसमें रिसीव कर्नल्स को तब तक वापस जाने की अनुमति नहीं होती है जब तक वे फ्लैग सेट न देख लें। सिंक्रनाइज़ेशन की विलंबता न्यूनतम है, क्योंकि यह वास्तविक संचार और गणना के साथ ओवरलैप होती है।

प्रत्येक डेटा-समानांतर रैंक पर टोकनों में विशेषज्ञ आउटपुट्स को संयोजित करना।

सेंडर पक्ष पर, कर्नल्स समानांतर में सभी स्थानीय एक्सपर्ट्स को सौंपे गए टोकन की सूची को पार (traverse) करते हैं, उन्हें गंतव्य रैंक पर एक बफ़र में लिखते हैं। टारगेट रैंक, एक्सपर्ट इंडेक्स और टोकन इंडेक्स को टोकन की प्राप्ति पर स्केटर (scatter) कर्नल्स द्वारा पॉप्युलेट किए गए प्रति-टोकन बफ़र्स से पढ़ा जाता है। प्रत्येक सेंडर के पास लिखने के लिए प्रति एक्सपर्ट अपना निजी मेमोरी क्षेत्र होता है, जैसा कि ऊपर चित्र में दर्शाया गया है, जिससे सिंक्रनाइज़ करने की आवश्यकता से बचा जा सके। डिस्पैच के समान, कंबाइन गंतव्य रैंक पर प्रति टोकन काउंटर्स को परमाणु रूप से बढ़ाता है ताकि डेटा की प्राप्ति का संकेत मिल सके: जब काउंटर उन एक्सपर्ट्स की संख्या से मेल खाता है जिन्हें एक टोकन डिस्पैच किया गया था, तो टोकन सामग्री तक पहुँचा जा सकता है।

रिसीव कर्नल में, टोकन की सूची को कई ब्लॉकों में समानांतर में पार किया जाता है, सिग्नलिंग ऑपरेशन द्वारा सेट फ्लैग को पोल (poll) करके उनकी सामग्री के आने की प्रतीक्षा की जाती है। आगमन पर, पेलोड को निजी बफ़र्स से पढ़ा जाता है, जिसमें रूटिंग तालिका यह इंगित करती है कि किस बफ़र से पढ़ना है और प्रत्येक एक्सपर्ट को क्या भार (weight) निर्दिष्ट करना है। परिणाम तब बाहरी रूप से आवंटित टेंसर्स में लिखे जाते हैं, जिसमें कर्नल सभी डिवाइसों के बैरियर पास करने के बाद निष्पादन समाप्त करता है। एक DP समूह में, सभी रैंक अपनी खुद की प्रतिकृतियां (replicas) की गणना करने के लिए प्रत्येक एक्सपर्ट सक्रियण की एक कॉपी प्राप्त करते हैं।

बेंचमार्क्स

हम InfiniBand और CX-7 NICs के माध्यम से जुड़े 16 नोड्स में फैले 128 GPU के क्लस्टर पर अपने कर्नल्स का मूल्यांकन करते हैं। हम CPU प्रॉक्सी के साथ GPUDirect Async (IBGDA) और रिलायबल कनेक्शन (RC) दोनों का प्रयास करते हैं। हम उनकी तुलना अत्यधिक अनुकूलित DeepSeek कार्यान्वयन, और PyTorch (NCCL के माध्यम से) या NVSHMEM द्वारा प्रदान किए गए घने (dense) प्रिमिटिव्स से करते हैं।

मल्टी-नोड EP128 कॉन्फ़िगरेशन में प्रति GPU 128 टोकनों के लिए संचार विलंबता।

हालाँकि घने NVSHMEM ऑल-टू-ऑल ऑपरेशन पर, IBRC और IBGDA का प्रदर्शन समान है (6378 µs बनाम 6180 µs), IBGDA विरल (sparse) कर्नल्स के साथ काफी तेज़ है। डिस्पैच और कंबाइन को जोड़ने पर, IBGDA 902 µs का उपयोग करता है जबकि IBRC 3223 µs लेता है - विलंबता में 3.6 गुना सुधार। जबकि ऑल-टू-ऑल बैंडविड्थ बाधित है, विरल कर्नल परिमाण के आदेश (orders of magnitudes) कम डेटा प्रसारित करते हैं, जो विलंबता से बाधित होते हैं। सीधे GPU से नेटवर्क ट्रांसफर को ट्रिगर करके, GPU और NIC को समन्वित करने के लिए CPU प्रॉक्सी की आवश्यकता के बिना, एंड-टू-एंड विलंबता काफी कम हो जाती है।

हालाँकि हमारे पोर्टेबल कर्नल्स अत्यधिक अनुकूलित DeepSeek कर्नल्स की तुलना में लगभग 2 गुना धीमे हैं, वे घने कर्नल्स की तुलना में विलंबता में 10 गुना सुधार करते हैं (902 µs बनाम 9944 µs)।

इसके अतिरिक्त, सेंडर-रिसीवर घटकों में विभाजन कुछ विलंबता को छिपाने की अनुमति देता है, लाइब्रेरी प्रिमिटिव्स के विपरीत।

सिंगल-नोड EP8 कॉन्फ़िगरेशन में प्रति GPU 128 टोकनों के लिए संचार विलंबता।

सिंगल-नोड (EP8) पर, NVSHMEM परिवहन के लिए NVLINK का उपयोग करता है, जो इंटर-नोड नेटवर्किंग की तुलना में कम विलंबता और उच्च थ्रूपुट प्रदान करता है। हमारे पोर्टेबल कर्नल्स सिंगल-नोड पर DeepEP की तुलना में लगभग 2.5 गुना तेज़ हैं (186µs बनाम 481 µs)।

आगे के सुधार

यहाँ वर्णित कर्नल्स ML फ्रेमवर्क्स के इन-बिल्ट प्रिमिटिव्स से बेहतर प्रदर्शन करते हैं और विशेष इंटर-डिवाइस ट्रांसपोर्ट, जैसे NVLink या InfiniBand के लिए अत्यधिक-विशिष्ट (over-specializing) हुए बिना उचित प्रदर्शन प्रदान करते हैं। पहले से उल्लिखित अवसरों के अलावा, संचार प्रिमिटिव्स को अधिक विशिष्ट संस्करणों के साथ प्रतिस्थापित करके आगे के प्रदर्शन लाभ प्राप्त किए जा सकते हैं। उदाहरण के लिए, NVLink पर, सममित मेमोरी के उपयोग को डिवाइसों के बीच साझा बफ़र्स के साथ प्रतिस्थापित किया जा सकता है, कुछ कॉपी करने को समाप्त किया जा सकता है और टोकन बैचों के बजाय व्यक्तिगत टोकन के स्तर पर अधिक बारीक सिंक्रनाइज़ेशन योजनाओं की अनुमति दी जा सकती है। InfiniBand पर, एक कार्यान्वयन सीधे अंतर्निहित क्यू पेयर्स (queue pairs) तक पहुँच सकता है, जिससे अपेक्षित और भेजे गए टोकन काउंट्स दोनों को स्थानांतरित करने की आवश्यकता समाप्त हो जाती है और क्यू से आउटपुट बफ़र्स में टोकन को फ्लश किया जा सकता है जैसे ही वे आते हैं। हालाँकि, ऐसे कार्यान्वयन पोर्टेबिलिटी की कीमत पर आते हैं, जबकि हमारे कार्यान्वयन विशिष्ट हार्डवेयर के लिए अनुकूलन करने के लिए विकास प्रयास करने से पहले विशिष्ट कॉन्फ़िगरेशन वाले क्लस्टर्स का मूल्यांकन करने की अनुमति देते हैं।

निष्कर्ष

हमने MoE संचार के लिए एक उच्च-प्रदर्शन, पोर्टेबल लाइब्रेरी प्रस्तुत की है जो विविध हार्डवेयर कॉन्फ़िगरेशन में संगतता बनाए रखते हुए मानक ऑल-टू-ऑल संचार की तुलना में 10 गुना तेज़ प्रदर्शन प्राप्त करती है। NVLink के साथ सिंगल-नोड परिनियोजन (deployments) पर, हमारा समाधान पिछले कार्यान्वयनों की तुलना में 2.5 गुना कम विलंबता प्रदर्शित करता है।

हमारा दृष्टिकोण प्रमुख नवाचारों के माध्यम से पोर्टेबिलिटी के साथ प्रदर्शन को संतुलित करता है, जिसमें GPU-आरंभित संचार समर्थन, एक स्प्लिट कर्नल आर्किटेक्चर जो गणना-संचार ओवरलैप को सक्षम बनाता है, और न्यूनतम NVSHMEM प्रिमिटिव्स का उपयोग करके कुशल टोकन डिस्पैच शामिल है। मल्टी-नोड सेटअप पर अत्यधिक विशिष्ट कार्यान्वयनों की तुलना में लगभग 2 गुना धीमा होने के बावजूद, हमारी लाइब्रेरी विभिन्न नेटवर्क वातावरणों (NVLink, CX-7, और EFA) में बेहतर लचीलापन प्रदान करती है।

जैसे-जैसे MoE मॉडल स्केल करना जारी रखते हैं, हमारे जैसी कुशल संचार रणनीतियां व्यावहारिक परिनियोजन के लिए तेजी से महत्वपूर्ण हो जाएंगी। हमारा पूर्णतः ओपन-सोर्स कार्यान्वयन https://github.com/ppl-ai/pplx-kernels पर उपलब्ध है।

संदर्भ

DeepEP

DeepSeek-V3 तकनीकी रिपोर्ट

NVIDIA NVSHMEM

उद्धरण

BibTeX
@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}