कार्यक्षम आणि पोर्टेबल मिक्चर-ऑफ-एक्सपर्ट्स कम्युनिकेशन

आम्ही मिक्चर-ऑफ-एक्सपर्ट्स (MoE) कम्युनिकेशनसाठी एक उच्च-कार्यक्षमता, पोर्टेबल, ओपन-सोर्स लायब्ररी सादर करतो जी मानक ऑल-टू-ऑल कम्युनिकेशन प्रिमिटिव्ह्जच्या तुलनेत १० पट वेगवान कार्यक्षमता साध्य करते.

आमच्या अंमलबजावणीमध्ये अनेक प्रमुख तांत्रिक नवकल्पना आहेत ज्या उत्कृष्ट MoE कम्युनिकेशन कार्यक्षमता देतात:

GPU-प्रारंभ केलेले कम्युनिकेशन (IBGDA): थेट GPU-ते-NIC कम्युनिकेशनला समर्थन देते, सीपीयू सहभाग टाळून लेटन्सी मोठ्या प्रमाणात कमी करते

संचार आणि संगणकीय आच्छादन: स्वतंत्र पाठवणे आणि प्राप्त करणे टप्प्यांसह विभाजित कर्नल आर्किटेक्चर नेटवर्क स्थानांतरांच्या समवेत संगणक कार्य एकाच वेळी सुरू ठेवण्यास सक्षम करते.

सर्वात वेगवान एकल-नोड कार्यप्रदर्शन: एकल-नोड कॉन्फिगरेशनवरील पूर्वीच्या सर्वात वेगवान अंमलबजावणीपेक्षा 2.5 पट कमी लेटन्सी.

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

ही लायब्ररी पूर्णपणे ओपन-सोर्स आहे आणि https://github.com/ppl-ai/pplx-kernels येथे उपलब्ध आहे.

या लेखामध्ये, आम्ही मोठ्या प्रमाणावरील MoE मॉडेल्समधील तज्ञ समांतरतेच्या आव्हानांचे अन्वेषण करतो आणि वितरीत GPU वातावरणात कार्यक्षम टोकन डिस्पॅच आणि संयोजनासाठी आमचा दृष्टिकोन वर्णवितो.

परिचय

DeepSeek R1 आणि Mixtral 8x7B सारखी मिक्सचर-ऑफ-एक्सपर्ट्स (MoE) मॉडेल्स, प्रत्येक टोकनसाठी सक्रिय केलेल्या वजनांची संख्या मर्यादित करून दाट मॉडेल्समध्ये सुधारणा करतात. प्रत्येक डीकोडर लेयरच्या मल्टी-लेयर पर्सेपट्रॉन (MLP) मधील थरांची संख्या वाढवून किंवा रेखीय थरांचा आकार वाढवून उच्च पॅरामीटर संख्या साध्य करण्याऐवजी, MoE मॉडेल्स पारंपारिक MLP ला अनेक तज्ञांसह आणि राउटरसह बदलतात. उदाहरणार्थ, DeepSeek R1 च्या 671B पॅरामीटर्सपैकी, इनफरन्स दरम्यान दिलेल्या टोकनसह फक्त 37B गुणले जातात. यामुळे टोकन डीकोड करण्यासाठी आवश्यक संगणकाचे प्रमाण मोठ्या प्रमाणात कमी होते आणि दाट मॉडेलच्या तुलनेत लेटन्सी सुधारते.

दाट मॉडेल्सच्या तुलनेत MoE मॉडेल्स इनफरन्ससाठी काही अतिरिक्त आव्हाने प्रस्तुत करतात. तज्ञ स्वतः लहान MLP लेयर असताना, प्रत्येक डीकोडर लेयरमध्ये एक राउटर समाविष्ट असतो जो टोकन कोणत्या तज्ञांकडे पाठवला जावा हे ठरवतो, प्रत्येक टोकन अनेक तज्ञांकडे पाठवला जातो. राउटर सामान्यत: संभाव्यता वितरण तयार करणारा एक लहान रेखीय लेयर असतो. सहसा टॉप-K स्कोअर असलेले तज्ञ निवडले जातात आणि अंतिम सक्रियकरण भारित सरासरी म्हणून मोजले जाते, राउटरच्या संभाव्यता वितरणापासून व्युत्पन्न केलेल्या वजनाने गुणित तज्ञ आउटपुट जोडले जातात.

वितरीत MoE प्रणालींमध्ये लेटन्सी कमी करण्यासाठी, अनेक डिव्हाइसेसवर समांतरता शोषली जाऊ शकते, परंतु यामुळे संवादाची आव्हाने निर्माण होतात. Mixtral 8x7B किंवा Llama-70B सारखी मॉडेल्स एकाच नोडमधील 8 डिव्हाइसेसमध्ये बसतात, जे जलद आणि कमी-लेटन्सी NVLink इंटरकनेक्शनचा (900Gbps पर्यंत) फायदा घेतात. तथापि, मोठ्या मॉडेल्सना InfiniBand (जास्तीत जास्त 400Gbps) वापरून अनेक नोड्समध्ये तज्ञ शार्ड करणे आवश्यक असते, ज्यामुळे अतिरिक्त लेटन्सी आव्हाने निर्माण होतात.

विस्तृत करा

या लेखामध्ये, आम्ही तज्ञ समांतरतेच्या समस्येचा शोध घेतो, GPU कर्नलच्या आमच्या अंमलबजावणीचे वर्णन करतो जे नियुक्त तज्ञ धारण करणाऱ्या डिव्हाइसेसवर टोकन कार्यक्षमतेने पाठवतात आणि तज्ञ आउटपुटची भारित बेरीज मोजून त्यांना परत गोळा करतात. विशिष्ट नेटवर्क हार्डवेअरला लक्ष hकरणार्‍या आक्रमकपणे ऑप्टिमाइझ केलेल्या अंमलबजावणी अस्तित्वात असताना, आम्ही अत्यंत पोर्टेबल अंमलबजावणी सादर करतो जी NVSHMEM प्रिमिटिव्हच्या किमान संचावर अवलंबून असते तरीही मानक सर्व-ते-सर्व संवादापेक्षा 10 पट वेगवान कार्यप्रदर्शन वितरीत करते. आमची अंमलबजावणी विविध नेटवर्क वातावरणात उत्कृष्ट पोर्टेबिलिटी राखून एकल-नोड कॉन्फिगरेशनवर अत्याधुनिक कार्यप्रदर्शन साध्य करते.

शार्डिंग आणि समांतरता

कार्यक्षम इनफरन्ससाठी, मॉडेलचे वजन डिव्हाइस मेमरीमध्ये ठेवणे आवश्यक आहे, त्याच वेळी मॉडेलच्या माध्यमातून फॉरवर्ड पाससाठी आवश्यक असलेले सक्रियकरण, KV कॅशे आणि इतर बफरसाठी पुरेशी स्पेस सोडली पाहिजे. सर्वात सक्षम मॉडेल्स सर्वात सक्षम GPU च्या क्षमतेपेक्षा जास्त असतात, त्यामुळे इनफरन्स अनेक डिव्हाइसेसवर पसरविला जावा, जे सामूहिकपणे मॉडेलचे वजन साठवू शकतात. वजनाच्या शार्डिंग योजनांच्या आधारे, संगणक कार्य ऑफलोड करण्यासाठी आणि डिव्हाइस सिंक्रोनाइझ करण्यासाठी विविध संवाद आणि संगणक योजना वापरल्या पाहिजेन्.

विस्तृत करा

एक्सपर्ट पॅरेललिझम (EP), पहिल्या आकृतीमध्ये चित्रित केले आहे, केवळ तज्ञ संगणक कार्यास समांतर करते. वेगवेगळ्या तज्ञांना वेगवेगळ्या डिव्हाइसेसना नियुक्त केले जाते, जे त्यांचे वजन धरून ठेवतात. रूटिंग नंतर, टोकन संबंधित डिव्हाइसवर पाठवले जातात, परिणाम गोळा केले जातात आणि नंतर जमा केले जातात. मॉडेलच्या इतर भागांमधील समांतरतेच्या अक्षावर रूटिंगची गुंतागुंत अवलंबून असते: इतर लेयरची प्रतिकृती बनवणे पूर्णपणे रूटिंगची गरज दूर करू शकते, कारण प्रत्येक रँक स्थानिक पातळीवर प्रतिकृती बनवलेल्या रूटिंग टेबलमधून टोकन निवडू शकते. तथापि, जर रँकमधील फक्त एक जण रूटिंग चालवत असेल, तर टोकन, अनुक्रमणिका आणि वजन त्यांच्या संबंधित तज्ञांकडे पाठवण्यासाठी ब्रॉडकास्ट आवश्यक आहे. शेवटी, ऑल-गेदर किंवा ऑल-टू-ऑल ब्रॉडकास्ट आउटपुट टोकन ज्या रँकसह मॉडेलची अंमलबजावणी चालू ठेवतो त्याच्याशी सिंक्रोनाइझ करतो. अशी अंमलबजावणी तुलनेने सोपी आहे कारण टॉर्च आधीपासून आवश्यक प्रिमिटिव्ह उघड करतो, जरी तज्ञांकडून टोकन गोळा करणाऱ्या गेदर ऑपरेशनसह एकत्रीकरण फ्यूज करून काही फायदे मिळू शकतात.

विस्तृत करा

तज्ञ-केवळ समांतरता आदर्शपणे स्केल करत नाही, कारण मॉडेल अटेंशन, नॉरम आणि सॅम्पलिंग सारखे तज्ञ-नसलेले लेयर चालू असताना क्लस्टरमधील नोड निष्क्रिय असू शकतात. तथापि, या थरांची गणना, प्रामुख्याने अटेंशन, देखील Tensor Parallelism (TP) द्वारे फायदा घेऊ शकते. बहुतेक मॉडेल्स मल्टी-हेड अटेंशनवर अवलंबून असतात, याचा अर्थ असा की अटेंशन हेड आणि त्यांचे संबंधित Q, K आणि V प्रक्षेपण डिव्हाइसेसवर शार्ड केले जाऊ शकतात, वेगवेगळ्या थरांमधील स्लाईसची प्रतिकृती बनवू शकतात किंवा गोळा करू शकतात. जर अटेंशन सर्व डिव्हाइसेसवर पसरले असेल, तर ऑल-गेदर सक्रियकरण सिंक्रोनाइझ करू शकते, ज्यामुळे रूटिंगची प्रतिकृती बनवता येते, कार्यक्षम अंमलबजावणीसाठी तज्ञ समांतरता केस प्रमाणेच सिंक्रोनाइझेशन प्रिमिटिव्हची आवश्यकता असते. तथापि, या स्तरावर समांतरतेला मर्यादा आहेत, कारण अटेंशन हेडची संख्या एका निश्चित उंबरठ्यापेक्षा कमी केल्याने कमी होणारे रिटर्न मिळतील.

विस्तृत करा

सर्व डिव्हाइसेसचा सर्वोत्तम वापर करण्यासाठी आणि 128 किंवा 256 GPU पर्यंत, तज्ञ समांतरतेची उच्च पदवी समर्थन करण्यासाठी, डेटा समांतरता (DP) आवश्यक आहे. या योजनेनुसार, मॉडेल ज्या डिव्हाइसेसवर विभागले गेले आहे त्यांना विनंत्या समवर्तीपणे हाताळण्यासाठी गटबद्ध केले जाते, अटेंशनची गणना करणे आणि त्यांच्या स्थानिक गटात शार्ड केलेले KV कॅशे राखणे. या समांतर गटांची एकाधिक उदाहरणे तज्ञ मूल्यमापनावर सहकार्य करतात, प्रत्येक तज्ञ वजनाचा एक वेगळा उपसंच होस्ट करतो. अटेंशन हेडच्या संख्येवरून, एखादा गट सामान्यत: संपूर्ण नोडच्या आकारापर्यंत पसरू शकतो, कारण नोडमधील संवाद जलद असतो. उदाहरणार्थ, वरील आकृतीत, एक DP रँक स्वतंत्रपणे दोन विनंत्या, A आणि B सेवा देतो, अटेंशन, नॉरम आणि विनंत्यांसाठी इतर कोणतेही बुककीपिंग हाताळतो. दुसरा DP रँक एक वेगळी विनंती C प्रक्रिया करतो. तथापि, पहिल्या नोडमध्ये अर्धे तज्ञ आहेत, तर इतर नोडमध्ये दुसरे सत्र आहे, अशा प्रकारे A आणि B मधील टोकन रूट केल्यानंतर दुसऱ्या नोडवर पाठवले जाऊ शकतात आणि उलट. यामुळे विरल संवाद समस्या उद्भवते: प्रत्येक डिव्हाइस इतर कोणत्याही गंतव्य रँकवर टोकनची भिन्न संख्या पाठवू शकते. टॉर्चमधील विद्यमान प्रिमिटिव्ह, प्रामुख्याने all_to_all, विशेषतः योग्य नाहीत, कारण त्यांना काही प्रकारच्या पॅडिंग किंवा GPU-से-CPU सिंक्रोनाइझेशन आणि मेटाडेटा ब्रॉडकास्टची आवश्यकता असू शकते. संवाद प्रभावीपणे अंमलात आणण्यासाठी, रूटिंग माहितीवर प्रक्रिया करण्यासाठी आणि CPU हस्तक्षेपाशिवाय डिव्हाइसेसवरून संवाद सुरू करण्यासाठी कस्टम कर्नल आवश्यक आहेत. रूटिंग नंतर, एक डिस्पॅच कर्नलने टोकन ज्या रँकवर रूट केले होते त्यांना पाठवले पाहिजे, तर combinar बाजूला वर्तमान DP गटातील विनंत्यांशी संबंधित सक्रियकरण गोळा केले पाहिजेत. याशिवाय, ब्रॉडकास्ट कमी करण्यासाठी आणि त्यानंतरच्या थरांसाठी सर्व रँकवर टोकन योग्यरित्या प्रतिकृती बनवण्यासाठी एकाधिक डिव्हाइसेस असू शकणाऱ्या DP ग्रुपमध्ये कार्य संतुलित केले पाहिजे.

NVSHMEM

NVSHMEM ही NVIDIA-विशिष्ट OpenSHMEM अंमलबजावणी आहे, जी अंतर्निहित हार्डवेअरची जटिलाता दूर करणारी पोर्टेबल डिव्हाइस-ते-डिव्हाइस संवाद सुविधा प्रदान करते. API डिव्हाइस-ते-डिव्हाइस वाचणे आणि लिहिणे व्यक्त करू शकते, जे वैयक्तिक वाहतूक स्तरांच्या प्रिमिटिव्हशी मॅप केले जातात. NVLink आणि RDMA दोन्ही समर्थित आहेत, पोर्टेबिलिटीची पदवी प्रदान करतात आणि NVSHMEM कर्नलला नोडमध्ये किंवा अनेक नोड्सच्या क्लस्टरमध्ये ऑपरेट करण्यास परवानगी देतात. आमच्या कर्नलमध्ये, आम्ही डेटा हस्तांतरित करण्यासाठी आणि सहभागी डिव्हाइसेस सिंक्रोनाइझ करण्यासाठी रिमोट मेमरी राइट्स आणि सिग्नलिंग ऑपरेशन्स वापरतो.

NVSHMEM ऑपरेशन्स सममितीय मेमरीच्या संकल्पनेभोवती तयार केल्या आहेत: ते डिव्हाइस-ते-डिव्हाइस संवादात सहभागी असलेल्या सर्व डिव्हाइसेसवर वाटप केलेल्या बफर्सवर ऑपरेट करतात. स्थानिक डिव्हाइस ऑपरेशनचे गंतव्यस्थान निर्दिष्ट करण्यासाठी लक्ष्य रँकच्या अनुक्रमणिकासह त्याच्या स्वतःच्या स्थानिक बफरमध्ये व्युत्पन्न केलेल्या पत्त्याचा वापर करू शकते. खालील आकृती ही संकल्पना स्पष्ट करते: दोन्ही GPU समान आकाराचे सममितीय बफर वाटप करतात, त्यांच्याशी src पॉइंटर राखून ठेवतात. पहिल्या रँकला दुसऱ्या रँकला 3 पूर्णांक पाठवायचे आहेत, त्यांना बफरच्या सुरूवातीस ठेवणे: nvshmem_int_put_nbi स्थानिक बफरवरून प्रारंभ पत्ता व्युत्पन्न करते, लक्ष्य डिव्हाइस निर्दिष्ट करते. दुसरा रँक त्याच्या स्वतःच्या बफरमधून ऑफसेट व्युत्पन्न करतो, एका घटकाने ऑफसेट करून पहिल्या डिव्हाइसला एक घटक पाठवतो. गंतव्य पत्ते नेहमी nvshmem_alloc वापरून वाटप केलेले सममितीय बफर असले पाहिजेत, तर स्त्रोत बफर डिव्हाइस मेमरीचे अनियंत्रित प्रदेश असू शकतात, बशर्ते ते NVSHMEM सह पूर्व-नोंदणीकृत असतील.

जरी NVSHMEM विस्तृत श्रेणीचे प्रिमिटिव्ह प्रदान करते, आमचे कर्नल केवळ 3 फंक्शन्सवर अवलंबून असतात, त्यावर सर्व सिंक्रोनाइझेशन आणि फेन्सिंग तयार करतात.

विस्तृत करा

nvshmemx_putmem_signal_nbi_warp: रिमोट डिव्हाइसवर फ्लॅग सेट करताना, डेटाचा एक ब्लॉक एका डिव्हाइसवरून दुसर्‍या डिव्हाइसवर हस्तांतरित करते. ऑपरेशन एकतर 64-बिट स्थान सेट करते (NVSHMEM_SIGNAL_SET) किंवा वाढवते (NVSHMEM_SIGNAL_ADD). मेमरीचा संपूर्ण ब्लॉक हस्तांतरित झाल्यानंतर फ्लॅग अपडेट केला जातो. रिमोट डिव्हाइसला फ्लॅगमधील बदल दिसल्यास, ते त्याच्या स्वतःच्या मेमरीमधील बफर सुरक्षितपणे प्रवेश करू शकते. हे फंक्शन डेटा हस्तांतरण आणि सिंक्रोनाइझेशन एकत्रित करण्यासाठी उपयुक्त आहे.

nvshmemx_signal_op एकाच मेमरी स्थानावर ऑपरेट करते, सामान्यतः 64-बिट फ्लॅग, ते अणुऊर्जेने सेट करणे किंवा वाढवणे. मेटाडेटा पाठवण्यात आणि डिव्हाइसेस सिंक्रोनाइझ करण्यात हे उपयुक्त आहे.

nvshmem_uint64_wait_until रिमोट अपडेट होईपर्यंत फ्लॅग पोल करण्यासाठी, सिग्नलच्या प्राप्त करणाऱ्या टोकाला वापरली जाते.

ऑपरेशन्समधील कोणतेही क्रमबद्धीकरण nvshmem_quiet किंवा nvshmem_fence द्वारे सुनिश्चित करणे महागडे ठरते; म्हणूनच आमची कर्नल्सची अंमलबजावणी इतर कोणत्याही NVSHMEM प्रिमिटिव्ह्जचा वापर टाळते आणि स्वतंत्र send-receive जोड्यांव्यतिरिक्त इतर सर्व अवलंबित्वे कमीतकमी ठेवते. डेटा पाठवताना, आम्ही नेहमी फंक्शन्सची नॉन-ब्लॉकिंग आवृत्ती वापरतो, स्थानिक रँकमधून डेटा प्रत्यक्ष पाठवला जाण्याचीही वाट न पाहता. कर्नल्सच्या समूहामधील इतर, अप्रत्यक्ष समक्रमण हे सुनिश्चित करते की NICs ना त्यांना प्रवेशाची गरज असताना हे बफर्स विध्वंसक रीतीने ओव्हरराइट होत नाहीत.

NVSHMEM GPUDirect RDMA चा वापर करते, ज्यामुळे NIC ला CPU चा सहभाग न घेता डेटा देवाणघेवाणीसाठी GPU मेमरी थेट वाचता आणि लिहिता येते. याव्यतिरिक्त, ConnectX NICs वर, NVSHEMEM GPU-initiated communication (जे Infiniband GPUDirect Async, किंवा IBGDA म्हणूनही ओळखले जाते) ला समर्थन देते, ज्यामुळे GPU थेट NIC कडे RDMA ऑपरेशन्स पोस्ट करू शकते आणि NIC च्या completion queue ला थेट पोल करू शकते. जे प्लॅटफॉर्म GPU-initiated communication ला समर्थन देत नाहीत, त्यावर NVSHMEM होस्ट CPU वर एक “proxy” थ्रेड सुरू करते, जे GPU च्या वतीने communication सुरू करणे आणि completion पोल करणे करते. GPU-initiated communication समर्थित असो वा नसो, NVSHMEM प्रोग्राम पोर्टेबल असतो. तथापि, GPU-initiated communication विलंब लक्षणीयरीत्या कमी करते, कारण ते CPU पर्यंतची वळवाट पूर्णपणे टाळते.

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

आम्ही MoE communication डिस्पॅच आणि combine कर्नल्सच्या एका जोडीद्वारे अंमलात आणतो. डिस्पॅच कर्नल्स प्रत्येक रँकवरील टोकन्स आणि routing माहिती वाचण्याचे, आणि त्यांना योग्य experts कडे पाठवण्याचे काम करतात. combine कर्नल्स experts ने तयार केलेली activations गोळा करतात आणि router ने गणना केलेल्या weights च्या आधारे निवडलेल्या experts कडून येणाऱ्या weighted average चीही गणना करत त्यांना त्यांच्या source रँक्सकडे परत पाठवतात. डेटा transfer ला computation सोबत overlap करता यावे म्हणून, कर्नल्सना पुढे send आणि receive component मध्ये विभागले जाते. send कर्नल्स सर्व non-blocking आणि non-synchronizing असतात: ते फक्त सर्व writes remotes कडे dispatch करतात. दुसऱ्या बाजूला, combine कर्नल्स फक्त memory मधून वाचतात, आवश्यक सर्व data transfer होईपर्यंत थांबतात. NICs कडे काम dispatch केल्यानंतर, डेटा wire वरून asynchronously transfer होत असताना, GPUs स्थानिक पातळीवर इतर उपयुक्त काम करू शकतात, जसे shared experts लागू करणे किंवा attention compute करणे.

प्रत्येक कर्नल जोडीचे स्वतःचे सममितीय मेमरी बफर आहेत, प्रत्येक रँक इतर सर्व रँकवरून डेटा प्राप्त करण्यासाठी खाजगी स्टोरेज वाटप करतो. DP गटातून पाठवता येणाऱ्या टोकन्सची कमाल संख्या रँकवर निश्चित केली जाते, जी प्रत्येक रँक प्रत्येक DP रँकवरील प्रत्येक स्थानिक तज्ञासाठी प्राप्त करू शकणारी उच्च मर्यादा देखील ठरवते. यामुळे प्रेषक रँकना त्यांच्यामध्ये कोणत्याही सिंक्रोनाइझेशनची आवश्यकता नस, लिहिण्यासाठी गंतव्य रँकवर एक अद्वितीय पत्ता तयार करण्याची परवानगी मिळते. डेटा प्राप्त झाल्यानंतर, प्राप्त कर्नलना आवश्यक असलेल्या इष्टतम स्वरूपात मेमरीमध्ये त्यांची मांडणी करण्यासाठी कर्नल टोकन इकडे तिकडे शफल करतात. बफरचा आकार मोठा असला तरी, मॉडेलच्या सर्व अनुक्रमिक थरांमध्ये त्यांचा पुन्हा वापर केला जातो.

विस्तृत करा

वरील आकृतीत चित्रित केल्याप्रमाणे, जागतिक सिंक्रोनाइझेशनचा एकमेव प्रकार, कम्बाइन-रिसीव्ह कर्नलमध्ये अंमलात आणला जातो. एकदा पाठवण्यापासून डेटा डिस्पॅच झाला की, रँक परत येतो आणि जेव्हा तो इतर रँककडून आवश्यक असलेले सर्व टोकन प्राप्त करतो आणि प्रक्रिया करताच कम्बाइन-सेंड कर्नलमध्ये प्रवेश करतो तेव्हा अंमलबजावणी सुरू ठेवतो. कम्बाइन-रिसीव्हमधील बॅरियर हे सुनिश्चित करतो की कोणतीही रँक पुढे जाऊ शकत नाही आणि डिस्पॅच-सेंड सुरू करू शकत नाही जेव्हा इतर कोणतीही रँक अद्याप डेटा प्राप्त करण्याची वाट पाहत आहे, ज्यामुळे संभाव्यपणे विध्वंसक आच्छादन होऊ शकते. साधेपणासाठी कम्बाइन कर्नलमध्ये सिंक्रोनाइझेशन केले जाते, एक ऑप्टिमाइजेशन म्हणून डिस्पॅच रिसीव्ह आणि कम्बाइन सेंड कर्नल दरम्यान समान बॅरियर लागू केला जाऊ शकतो, ज्यामुळे GPU इतर कर्नल चालवत असताना NIC ला ट्रान्सफर स्थिरावण्यासाठी अधिक वेळ देऊन लेटन्सी आणखी कमी होते.

दोनों कर्नल डिव्हाइसच्या सर्व उपलब्ध SM वर विभागलेले आहेत: डिस्पॅच सेंड आणि कम्बाइन रिसीव्ह कर्नलना प्रति-रँक कमाल टोकन मोजणी (max_m) वर समांतर करणे आवश्यक असताना, डिस्पॅच रिसीव्ह आणि कम्बाइन सेंड कर्नलना प्रत्येक रँक आणि स्थानिक तज्ञाकडून इतकी टोकन हाताळणे आवश्यक आहे, कमाल max_m * (num_experts // EP) * (EP // TP) साठी.

डिस्पॅच

डिस्पॅच कर्नल टोकन रूट केले गेलेले तज्ञ होस्ट करणार्‍या सहकाऱ्यांना टोकन आणि ब्लॉक स्केलिंग फॅक्टर पाठवण्यासाठी जबाबदार आहे. प्रेषक बाजू दोन कार्यांना समांतर करण्यासाठी वार्प स्पेशलायझेशनवर अवलंबून असते: प्रत्येक रँकवरून किती टोकन पाठवले जातात हे ठरवण्यासाठी रूटिंग माहिती एकत्र करणे आणि रिमोटवर पाठवण्यासाठी मेसेजमध्ये टोकन पॅक करणे. प्राप्त करणारी बाजू प्रथम सर्व टोकन मोजणी प्राप्त होण्याची प्रतीक्षा करते, त्यानंतर ते डाऊनस्ट्रीम ग्रुप केलेले GEMM कर्नलना पास केले जाणार्‍या टेन्सरमध्ये सामायिक मेमरीमधून पेलोड वाचते आणि अनपॅक करते. हे प्रत्येक टोकन कुठे परत पाठवले जावे हे दर्शवण्यासाठी कम्बाइन कर्नल सामायिक केलेल्या बफरमध्ये माहिती देखील संग्रहित करते. प्राप्त कर्नल नॉन-डिटर्मिनिस्टिक क्रमाने सलग बफरमध्ये टोकन इकडे तिकडे शफल करत असल्यामुळे ही माहिती आवश्यक असते.

विस्तृत करा

प्रेषक भागात, प्रत्येक ब्लॉकमधील 32 थ्रेड्सचा एकच वार्प रूटिंग माहिती वाचण्यासाठी आणि प्रत्येक गंतव्य तज्ञाला या रँककडून प्राप्त होण्याची अपेक्षा असलेल्या टोकनची संख्या मोजण्यासाठी जबाबदार आहे, nvshmem_signal_op वापरून संख्या अधिक एक पाठवत आहे. संख्या एकाने वाढवली जाते, कारण रिमोट टोकावर शून्यापासून शून्याशिवायचे संक्रमण स्त्रोत रँकवरून मोजणी प्राप्त झाल्याचे सूचित करते. समांतरपणे, उर्वरित वार्प्स टोकन सर्व ब्लॉकमध्ये सममितीय मेमरीमध्ये समांतरपणे कॉपी करण्यासाठी सहकार्य करतात, सक्रियकरण, स्केलिंग फॅक्टर आणि स्थानिक रँकवरील त्यांची अनुक्रमणिका एका सलग भागामध्ये पॅक करतात. टोकन कुठे लिहिले जाईल हे ठरवण्यासाठी कम्बाइन प्रेषकाला अनुक्रमणिका आवश्यक असते. पुढे, वार्प गटांमधील बॅरियरद्वारे सर्व डेटा कॉपी केला गेला असल्याची खात्री केल्यानंतर, वार्प पुन्हा स्वतंत्रपणे ऑपरेट करतात, प्रत्येक समान बफर समांतरपणे वेगवेगळ्या तज्ञाला पाठवतात. nvshmemx_putmem_signal_nbi_warp वापरून टोकन पाठवले जातात, जे रिमोट डिव्हाइसवर स्थानिक रँकवरून पाठवलेल्या टोकनची संख्या देखील अनुक्रमे वाढवतात. DP गटात, पाठवले जाणारे टोकन प्रत्येक रँककडे असल्यामुळे, प्रत्येक डिव्हाइस टोकनचा उपसंच बाहेर पाठवून डिस्पॅच समान रीतीने संतुलित केले जाते.

प्राप्त करण्याच्या टोकाला, सर्व कर्नल प्रथम पाठवलेल्या टोकन मोजणी पाठवण्याची प्रतीक्षा करून प्रेषक रँकसह सिंक्रोनाइझ होतात. त्यानंतर, ते सर्व एकूण मोजणीवर स्थिरावण्यासाठी अनुक्रमे वाढवलेल्या पाठवलेल्या टोकन मोजणीची प्रतीक्षा करतात, जे putmem कॉलच्या सिमेंटिक्समुळे स्त्रोत रँकवरील सर्व पेलोड देखील पाठवले गेले आहेत हे दर्शवतात. कर्नल सर्व ब्लॉक आणि थ्रेड्समध्ये ऑपरेशन समांतर करून, nvshmem_uint64_wait_until वापरून मोजणीवर मतदान करतात. त्यानंतर, क्रॉस-ब्लॉक बॅरियर हे सुनिश्चित करतो की सर्व डेटा योग्यरित्या प्राप्त झाल्याशिवाय कोणताही ब्लॉक बफर वाचत नाही. ब्लॉकमध्ये पसरलेले आणि अनुक्रमे वाढवलेल्या काउंटरद्वारे सिंक्रोनाइझ केलेले, टोकन सममितीय मेमरीमधून नियमित सलग बफरमध्ये कॉपी केले जातात जे तज्ञांची अंमलबजावणी करणारे डाऊनस्ट्रीम कर्नलचे इनपुट बनतात. स्त्रोत रँक, तज्ञ अनुक्रमणिका आणि टोकन अनुक्रमणिका स्वतंत्रपणे साठवली जाते, कम्बाइन कर्नलला सक्रियकरण पाठवावे लागेल नेमके स्थान दर्शवते. जरी DP गटातील टोकन वेगवेगळ्या डिव्हाइसेसवरून पाठवले गेले असले तरी, संबंधित तज्ञाला पास करण्यासाठी ते सर्व एकत्र गटबद्ध केले जातात.

कम्बाइन

combine कर्नल्स पुन्हा एकदा send आणि receiver अशा अर्ध्यांमध्ये विभागले जातात: senders 16-bit un-quantized activations remote devices कडे copy करतात, तर receive कर्नल्स डेटा पाठवला जाण्याची वाट पाहतात आणि सर्व expert contributions चा weighted average स्थानिक पातळीवर compute करतात. याव्यतिरिक्त, ते dispatch-combine sequence समक्रमित करण्यासाठी barrier म्हणूनही काम करतात: प्रत्येक rank send कर्नल्समध्ये प्रवेश करताना प्रत्येक peer वर एक flag सेट करते, आणि receive कर्नल्सना तो flag सेट झालेला दिसल्याशिवाय परत येण्याची परवानगी नसते. समक्रमणाचा latency अत्यल्प असतो, कारण तो प्रत्यक्ष communication आणि computation सोबत overlap होतो.

विस्तृत करा

प्रेषक बाजूला, कर्नल सर्व स्थानिक तज्ञांना नियुक्त केलेल्या टोकनची सूची समांतरपणे पार करतात, त्यांना गंतव्य रँकवरील बफरवर लिहितात. लक्ष्य रँक, तज्ञ अनुक्रमणिका आणि टोकन अनुक्रमणिका टोकन प्राप्त झाल्यानंतर स्कॅटर कर्नल द्वारे पॉप्युलेट केलेल्या प्रति-टोकन बफरमधून वाचली जातात. आकृतीत दर्शविल्याप्रमाणे, प्रत्येक प्रेषकाला लिहिण्यासाठी प्रति तज्ञ स्वतःचा खाजगी मेमरी प्रदेश असतो, ज्यामुळे सिंक्रोनाइझ करण्याची आवश्यकता टाळली जाते. डिस्पॅच प्रमाणेच, कम्बाइन डेटा प्राप्त झाल्याचे दर्शवण्यासाठी गंतव्य रँकवरील प्रति टोकन काउंटर अनुक्रमे वाढवते: जेव्हा काउंटर टोकन ज्या तज्ञांकडे पाठवले होते त्यांच्या संक्रियेशी जुळतो, तेव्हा टोकन सामग्रीमध्ये प्रवेश केला जाऊ शकतो.

रिसीव्ह कर्नलमध्ये, सिग्नलिंग ऑपरेशनद्वारे सेट केलेल्या फ्लॅगवर मतदान करून त्यांची सामग्री येण्याची प्रतीक्षा करून, टोकनची सूची एकाधिक ब्लॉकमध्ये समांतरपणे पार केली जाते. आगमनावर, खाजगी बफरमधून पेलोड वाचले जातात, रूटिंग टेबल दर्शवते की कोणते बफर वाचायचे आहे आणि प्रत्येक तज्ञाला कोणते वजन नियुक्त करायचे आहे. परिणाम नंतर बाह्यरीत्या वाटप केलेल्या टेन्सरवर लिहिले जातात, सर्व डिव्हाइसेस बॅरियर पास केल्यानंतर कर्नल अंमलबजावणी समाप्त करते. DP गटात, सर्व रँक स्वतःची प्रतिकृती मोजण्यासाठी प्रत्येक तज्ञ सक्रियकरणाची प्रत प्राप्त करतात.

बेंचमार्क

InfiniBand आणि CX-7 NIC द्वारे जोडलेल्या 16 नोड्समध्ये पसरलेल्या 128 GPU च्या क्लस्टरवर आम्ही आमच्या कर्नलचे मूल्यांकन करतो. आम्ही CPU प्रॉक्सीसह GPUDirect Async (IBGDA) आणि Reliable Connection (RC) दोन्ही वापरून पाहतो. आम्ही त्यांची अत्यंत ऑप्टिमाइझ केलेली DeepSeek अंमलबजावणीशी तुलना करतो, तसेच PyTorch (NCCL द्वारे) किंवा NVSHMEM द्वारे प्रदान केलेल्या दाट प्रिमिटिव्हशी तुलना करतो.

विस्तृत करा

जरी दाट NVSHMEM ऑल-टू-ऑल ऑपरेशनवर, IBRC आणि IBGDA चे कार्यप्रदर्शन समान आहे (6378 µs विरुद्ध 6180 µs), विरल कर्नलसह IBGDA लक्षणीयरीत्या वेगवान आहे. डिस्पॅच आणि कम्बाइन जोडून, IBGDA 902 µs वापरते तर IBRC 3223 µs घेते - लेटन्सीमध्ये 3.6x सुधारणा. ऑल-टू-ऑल बँडविड्ह बाउंड असताना, विरल कर्नल परिमाणांचे ऑर्डर्स कमी डेटा ब्रॉडकास्ट करतात, लेटन्सी द्वारे बाउंड असतात. GPU आणि NIC समन्वयित करण्यासाठी CPU प्रॉक्सीची आवश्यकता ሳይ, थेट GPU वरून नेटवर्क ट्रान्सफर ट्रिगर करून, एंड-टू-एंड लेटन्सी लक्षणीयरीत्या कमी होते.

आमचे पोर्टेबल कर्नल अत्यंत ऑप्टिमाइझ केलेल्या DeepSeek कर्नलपेक्षा सुमारे 2 पट धीमे असताना, ते दाट कर्नलच्या तुलनेत लेटन्सी 10x सुधारतात (902 µs विरुद्ध 9944 µs).

याव्यतिरिक्त, प्रेषक-प्राप्तकर्ता घटकांमध्ये विभागणी लायब्ररी प्रिमिटिव्हच्या विपरीत, काही लेटन्सी लपविण्यास देखील अनुमती देते.

विस्तृत करा

एकल-नोड (EP8) वर, NVSHMEM वाहतुकीसाठी NVLINK वापरते, नोड-मधील नेटवर्किंगपेक्षा कमी लेटन्सी आणि उच्च थ्रूपुट वितरीत करते. आमचे पोर्टेबल कर्नल एकल-नोडवर DeepEP पेक्षा सुमारे 2.5x वेगवान आहेत (186µs विरुद्ध 481 µs).

पुढील सुधारणा

येथे वर्णलेले कर्नल ML फ्रेमवर्कच्या इन-बिल्ट प्रिमिटिव्हपेक्षा श्रेष्ठ आहेत आणि NVLink किंवा InfiniBand सारख्या विशिष्ट डिव्हाइस-मधील वाहतुकीसाठी अति-विशेष न करता योग्य कार्यप्रदर्शन देतात. आधीच नमूद केलेल्या संधींच्या बाजूला, अधिक विशेष आवृत्त्यांसह संवाद प्रिमिटिव्ह बदलून पुढील कार्यप्रदर्शन लाभ प्राप्त करता येतात. उदाहरणार्थ, NVLink वर, सममितीय मेमरीचा वापर डिव्हाइसवर सामायिक केलेल्या बफरसह बदलला जाऊ शकतो, काही प्रत बनवणे दूर करते आणि टोकन बॅचऐवजी वैयक्तिक टोकनवर सूक्ष्म सिंक्रोनाइझेशन योजनांना अनुमती देते. InfiniBand वर, अंमलबजावणी अंतर्निहित रांग जोड्या थेट प्रवेश करू शकते, दोन्ही अपेक्षित आणि पाठवलेल्या टोकन मोजणी हस्तांतरित करण्याची आवश्यकता दूर करते आणि टोकन येताच रांगेतून आउटपुट बफरमध्ये फ्लश करते. तथापि, अशी अंमलबजावणी पोर्टेबिलिटीच्या किंमतीवर येते, तर आमची अंमलबजावणी विशिष्ट हार्डवेअरसाठी ऑप्टिमाइझ करण्यासाठी विकास प्रयत्न प्रतिबद्ध करण्यापूर्वी विशिष्ट कॉन्फिगरेशनसह क्लस्टरचे मूल्यांकन करण्यास अनुमती देते.

निष्कर्ष

आम्ही MoE संवादासाठी उच्च-कार्यप्रदर्शन, पोर्टेबल लायब्ररी सादर केली आहे जी विविध हार्डवेअर कॉन्फिगरेशनमध्ये सुसंगतता राखून मानक ऑल-टू-ऑल संवादाच्या तुलनेत 10x वेगवान कार्यप्रदर्शन साध्य करते. NVLink सह एकल-नोड तैनातीवर, आमचे समाधान मागील अंमलबजावणीपेक्षा 2.5x कमी लेटन्सी प्रदर्शित करते.

आमचा दृष्टिकोन GPU-प्रारंभित संवाद समर्थन, संगणक-संवाद आच्छादन सक्षम करणारे विभाजित कर्नल आर्किटेक्चर, आणि किमान NVSHMEM प्रिमिटिव्ह वापरून कार्यक्षम टोकन डिस्पॅच यासह मुख्य नवकल्पनांद्वारे कार्यक्षमतेला पोर्टेबिलिटीसह संतुलित करतो. मल्टि-नोड सेटअपवर अत्यंत विशिष्ट अंमलबजावणीपेक्षा अंदाजे 2 पट धीमे असताना, आमची लायब्ररी विविध नेटवर्क वातावरणात (NVLink, CX-7, आणि EFA) ανώτερη लवचिकता प्रदान करते.

MoE मॉडेल्स स्केल करणे सुरू ठेवल्यामुळे, आमच्यासारख्या कार्यक्षम संवाद धोरणे व्यावहारिक तैनातीसाठी वाढत्या प्रमाणात महत्त्वाची बनतील. आमची पूर्णपणे ओपन-सोर्स अंमलबजावणी https://github.com/ppl-ai/pplx-kernels येथे उपलब्ध आहे.

संदर्भ

DeepEP

DeepSeek-V3 तांत्रिक अहवाल

NVIDIA NVSHMEM

उद्धरण

@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://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}