AWS EFA पर ट्रिलियन-पैरामीटर मॉडल को सक्षम करना

Perplexity में, हम अपने उत्पाद, अपनी API और अपनी शोध टीमों (research teams) के लिए बेहतरीन मॉडलों का उपयोग करते हैं। बड़े ओपन-सोर्स मिक्सचर-ऑफ-एक्सपर्ट्स (Mixture-of-Experts) मॉडल, जैसे कि Kimi-K2, विशेष चुनौतियां पेश करते हैं, क्योंकि 8x NVIDIA H200 GPU वाले सबसे बड़े इन्फरेन्स नोड्स उन्हें कुशलतापूर्वक समायोजित नहीं कर सकते, जिसके लिए मल्टी-नोड परिनियोजन (multi-node deployments) की आवश्यकता होती है। हम एक्सपर्ट पैरेललिज्म के लिए कर्नलों का एक सेट प्रस्तुत करते हैं जो ConnectX-7 पर अत्याधुनिक लेटेंसी प्राप्त करते हैं, जो DeepEP के प्रदर्शन से कहीं बेहतर है। यही कर्नेल AWS इलास्टिक फैब्रिक एडाप्टर (EFA) पर व्यवहार्य लेटेंसी प्राप्त करने वाले पहले कर्नल भी हैं, जो ट्रिलियन-पैरामीटर मॉडल परिनियोजन को सक्षम बनाते हैं।

GitHub पर हमारे कर्नल आज़माएं और arXiv पर पूर्ण शोध पत्र पढ़ें।

परिचय

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

पैरेललिज्म के अन्य रूपों, जैसे कि टेंसर पैरेललिज्म (TP) और डेटा पैरेललिज्म (DP), जिन्हें torch.distributed और NCCL जैसी सामूहिक संचार पुस्तकालयों का उपयोग करके आसानी से लागू किया जा सकता है, के विपरीत, MoE रूटिंग में विरल पीयर-टू-पीयर संचार शामिल है। एक डिस्पैच कर्नेल को आने वाले टोकन के सेट को विभाजित करना होगा और उन्हें विशेषज्ञों को होस्ट करने वाले रैंकों तक भेजना होगा, जबकि इसकी दोहरी समस्या, कम्बाइन कर्नेल को उनके संसाधित समकक्षों को प्राप्त करना होगा और उनके भारित औसत की गणना करनी होगी। ये ट्रांसफर कलेक्टिव्स के साथ मैप नहीं होते हैं और कम लेटेंसी प्राप्त करने के लिए अनुकूलित कर्नलों की आवश्यकता होती है।

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

इस लेख में, हम अपने नए इंटर-नोड कर्नलों के कार्यान्वयन में गहराई से उतरेंगे जो हमारे NVSHMEM-आधारित पिछले काम के प्रदर्शन में काफी सुधार करते हैं। ConnectX-7 पर अत्याधुनिक परिणामों को प्रदर्शित करने के अलावा, हम यह भी प्रदर्शित करते हैं कि हमारे इन्फरेन्स रनटाइम के भीतर, कर्नलों का उपयोग सबसे बड़े ओपन-सोर्स मॉडलों को प्रभावी ढंग से सर्व करने के लिए किया जा सकता है।

InfiniBand पर इन्फरेन्स

टेंसर पैरेललिज्म और एक्सपर्ट पैरेललिज्म एक सिंगल नोड के भीतर प्रभावी रूप से स्केल कर सकते हैं, क्योंकि NVLink माइक्रोसेकंड की लेटेंसी पर उच्च थ्रूपुट (900GB/s) प्रदान करता है। जबकि NVIDIA GB200 और GB300 सिस्टम ट्रिलियन-पैरामीटर मॉडलों के वेट्स और KV कैश को रखने के लिए पर्याप्त मेमोरी से लैस हैं, 8 H200 GPU तक वाले AWS p5en इंस्टेंस पर परिनियोजन 1120GB HBM तक सीमित हैं, जिसे मॉडल वेट्स और KV कैश दोनों के बीच विभाजित किया जाना चाहिए। परिणामस्वरूप, कुछ मॉडलों को InfiniBand पर इंटर-नोड परिनियोजन की आवश्यकता होती है।

हालाँकि InfiniBand अभी भी बहुत तेज़ है, यह केवल 400Gbps (50GB/s) का थ्रूपुट प्रदान करता है, जो MoE रूटिंग के लिए लेटेंसी में दर्जनों से सैकड़ों माइक्रोसेकंड जोड़ता है। सौभाग्य से, साझा विशेषज्ञों (shared experts), माइक्रो-बैचिंग और कंप्यूटेशन-कम्युनिकेशन ओवरलैपिंग के उपयोग के माध्यम से अतिरिक्त लेटेंसी को छिपाया जा सकता है ताकि प्रतिस्पर्धी लागत पर व्यवहार्य थ्रूपुट प्रदान किया जा सके। DeepSeek ने H800 GPU और NVIDIA ConnectX-7 InfiniBand एडाप्टर के क्लस्टर पर तैनात DeepSeek-R1 के माध्यम से इन तकनीकों का बीड़ा उठाया।

AWS क्लस्टर पर, उनके कस्टम इलास्टिक फैब्रिक एडाप्टर (EFA) के उपयोग के कारण स्थिति अलग है। भले ही ये प्रशिक्षण वर्कलोड में आमतौर पर उपयोग किए जाने वाले कलेक्टिव्स पर 400Gbps का पीक थ्रूपुट प्रदान करते हैं, लेकिन MoE डिस्पैच और कम्बाइन के दौरान एक्सचेंज किए गए संदेश आकारों पर ये ConnectX-7 से थोड़े कम पड़ जाते हैं। इसके अतिरिक्त, EFA GPUDirect Async का समर्थन नहीं करता है, जिसके लिए ट्रांसफर शुरू करने के लिए GPU को NIC से जोड़ने वाले एक CPU थ्रेड की आवश्यकता होती है। अतिरिक्त जटिलता के अलावा, अतिरिक्त PCIe लेनदेन सभी ट्रांसफर में माइक्रोसेकंड के बराबर ओवरहेड जोड़ते हैं। NVSHMEM जैसे फ्रेमवर्क कर्नेल लेखकों को डिवाइस-स्वतंत्र कार्यान्वयन प्रदान करने की अनुमति देते हैं, हालांकि सामान्य API लागत और लेटेंसी दोनों के संदर्भ में परिनियोजन को अप्रभावी बना देते हैं।

NVSHMEM-आधारित pplx-kernels

हमारे पिछले कर्नलों ने NVSHMEM के माध्यम से अंतर्निहित ट्रांसपोर्ट कार्यान्वयन को हटा दिया था। वे साथियों को सिंक्रनाइज़ करने के लिए एटॉमिक काउंटर्स का उपयोग करके टोकन को एक-एक करके ट्रांसफर करने के लिए nvshmemx_putmem_signal_nbi_warp पर निर्भर थे। जबकि इसने IBGDA के साथ ConnectX-7 NIC पर उचित प्रदर्शन की पेशकश की, प्रॉक्सी-आधारित ConnectX-7 IBRC कार्यान्वयन काफी धीमा था, जबकि EFA कार्यान्वयन एक मिलीसेकंड से कम लेटेंसी प्रदान करने में विफल रहा।

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

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

EFA पर MoE ऑल-टू-ऑल

बड़ा करें

नए कर्नेल एक हाइब्रिड CPU-GPU आर्किटेक्चर का पालन करते हैं, जहाँ GPU कर्नेल डिवाइस पर मॉडल के साथ इंटरैक्ट करते हैं और होस्ट CPU पर चलने वाला एक प्रॉक्सी थ्रेड NIC के साथ इंटरैक्शन को संभालता है। नए कर्नेल उसी TransferEngine का पुन: उपयोग करते हैं जिसे हमने KV कैश ट्रांसफर के लिए बनाया था।

MoE रूटिंग एक डिस्पैच कर्नेल द्वारा की जाती है, जो मॉडल से रूट करने के लिए टोकन और प्रति-टोकन विशेषज्ञों को प्राप्त करता है, उन्हें उन रैंकों पर नेटवर्क के माध्यम से भेजता है जो विशेषज्ञों को होस्ट करते हैं। समूहीकृत GEMM कर्नलों के टोकन को विशेषज्ञ वेट्स के साथ गुणा करने के बाद, कम्बाइन कर्नेल टोकन को उनके स्रोत रैंकों पर वापस ट्रांसफर करता है और उन्हें मॉडल पर लौटने से पहले उनके भारित औसत की गणना करता है।

डिस्पैच और कम्बाइन दोनों कर्नेल को ओवरलैपिंग और माइक्रो-बैचिंग की सुविधा के लिए एक प्रेषक (sender) और एक रिसीवर (receiver) आधे हिस्से में विभाजित किया गया है। प्रेषक टोकन को सेंड बफ़र्स में शफल करता है ताकि उन्हें एक ही RDMA राइट ऑपरेशन के माध्यम से साथियों को लिखा जा सके। प्रेषक को निष्पादित करने के बाद, GPU अन्य कार्य (जैसे साझा विशेषज्ञ गुणन) कर सकता है जबकि नेटवर्क पर ट्रांसफर बैकग्राउंड में लंबित होते हैं। रिसीवर कर्नेल तब तक ब्लॉक रहते हैं जब तक कि सभी ट्रांसफर प्राप्त नहीं हो जाते। बैकग्राउंड में, एक होस्ट-साइड वर्कर थ्रेड ट्रांसफर को संभालता है: सेंड बफ़र्स पॉप्युलेट होने पर इसे GPU द्वारा सिग्नल दिया जाता है, जिसके बाद यह RDMA राइट्स को डिस्पैच करता है। एक बार जब NIC पूर्णता का संकेत देता है, तो यह रिसीवर को अनब्लॉक करता है, जो टोकन को मॉडल पर लौटा देता है।

ये कर्नेल होस्ट के साथ कसकर जुड़े हुए हैं, जो यूनिफाइड मेमोरी और GDRCopy दोनों के माध्यम से एक-दूसरे के साथ संचार करते हैं। CUDA यूनिफाइड मेमोरी डिवाइस मेमोरी या होस्ट मेमोरी को दोनों के एड्रेस स्पेस में मैप करने की अनुमति देती है, जो PCIe के माध्यम से संचार को पारदर्शी रूप से संभालती है।

जबकि यूनिफाइड मेमोरी बल्क ट्रांसफर के लिए उपयुक्त है, कम-लेटेंसी पोलिंग GDRCopy के माध्यम से की जाती है। GDRCopy माइक्रोसेकंड लेटेंसी पर GPU मेमोरी से पढ़ने और लिखने के लिए GPUDirect RDMA का उपयोग करता है।

RDMA के अतिरिक्त, एक नोड के भीतर कर्नल टोकन का आदान-प्रदान करने के लिए NVLink का भी लाभ उठाते हैं। यह ट्रैफिक के एक महत्वपूर्ण हिस्से को ऑफलोड कर सकता है, EP32 के लिए एक चौथाई तक और EP64 के लिए आठवें हिस्से तक, जिससे कुल मिलाकर कम लेटेंसी प्राप्त होती है।

सेंड और रिसीव बफ़र्स

डिस्पैच कर्नेल के लिए सेंड बफ़र्स का डिज़ाइन होस्ट प्रॉक्सी से जुड़ी व्यवहार्य लेटेंसी प्राप्त करने के लिए महत्वपूर्ण है। यदि N रैंक हैं, जिनमें से प्रत्येक T टोकन डिस्पैच कर रहा है, जिन्हें E में से R विशेषज्ञों के लिए रूट किया गया है, तो एक रैंक सबसे खराब स्थिति में N * T * max(R, E / N) टोकन प्राप्त कर सकता है, यदि सभी को उसी पर रूट किया गया है। इसके अलावा, यह न्यूनतम सीमा केवल तभी प्राप्त की जा सकती है जब प्रेषक अपने राइट्स को लगातार व्यवस्थित करें, क्योंकि निजी रिसीव बफ़र्स बहुत अधिक मेमोरी का उपयोग करेंगे। प्रेषकों के लिए बिना किसी संघर्ष के लगातार अपने राइट्स को व्यवस्थित करने में सक्षम होने के लिए, रूटिंग जानकारी का पूर्ण ज्ञान प्रत्येक रैंक में वितरित किया जाना चाहिए। परिणामस्वरूप, हमारे डिस्पैच कर्नेल पहले प्रति-विशेषज्ञ टोकन गणना का आदान-प्रदान करते हैं, जिसके बाद वे प्रत्येक निर्धारित कर सकते हैं कि उनके राइट्स कहाँ जाते हैं और वे प्रत्येक रैंक से किस ऑफसेट पर टोकन प्राप्त करते हैं।

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

बड़ा करें

कम्बाइन कर्नेल काफी सरल हैं, क्योंकि रूटिंग जानकारी उनके लिए डिस्पैच चरण से पहले से ही उपलब्ध है। वे टोकन को प्रति पीयर एक सिंगल राइट में कॉपी कर सकते हैं, जिसके बाद वे टोकन के भारित औसत की गणना करके आउटपुट टेंसर में करने के लिए रिसीव बफ़र से पढ़ते हैं।

NVLink के लिए, एक नोड के भीतर सेंड और रिसीव बफ़र्स को साथियों में मैप किया जाता है ताकि वे पड़ोसी रैंकों से सीधे टोकन को पुश या पुल कर सकें।

डिस्पैच और कम्बाइन

नीचे दिए गए आंकड़े GPU, CPU और NIC के बीच के समय और इंटरैक्शन का विवरण देते हैं, जो मॉडल के कर्नलों के साथ इंटरलीव्ड डिस्पैच और कम्बाइन के एक पुनरावृत्ति (iteration) को दिखाते हैं।

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

बड़ा करें

बड़ा करें

एक बार जब कोई रैंक सभी साथियों से रूटिंग जानकारी प्राप्त कर लेता है, तो यह इसे केंद्रीकृत करता है और डिस्पैच और कम्बाइन के लिए टोकन की पहचान करने के लिए आवश्यक विभिन्न ऑफसेट की गणना करता है। यह कार्य CPU पर तब किया जाता है जब ट्रांसफर पहले से ही लंबित होते हैं, इसलिए यह लेटेंसी के प्रति संवेदनशील नहीं है। बाद में, एक बार सभी ऑफसेट ज्ञात हो जाने के बाद, शेष टोकन के लिए राइट्स, जो निजी बफ़र आकार से अधिक हैं, पोस्ट किए जाते हैं और डिस्पैच का सेंड चरण पूरा हो जाता है।

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

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

MoE रूटिंग के लिए TransferEngine

अपने कार्यान्वयन के प्रोटोटाइप के लिए, हमने शुरू में RDMA राइट्स को जारी करने और ट्रैक करने के लिए डिसएग्रीगेट प्रीफिल और डिकोड के लिए बनाए गए Transfer Engine का पुन: उपयोग किया। हालांकि, डिसएग्रिगेट डिकोडिंग के विपरीत जहां प्रत्येक नोड एक ही पीयर को बड़ी संख्या में राइट्स पोस्ट करता है, MoE डिस्पैच एक ऑपरेशन में सभी साथियों को छोटे ट्रांसफर पोस्ट करता है। आउट-ऑफ-द-बॉक्स, प्रदर्शन इष्टतम नहीं था और हमारा कार्यान्वयन ConnectX-7 पर DeepEP से लगभग 20us पीछे चल रहा था। परिणामस्वरूप, हमने बड़ी संख्या में साथियों के साथ इंटरैक्शन को अनुकूलित करने के लिए दो ऑपरेशनों, स्कैटर और बैरियर का समर्थन करने के लिए TransferEngine को विशेष रूप से तैयार किया।

स्कैटर और बैरियर दोनों ऑपरेशन पीयर समूहों को उजागर करते हैं जिन्हें अंतर्निहित ट्रांसपोर्ट कार्यान्वयन (वर्क रिक्वेस्ट (WR) टेम्प्लेटिंग) द्वारा उपयोग की जाने वाली डेटा संरचनाओं को पहले से पॉप्युलेट करने के लिए पहले से पंजीकृत किया जा सकता है। स्कैटर एक ही स्रोत बफ़र से विभिन्न स्लाइस को पीयर समूह के सभी सदस्यों पर अलग-अलग ऑफसेट में कॉपी करता है। बैरियर एक इमीडिएट के माध्यम से सभी साथियों को संकेत देता है। महत्वपूर्ण पथ पर इन ऑपरेशनों को बल्क में इनक्यू और हैंडल करने से कीमती माइक्रोसेकंड की बचत होती है। EFA पर, जहां दो 200Gbps NICs 400Gbps बैंडविड्थ प्रदान करने के लिए एग्रीगेट होते हैं, हम प्रत्येक ट्रांसफर के बाइट्स को शार्ड करने के बजाय दो NICs के बीच ग्रुप पीयर्स को शार्ड करते हैं।

हालांकि शुरू में हमने TransferEngine को पूरी तरह से libfabric का उपयोग करके EFA पर विकसित किया, MoE कर्नेल TransferEngine में libibverbs का उपयोग करके ConnectX-7 समर्थन जोड़ने के बाद कोड की एक भी पंक्ति बदले बिना ConnectX-7 पर काम करते हैं। स्केलेबल रिलायबल डेटाग्राम (SRD) आधारित EFA की तुलना में, ConnectX-7 को कनेक्शन सेटअप और पीयर प्रबंधन के लिए अतिरिक्त कार्य की आवश्यकता होती है। हम रिलायबल कनेक्शन (RC) को बूटस्ट्रैप करने के लिए अनरिलायबल डेटाग्राम (UD) क्यू पेयर का उपयोग करते हैं। हम प्रत्येक रिमोट पीयर के लिए दो क्यू पेयर बनाते हैं, एक टू-साइडेड SEND/RECV के लिए और एक वन-साइडेड WRITE_IMM के लिए, क्योंकि RECV और WRITE_IMM दोनों क्रम में RECV ऑपरेशन का उपभोग करते हैं जबकि दोनों ऑपरेशनों को अलग-अलग बफ़र आकारों की आवश्यकता होती है। इसके अलावा, हम ConnectX-7 और GPU मेमोरी के बीच PCIe लेनदेन ऑर्डरिंग में ढील देते हैं। EFA के समान, हम WR टेम्प्लेटिंग ऑप्टिमाइज़ेशन को लागू करते हैं। इसके अतिरिक्त, जब एक ही पीयर को कई छोटे संदेश भेजते हैं, तो हम संदेश दर में सुधार करने के लिए WR चेनिंग का उपयोग करते हैं।

चूंकि GPU-साइड कार्य दोनों एडाप्टर में समान है, हमने DeepEP के प्रदर्शन से आगे निकलने के लिए अंतर्निहित ट्रांसपोर्ट कार्यान्वयन को अपनी सीमाओं तक धकेल दिया ताकि यह उजागर किया जा सके कि EFA और ConnectX-7 MoE रूटिंग के बीच का अंतर अंतर्निहित हार्डवेयर की क्षमताओं द्वारा निर्धारित होता है।

मूल्यांकन

डिस्पैच/कम्बाइन लेटेंसी

हमने 7168 के हिडन डाइमेंशन, ब्लॉक-स्केल्ड fp8 डिस्पैच और bf16 कम्बाइन के साथ मानक DeepSeek-V3 और Kimi-K2 कॉन्फ़िगरेशन पर ConnectX-7 और EFA एडाप्टर से लैस H200 सिस्टम पर डिस्पैच और कम्बाइन लेटेंसी का मूल्यांकन किया। हम डिस्पैच के लिए 128 और कम्बाइन के लिए 4096 के बैच आकार पर विचार करते हैं। ConnectX-7 पर, हम GPUDirect Async (IBGDA) और होस्ट प्रॉक्सी (IBRC) दोनों का उपयोग करके DeepEP और हमारे पिछले NVSHMEM-आधारित pplx-kernels के साथ तुलना करते हैं। हम क्रमशः EP16, EP32 और EP64 पर 459us, 582us और 692us की संयुक्त लेटेंसी प्राप्त करते हैं। हम UCCL-EP जैसे सामान्य प्रॉक्सी पर निर्भर अन्य EFA-आधारित कर्नलों से बेहतर प्रदर्शन करते हैं, जो 519us, 966us और 1159us प्राप्त करते हैं।

बड़ा करें

डिकोड पर, हमारे कर्नेल के DeepEP की तुलना में कुछ माइक्रोसेकंड धीमे होने के बावजूद, हम कुल मिलाकर कम लेटेंसी प्राप्त करते हैं क्योंकि हमारे कम्बाइन कर्नेल काफी तेज़ हैं। हालांकि हम एक होस्ट प्रॉक्सी का उपयोग करते हैं, हमारे बल्क ट्रांसफर के कारण हम सभी IBGDA-आधारित प्रॉक्सीलेस कार्यान्वयनों के प्रदर्शन से आगे निकल जाते हैं। हम यह भी दिखाते हैं कि प्रॉक्सी कार्यान्वयन मायने रखता है, क्योंकि हमारा अत्यधिक विशेष प्रॉक्सी सामान्य NVSHMEM कार्यान्वयन की तुलना में दस गुना तेज़ है। इस तथ्य के बावजूद कि 256KB पैकेट ट्रांसफर करते समय EFA, ConnectX-7 की तुलना में दोगुना धीमा है, जो MoE रूटिंग द्वारा एक्सचेंज किए गए सामान्य पैकेट आकार के करीब हैं, दोनों एडाप्टर के बीच का अंतर कुल मिलाकर कम महत्वपूर्ण है।

बड़ा करें

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

Kimi-K2 / DeepSeek-V3 थ्रूपुट

अपने कर्नलों की व्यवहार्यता साबित करने के लिए, हम Elastic Fabric Adapters (EFA) और H200 GPU से लैस कई p5en इंस्टेंस पर DeepSeek-V3 (671B पैरामीटर) और Kimi-K2 (1T पैरामीटर) की सेवा करने वाले अपने इन-हाउस इन्फरेन्स इंजन के एंड-टू-एंड प्रदर्शन को बेंचमार्क करते हैं। हम मॉडल को ड्राफ्ट वेरिफिकेशन मोड में चलाते हैं और एक सिंगल MTP लेयर का उपयोग करके 80% की स्वीकृति दर मानकर लेटेंसी की रिपोर्ट करते हैं।

DeepSeek-V3 पर, हम कुशल NVLink-ओनली डिस्पैच और कम्बाइन कर्नेल चलाने वाले EP=8 इंट्रा-नोड बेसलाइन के साथ तुलना करते हैं। इंटर-नोड परिनियोजन कुल मिलाकर सिंगल-नोड परिनियोजन के प्रदर्शन के बराबर या उससे बेहतर है, विशेष रूप से मध्यम आकार के बैचों पर, यह दर्शाता है कि नोड्स में स्केलिंग करने से कुल थ्रूपुट में सुधार होता है।

Kimi-K2 पर, कर्नेल सेवा को सक्षम करते हैं। इसके आकार के कारण, मॉडल H200 नोड पर डेटा पैरेललिज्म (DP) का उपयोग करके एक ही नोड के भीतर फिट नहीं हो सकता है, इसलिए मॉडल की सेवा के लिए एक इंटर-नोड परिनियोजन आवश्यक है। चूंकि मॉडल में कम अटेंशन हेड्स हैं, हम DeepSeek-V3 की तुलना में बेहतर लेटेंसी प्राप्त करते हैं और हम व्यवहार्य लेटेंसी पर EFA पर मॉडल की सेवा कर सकते हैं।

बड़ा करें

बड़ा करें

हमने माइक्रो-बैचिंग पर भी विचार किया है, हालाँकि इसने केवल बहुत उच्च बैच आकारों पर या एक से अधिक ड्राफ्ट टोकन का उपयोग करते समय लेटेंसी में सुधार किया है। समूहीकृत GEMM समस्याओं को दो में विभाजित करने से लाभ उठाने के लिए, टोकन बैच के आकार पर्याप्त रूप से सघन होने चाहिए। अधिकांश समय अभी भी MoE डिस्पैच और कम्बाइन पर खर्च किया जाता है, जो यह दर्शाता है कि कर्नलों में कोई भी और सुधार लेटेंसी में सुधार करेगा और माइक्रो-बैचिंग की सुविधा प्रदान करेगा।

भविष्य का कार्य

हम EFA पर प्रदर्शन में सुधार के लिए AWS इंजीनियरों के साथ मिलकर काम कर रहे हैं। अपनी तरफ से, हम libfabric के अपडेट का पालन करेंगे जो डेटा प्लेन ओवरहेड को कम करेगा। इसके अतिरिक्त, हम यूज़रस्पेस ओवरहेड को और कम करने और सीधे हार्डवेयर के साथ इंटरफ़ेस करने के लिए efa-direct के साथ प्रयोग करेंगे।

संदर्भ

DeepSeek R1 तकनीकी रिपोर्ट

Kimi K2 तकनीकी रिपोर्ट

DeepEP

UCCL-EP

NVIDIA Magnum IO NVSHMEM और GPUDirect Async का उपयोग करके HPC सिस्टम का नेटवर्क प्रदर्शन सुधारना

उद्धरण

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}