AWS EFA वर ट्रिलियन-पॅरामीटर मॉडेल्स सक्षम करणे

Perplexity मध्ये, आम्ही आमच्या उत्पादन, आमच्या API आणि आमच्या संशोधन संघांसाठी सर्वोत्तम मॉडेल्स वापरतो. किमी-के२ (Kimi-K2) सारखी मोठी ओपन-सोर्स मिक्सचर-ऑफ-एक्सपर्ट्स (MoE) मॉडेल्स विशेष आव्हाने उभी करतात, कारण ८x NVIDIA H200 GPU असलेले सर्वात मोठे इन्फरन्स नोड्स त्यांना कार्यक्षमतेने सामावून घेऊ शकत नाहीत, ज्यामुळे बहु-नोड उपयोगांची आवश्यकता असते. आम्ही एक्सपर्ट पॅरॅलेलिझमसाठी कर्नलचा एक संच सादर करतो जो ConnectX-7 वर अत्याधुनिक लेटन्सी प्राप्त करतो, जी DeepEP च्या कार्यक्षमतेपेक्षा जास्त आहे. हेच कर्नल AWS इलास्टिक फॅब्रिक अडॅप्टर (EFA) वर व्यवहार्य लेटन्सी साध्य करणारे पहिले कर्नल आहेत, ज्यामुळे ट्रिलियन-पॅरामीटर मॉडेल तैनात करणे शक्य होते.

आमचे कर्नल GitHub वर वापरून पहा आणि arXiv वर संपूर्ण संशोधन पेपर वाचा.

परिचय

मिक्सचर-ऑफ-एक्सपर्ट्स (MoE) हे शेकडो अब्ज किंवा ट्रिलियन पॅरामीटर्सपर्यंत स्केल करण्यासाठी मानक आर्किटेक्चर म्हणून उदयास आले आहे, तसेच इन्फरन्सच्या वेळी वाजवी लेटन्सी देखील देते. MoE मॉडेल्स ट्रान्सफॉर्मरच्या घट्ट स्तराला काही तज्ञांसह (एक्सपर्ट्स) आणि एका लहान लीनियर राऊटिंग स्तरासह बदलतात जो ठरवतो की टोकनला कोणत्या तज्ञाने गुणिले जावे. समान तज्ञांशी संबंधित टोकन्स गटबद्ध करून, समस्या विरल आणि सहज समांतर बनते, कारण आता विविध तज्ञ एकाधिक नोड्सवर पसरलेल्या वेगवेगळ्या GPU वर होस्ट केले जाऊ शकतात.

टॉन्सर पॅरॅलेलिझम (TP) आणि डेटा पॅरॅलेलिझम (DP) सारख्या समांतरतेच्या इतर रूपांच्या विपरीत, ज्या torch.distributed आणि NCCL सारख्या सामूहिक संप्रेषण लायब्ररी वापरून सहजपणे लागू केल्या जाऊ शकतात, MoE राऊटिंगमध्ये विरल पीअर-ट्यू-पीअर संप्रेषण समाविष्ट असते. एका डिस्पॅच कर्नलने येणाऱ्या टोकन्सचा संच विभाजित केला पाहिजे आणि त्यांना तज्ञांना होस्ट करणाऱ्या रँक्सवर पाठवले पाहिजे, तर त्याची द्विर्मितीय समस्या, कम्बाइन कर्नलने, त्यांच्या प्रक्रिया केलेल्या समकक्षांना प्राप्त करणे आणि त्यांची भारित सरासरी मोजणे आवश्यक आहे. हे हस्तांतरण सामूहिकतेशी जुळत नाहीत आणि कमी लेटन्सी साध्य करण्यासाठी सानुकूलित कर्नल आवश्यक आहेत.

Perplexity मध्ये, आम्ही पोर्टेबल आंतर-नोड कर्नल आणि विशेषित इंट्रा-नोड कर्नल दोन्ही विकसित करत आहोत, आमचे काम MoE संप्रेषण वरील मागील लेखात अधिक खोलीत सादर करत आहोत. आम्ही दुसर्‍या ब्लॉग पोस्टमध्ये हे देखील दाखवले आहे की बहु-नोड तैनाती मोठ्या MoE मॉडेल्ससाठी कमी लेटन्सी आणि उच्च थ्रुपुट आणू शकते.

या लेखात, आम्ही आमच्या नवीन आंतर-नोड कर्नल्सच्या अंमलबजावणीचा सखोल अभ्यास करतो जे आमच्या NVSHMEM-आधारित मागील कामाच्या कार्यक्षमतेमध्ये मोठ्या प्रमाणात सुधारणा करतात. ConnectX-7 वर अत्याधुनिक परिणाम दर्शविण्याव्यतिरिक्त, आम्ही हे देखील प्रदर्शित करतो की आमच्या इन्फरन्स रनटाइममध्ये, सर्वात मोठी ओपन-सोर्स मॉडेल्स प्रभावीपणे सर्व्ह करण्यासाठी कर्नल वापरले जाऊ शकतात.

InfiniBand वर इन्फरन्स

NVLink उच्च लेटन्सीच्या मायक्रोसेकंदमध्ये उच्च थ्रुपुट (900GB/s) ऑफर करत असल्याने, टॉन्सर पॅरॅलेलिझम आणि एक्सपर्ट पॅरॅलेलिझम एकाच नोडमध्ये प्रभावीपणे स्केल करू शकतात. NVIDIA GB200 आणि GB300 सिस्टीम ट्रिलियन-पॅरामीटर मॉडेल्सचे वजन आणि KV कॅशे धारण करण्यासाठी पुरेशी मेमरी सुसज्ज असताना, ८ H200 GPU पर्यंतच्या AWS p5en इन्स्टन्सवरील तैनाती HBM च्या 1120GB पर्यंत मर्यादित आहेत, जे मॉडेल वजन आणि KV कॅशे दोन्हीमध्ये विभागले जाणे आवश्यक आहे. परिणामी, विशिष्ट मॉडेल्सना InfiniBand वर आंतर-नोड तैनाती आवश्यक असते.

InfiniBand अद्याप अत्यंत जलद असले तरी, ते केवळ 400Gbps (50GB/s) थ्रुपुट ऑफर करते, ज्यामुळे MoE राऊटिंगसाठी लेटन्सीमध्ये शेकडो मायक्रोसेकंदांची भर पडते. सुदैवाने, स्पर्धात्मक किमतीत व्यवहार्य थ्रुपुट वितरीत करण्यासाठी सामायिक तज्ञ, मायक्रो-बॅचिंग आणि कॉम्प्युटेशन-कम्युनिकेशन ओव्हरलॅपिंग वापरून जोडलेली लेटन्सी लपवली जाऊ शकते. 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 सह संवाद हाताळतो. नवीन कर्नल आम्ही KV कॅशे ट्रान्सफर साठी तयार केलेले समान TransferEngine पुन्हा वापरतात.

MoE राऊटिंग डिस्पॅच कर्नलद्वारे केले जाते, जे मॉडेलकडून टोकन आणि प्रति-टोकन तज्ञांना त्यांच्यावर राऊत करण्यासाठी प्राप्त करते, तज्ञांना होस्ट करणाऱ्या रँक्सवर नेटवर्कवर पाठवते. ग्रुप केलेले GEMM कर्नल तज्ञांच्या वजनासह टोकन्स गुणाकार केल्यानंतर, कम्बाइन कर्नल टोकन्स त्यांच्या स्त्रोत रँक्सवर हस्तांतरित करतात आणि त्यांना मॉडेलवर परत करण्यापूर्वी त्यांची भारित सरासरी मोजतात.

ओव्हरलॅपिंग आणि मायक्रो-बॅचिंग सुलभ करण्यासाठी डिस्पॅच आणि कम्बाइन कर्नल दोन्ही प्रेषक आणि प्राप्तकर्ता अर्ध्या भागात विभाजित आहेत. प्रेषक टोकन्सना पाठवण्याच्या बफरमध्ये शफल करतो जेणेकरून ते एकाच 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) टोकन्स प्राप्त करू शकते, जर सर्व त्यावर राऊत केले गेले. शिवाय, हे किमान बंधन केवळ तेव्हाच साध्य केले जाऊ शकते जेव्हा प्रेषक त्यांचे लेखन सलगपणे मांडतात, कारण खाजगी प्राप्तकर्ता बफर खूप जास्त मेमरी वापरतील. प्रेषकांसाठी संघर्षाशिवाय त्यांचे लेखन सलगपणे मांडण्यास सक्षम होण्यासाठी, राऊटिंग माहितीचे संपूर्ण ज्ञान प्रत्येक रँकला वितरित केले जाणे आवश्यक आहे. परिणामी, आमचे डिस्पॅच कर्नल प्रथम प्रति-तज्ञ टोकन मोजणीची देवाणघेवाण करतात, त्यानंतर ते प्रत्येक ठरवू शकतात की त्यांचे लेखन कुठे जाते आणि त्यांना कोणत्या रँकवरून टोकन्स मिळतात.

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

विस्तृत करा

कम्बाइन कर्नल sostanzially सोपे आहेत, कारण डिस्पॅच टप्प्यावरून राऊटिंग माहिती त्यांच्यासाठी आधीच उपलब्ध आहे. ते प्रति पीअर एका लेखनात टोकन्स कॉपी करू शकतात, त्यानंतर ते आउटपुट टेन्सरमध्ये टोकन्सची भारित सरासरी मोजण्यासाठी प्राप्त बफरवरून वाचतात.

NVLink साठी, एका नोडमधील पाठवा आणि प्राप्त बफर सहकाऱ्यांमध्ये मॅप केले जातात जेणेकरून त्यांना शेजारच्या रँकमधून थेट टोकन्स पुश किंवा पुल करता येतील.

डिस्पॅच आणि कम्बाइन

खालील आकृती GPU, CPU आणि NIC मधील वेळ आणि संवाद तपशीलवार दर्शवते, मॉडेलच्या कर्नलसह इंटरलेव्ह केलेले डिस्पॅच आणि कम्बाइनची एक पुनरावृत्ती दर्शवते.

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

विस्तृत करा

विस्तृत करा

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

प्राप्तकर्त्याच्या शेवटी, कर्नल प्रथम NVLink द्वारे खाजगी बफरच्या क्षमतेच्या पलीकडे असलेल्या त्याच नोडमधून टोकन्स पुल करतो. पुढे, होस्ट प्रॉक्सी सर्व ट्रान्सफर पूर्ण झाल्याचा सिग्नल देईपर्यंत तो ब्लॉक होतो. एकदा सिग्नल प्राप्त झाल्यानंतर, पुढील कर्नलला त्यांना ओव्हरराईड करण्याची परवानगी देण्यापूर्वी, बफर डिस्पॅच आणि कम्बाइन दरम्यान पुन्हा वापरले जात असल्याने, NVLink आणि RDMA बॅरियर दोन्ही प्राप्त बफर वाचले गेले आहेत याची खात्री करतात, टोकन्स प्राप्त बफरमधून ग्रुप GEMM ला पास केलेल्या आउटपुट टेन्सरमध्ये शफल केले जातात.

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

MoE राऊटिंगसाठी TransferEngine

आमच्या अंमलबजावणीचे प्रोटोटाइप करण्यासाठी, आम्ही सुरुवातीला RDMA लेखन जारी करण्यासाठी आणि ट्रॅक करण्यासाठी डिसेग्रीगेट प्रीफिल आणि डीकोडसाठी तयार केलेले ट्रान्सफर इंजिन पुन्हा वापरले. तथापि, डिसेग्रीगेटिंग डीकोडिंगच्या विपरीत जिथे प्रत्येक नोड समान पीअरला मोठ्या संख्येने लेखन पोस्ट करतो, MoE डिस्पॅच एकाच ऑपरेशनमध्ये सर्व सहकाऱ्यांना लहान ट्रान्सफर पोस्ट करतो. आउट-ऑफ-द बॉक्स, कार्यप्रदर्शन इष्टतम नव्हते आणि आमची अंमलबजावणी ConnectX-7 वर DeepEP च्या मागे सुमारे 20us मागे होती. परिणामी, आम्ही मोठ्या संख्येने सहकाऱ्यांशी संवाद ऑप्टिमाइझ करण्यासाठी, स्कॅटर आणि बॅरियर या दोन ऑपरेशन्सना समर्थन देण्यासाठी TransferEngine विशेषित केले.

स्कॅटर आणि बॅरियर दोन्ही ऑपरेशन्स पीअर ग्रुप्स उघड करतात जे अंतर्निहित ट्रान्सफर्ट अंमलबजावणीद्वारे वापरल्या जाणाऱ्या डेटा संरचनांना पूर्व-लोकसंख्या करण्यासाठी पूर्व-नोंदणीकृत केले जाऊ शकतात (वर्क रिक्वेस्ट (WR) टेम्पलेटिंग). स्कॅटर पीअर ग्रुपच्या सर्व सदस्यांवर वेगवेगळ्या ऑफसेटवर समान स्त्रोत बफरमधून वेगवेगळ्या स्लाईसेस कॉपी करतो. बॅरियर तात्काळ माध्यमातून सर्व सहकाऱ्यांना सिग्नल करते. मोठ्या प्रमाणात या ऑपरेशन्स इनक्युइंग आणि हाताळल्याने गंभीर मार्गावर मौल्यवान मायक्रोसेकंद वाचतात. EFA वर, जिथे दोन 200Gbps NIC 400Gbps बँडविड्थ प्रदान करण्यासाठी एकत्रित होतात, आम्ही प्रत्येक ट्रान्सफरचे बाइट्स शार्ड करण्याऐवजी दोन NIC मध्ये ग्रुप पीअर शार्ड करतो.

जरी सुरुवातीला आम्ही libfabric वापरून केवळ EFA वर TransferEngine विकसित केले असले, तरी TransferEngine मध्ये libibverbs वापरून ConnectX-7 समर्थन जोडल्यानंतर MoE कर्नल कोडची एक ओळ न बदलता ConnectX-7 वर कार्य करतात. स्केलेबल विश्वासार्ह डेटाग्राम (SRD) आधारित EFA च्या तुलनेत, ConnectX-7 ला कनेक्शन सेटअप आणि पीअर व्यवस्थापनासाठी अतिरिक्त कामाची आवश्यकता असते. आम्ही विश्वसनीय कनेक्शन (RC) बूस्टस्ट्रॅप करण्यासाठी अविश्वसनीय डेटाग्राम (UD) राऊंड जोडणी (queue pair) वापरतो. आम्ही प्रत्येक रिमोट पीअरसाठी दोन राऊंड जोडणी तयार करतो, एक दोन-बाजूंनी SEND/RECV साठी आणि एक एक-बाजूंनी WRITE_IMM साठी, कारण RECV आणि WRITE_IMM दोन्ही क्रमाने RECV ऑपरेशन वापरतात तर दोन ऑपरेशन्सना वेगवेगळ्या बफर आकारांची आवश्यकता असते. शिवाय, आम्ही ConnectX-7 आणि GPU मेमरीमधील PCIe व्यवहार क्रमवारी सैल करतो. EFA प्रमाणे, आम्ही WR टेम्पलेटिंग ऑप्टिमायझेशन लागू करतो. याव्यतिरिक्त, त्याच पीअरला अनेक लहान संदेश पाठवताना, आम्ही संदेश दर सुधारण्यासाठी WR चेनिंग वापरतो.

GPU-बाजूचे कार्य दोन्ही अडॅप्टरमध्ये समान असल्यामुळे, EFA आणि ConnectX-7 MoE राऊटिंगमधील अंतर अंतर्निहित हार्डवेअरच्या क्षमतांनी निर्धारित केले जाते हे हायलाइट करण्यासाठी DeepEP च्या कार्यक्षमतेपेक्षा जास्त करण्यासाठी आम्ही अंतर्निहित वाहतूक अंमलबजावणी त्याच्या मर्यादांपर्यंत ढकलली.

मूल्यांकन

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

आम्ही मानक DeepSeek-V3 आणि Kimi-K2 कॉन्फिगरेशनवर ConnectX-7 आणि EFA अडॅप्टर सुसज्ज H200 सिस्टीमवर डिस्पॅच आणि कम्बाइन लेटन्सीचे मूल्यांकन केले ज्यामध्ये ७१६८ चे लपलेले परिमाण, ब्लॉक-स्केल केलेले fp8 डिस्पॅच आणि bf16 कम्बाइन आहे. आम्ही डिस्पॅचसाठी 128 आणि कम्बाइनसाठी 4096 चे बॅच आकार विचारात घेतो. ConnectX-7 वर, आम्ही GPUDirect Async (IBGDA) आणि होस्ट प्रॉक्सी (IBRC) दोन्ही वापरून DeepEP आणि आमच्या मागील NVSHMEM-आधारित pplx-kernels शी तुलना करतो. आम्ही अनुक्रमे EP16, EP32 आणि EP64 वर 459us, 582us आणि 692us ची एकत्रित लेटन्सी साध्य करतो. आम्ही UCCL-EP सारख्या जेनेरिक प्रॉक्सीवर अवलंबून असलेल्या इतर EFA-आधारित कर्नलपेक्षा श्रेष्ठ कामगिरी करतो, जे 519us, 966us आणि 1159us साध्य करतात.

विस्तृत करा

डीकोडवर, आमचे kernels DeepEP पेक्षा काही मायक्रोसेकंदांनी मंद असले तरी, आमच्या combine kernels लक्षणीयरीत्या जलद असल्यामुळे आम्ही एकूण कमी latencies साध्य करतो. जरी आम्ही host proxy वापरत असलो, तरी आमच्या bulk transfers मुळे आम्ही सर्व IBGDA-based proxyless implementations च्या कार्यक्षमतेपेक्षा पुढे जातो. आम्ही हेही दाखवतो की proxy implementation महत्त्वाची आहे, कारण आमचा अत्यंत specialized proxy generic NVSHMEM implementation पेक्षा एक order of magnitude ने जलद आहे. EFA हे 256KB packets हस्तांतरित करताना ConnectX-7 पेक्षा दुप्पट मंद असले, आणि हे packets MoE routing द्वारे देवाणघेवाण होणाऱ्या सामान्य packet size च्या जवळ असले, तरीही दोन्ही adapters मधील gap एकूणात इतका लक्षणीय नाही.

विस्तृत करा

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

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

आमच्या कर्नलची व्यवहार्यते सिद्ध करण्यासाठी, आम्ही इलास्टिक फॅब्रिक अडॅप्टर (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}
}