AWS Sagemaker Hyperpod पर उच्च-प्रदर्शन GPU मेमोरी ट्रांसफर

आधुनिक डीप लर्निंग इंफ्रास्ट्रक्चर के लिए अक्सर मशीनों के बीच GPU के मध्य बड़ी मात्रा में डेटा ट्रांसफर करने की आवश्यकता होती है। Perplexity में, हमें एक अद्वितीय तकनीकी चुनौती का सामना करना पड़ा: मशीनों के बीच गैर-सन्निकट GPU मेमोरी क्षेत्रों को अधिकतम संभव गति से कुशलतापूर्वक ट्रांसफर करना। हमारा लक्षित प्लेटफ़ॉर्म, AWS p5 इंस्टेंस, 32 नेटवर्क कार्ड के माध्यम से 3200 Gbps की प्रभावशाली नेटवर्क बैंडविड्थ प्रदान करता है। यह लेख एक कस्टम उच्च-प्रदर्शन नेटवर्किंग समाधान बनाने की हमारी यात्रा को साझा करता है जो इस सैद्धांतिक बैंडविड्थ का 97.1% प्राप्त करता है।

तकनीकी चुनौती

हमारे उपयोग के मामले में कई प्रमुख तकनीकी आवश्यकताएं थीं:

• दूरस्थ GPU के बीच गैर-सन्निकट मेमोरी चंक्स का उच्च-बैंडविड्थ ट्रांसफर
• चल रहे कार्यों को बाधित किए बिना कुबेरनेट्स (Kubernetes) डिप्लॉयमेंट से नोड्स को गतिशील रूप से जोड़ने या हटाने की क्षमता
• पीयर-टू-पीयर संचार पैटर्न के लिए समर्थन

जबकि NVIDIA की NCCL लाइब्रेरी वितरित डीप लर्निंग के लिए डी-फैक्टो मानक है, यह हमारे उपयोग के मामले के लिए आदर्श नहीं थी:

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

आधुनिक उच्च-प्रदर्शन नेटवर्क

हमारे समाधान को समझने के लिए, आइए पहले देखें कि आधुनिक उच्च-प्रदर्शन नेटवर्क पारंपरिक नेटवर्किंग से कैसे भिन्न हैं।

हम जो नेटवर्क दैनिक उपयोग करते हैं, उनमें से अधिकांश TCP/IP प्रोटोकॉल पर निर्भर करते हैं, जहाँ एप्लिकेशन सॉकेट का उपयोग करके ऑपरेटिंग सिस्टम कर्नेल के माध्यम से नेटवर्क कार्ड के साथ संचार करते हैं। हालाँकि, उच्च-प्रदर्शन नेटवर्क RDMA (रिमोट डायरेक्ट मेमोरी एक्सेस) का उपयोग करते हैं - एक पूरी तरह से अलग हार्डवेयर और सॉफ़्टवेयर स्टैक जो CPU को शामिल किए बिना मशीनों के बीच सीधे मेमोरी एक्सेस को सक्षम बनाता है।

AWS इलास्टिक फैब्रिक एडेप्टर (EFA) प्रदान करता है, जो एक कस्टम नेटवर्क इंटरफ़ेस है जो Amazon के स्केलेबल रिलायबल डेटाग्राम (SRD) नामक कस्टम प्रोटोकॉल को लागू करता है। पारंपरिक TCP/IP नेटवर्किंग के विपरीत, जहाँ डेटा को यूज़र स्पेस, कर्नेल स्पेस और नेटवर्क बफ़र्स के बीच कई बार कॉपी किया जाना चाहिए, RDMA के साथ EFA सीधे GPU मेमोरी और नेटवर्क कार्ड के बीच डेटा ट्रांसफर की अनुमति देता है, जो पूरी तरह से CPU को बायपास कर देता है।

उच्च-प्रदर्शन नेटवर्क डिज़ाइन का दर्शन

उच्च-प्रदर्शन नेटवर्किंग सिस्टम बनाने के लिए कई मौलिक धारणाओं पर पुनर्विचार करने की आवश्यकता होती है:

1. बफ़र स्वामित्व: पारंपरिक सॉकेट के विपरीत, जहाँ कर्नेल नेटवर्क बफ़र्स का प्रबंधन करता है और यूज़र स्पेस और कर्नेल स्पेस के बीच कॉपी करने की आवश्यकता होती है, RDMA के लिए आवश्यक है कि एप्लिकेशन अपने स्वयं के बफ़र्स का प्रबंधन करें। जब कोई एप्लिकेशन नेटवर्क ऑपरेशन शुरू करता है, तो वह ऑपरेशन पूरा होने तक बफ़र का स्वामित्व नेटवर्क कार्ड को हस्तांतरित कर देता है, जिससे डेटा कॉपी करने की आवश्यकता समाप्त हो जाती है।
2. मेमोरी पंजीकरण: एप्लिकेशनों को ऑपरेटिंग सिस्टम कर्नेल के साथ मेमोरी क्षेत्रों को पंजीकृत करना होगा। कर्नेल वर्चुअल एड्रेस मैपिंग स्थापित करता है जो CPU, GPU और नेटवर्क कार्ड को एक ही वर्चुअल एड्रेस को समझने की अनुमति देता है। यह पंजीकरण एक बार का ऑपरेशन है जो बाद में ज़ीरो-कॉपी डेटा ट्रांसफर को सक्षम बनाता है।
3. नियंत्रण तल बनाम डेटा तल: उच्च-प्रदर्शन नेटवर्क संचालन को दो श्रेणियों में विभाजित करते हैं:
   1. नियंत्रण तल (control plane) संचालन (जैसे कनेक्शन सेटअप और मेमोरी पंजीकरण) सुरक्षा सुनिश्चित करने के लिए कर्नेल के माध्यम से जाते हैं
   2. डेटा तल (data plane) संचालन (वास्तविक डेटा ट्रांसफर) अधिकतम प्रदर्शन के लिए कर्नेल को बायपास करते हैं
4. ट्रांसमिशन से पहले रिसेप्शन: कर्नेल-प्रबंधित बफ़र्स के बिना, एप्लिकेशनों को पहले से रिसीव ऑपरेशन पोस्ट करने होंगे, यह निर्दिष्ट करते हुए कि आने वाले डेटा को कहाँ रखा जाना चाहिए। यह सॉकेट मॉडल से एक मूलभूत बदलाव है जहाँ एप्लिकेशन किसी भी समय डेटा प्राप्त कर सकते हैं।
5. पोल-आधारित पूर्णता: epoll जैसे तंत्रों के माध्यम से कर्नेल सूचनाओं की प्रतीक्षा करने के बजाय, एप्लिकेशन सीधे हार्डवेयर कंप्लीशन क्यू को पोल करते हैं। यह सिस्टम कॉल ओवरहेड को समाप्त करता है और पूर्ण किए गए ऑपरेशनों पर तत्काल प्रतिक्रिया की अनुमति देता है।
6. हार्डवेयर टोपोलॉजी जागरूकता: अधिकतम प्रदर्शन प्राप्त करने के लिए हार्डवेयर टोपोलॉजी को समझना और अनुकूलित करना महत्वपूर्ण है।

हार्डवेयर टोपोलॉजी को समझना

AWS p5 इंस्टेंस में एक परिष्कृत आंतरिक आर्किटेक्चर होता है। जैसा कि नीचे दिखाया गया है, प्रत्येक इंस्टेंस में दो CPU सॉकेट होते हैं जो दो NUMA नोड्स बनाते हैं, जिसमें प्रत्येक NUMA नोड चार PCIe स्विच से जुड़ता है:

बड़ा करें

प्रत्येक PCIe स्विच के अंतर्गत, हमें ये मिलते हैं:

• चार 100 Gbps EFA नेटवर्क कार्ड
• एक NVIDIA H100 GPU
• एक NVMe SSD

TCP/IP और RDMA ट्रांसफर के लिए डेटा पथ उनके दृष्टिकोण में मूलभूत अंतर प्रदर्शित करते हैं:

बड़ा करें

TCP/IP (बाईं ओर) के साथ, डेटा को कई बार कॉपी किया जाना चाहिए:

1. GPU मेमोरी से यूज़र स्पेस में एप्लिकेशन बफ़र तक
2. यूज़र स्पेस बफ़र से कर्नेल बफ़र तक
3. कर्नेल बफ़र से नेटवर्क कार्ड तक
4. प्राप्त करने वाली तरफ विपरीत प्रक्रिया होती है

प्रत्येक कॉपी ऑपरेशन CPU चक्र और मेमोरी बैंडविड्थ की खपत करता है। एप्लिकेशन को प्रत्येक नेटवर्क ऑपरेशन के लिए यूज़र स्पेस और कर्नेल स्पेस के बीच संदर्भ स्विच (context switch) भी करना पड़ता है।

इसके विपरीत, RDMA (दाईं ओर) वास्तविक ज़ीरो-कॉपी डेटा ट्रांसफर प्रदान करता है:

1. नेटवर्क कार्ड सीधे GPU मेमोरी से पढ़ता है
2. डेटा सीधे रिमोट नेटवर्क कार्ड तक जाता है
3. रिमोट नेटवर्क कार्ड सीधे डेस्टिनेशन GPU मेमोरी में लिखता है
4. एप्लिकेशन को केवल यह जानने के लिए यूज़र स्पेस में एक कंप्लीशन क्यू (CQ) की जांच करनी होती है कि ट्रांसफर कब पूरा हो गया है

RDMA और उचित हार्डवेयर पेयरिंग के साथ, दो GPU के बीच डेटा ट्रांसफर करने के लिए केवल लोकल PCIe स्विच और नेटवर्क को पार करना होता है:

बड़ा करें

इसके विपरीत, TCP/IP ट्रांसफर को मुख्य मेमोरी के माध्यम से डेटा को कई बार कॉपी करना पड़ता है, जिससे PCIe बस में भारी भीड़ होती है:

बड़ा करें

libfabric के साथ निर्माण

हमने libfabric का उपयोग किया - एक फ्रेमवर्क जो फैब्रिक सेवाओं के लिए एक सामान्य इंटरफ़ेस प्रदान करता है। हमारा कार्यान्वयन दो प्रकार के RDMA संचालन का उपयोग करता है:

• नियंत्रण संदेशों के लिए टू-साइडेड RDMA (SEND/RECV) जो मेमोरी क्षेत्रों (जैसे, ऑफ़सेट और आकार) के बारे में मेटाडेटा ले जाते हैं।
• वास्तविक डेटा ट्रांसफर के लिए वन-साइडेड RDMA WRITE, जहाँ प्रत्येक WRITE ऑपरेशन एक सन्निकट मेमोरी चंक को संभालता है।

शीर्ष प्रदर्शन की यात्रा

हमारा विकास कई चरणों से होकर गुजरा:

1. SEND/RECV का उपयोग करके बुनियादी यूनिडायरेक्शनल संदेश ट्रांसफर लागू किया
2. बीडायरेक्शनल संचार तक विस्तारित किया
3. सीधे GPU-GPU ट्रांसफर के लिए GPUDirect RDMA WRITE जोड़ा
4. एकाधिक समवर्ती ट्रांसफर को संभालने के लिए विस्तारित किया
5. मजबूती के लिए ऑपरेशन क्यूइंग पेश की
6. एकल नेटवर्क कार्ड पर 97.4% बैंडविड्थ उपयोग प्राप्त किया

32 नेटवर्क कार्ड तक स्केल करते समय, हमने कई महत्वपूर्ण अनुकूलन लागू किए:

1. ऑपरेशन क्यूइंग: नेटवर्क कार्ड पर सीधे ऑपरेशन सबमिट करने के बजाय, हम एक एप्लिकेशन-स्तरीय कतार बनाए रखते हैं। यह नेटवर्क कंजेशन के खिलाफ मजबूती प्रदान करता है और प्रोग्रामिंग मॉडल को सरल बनाता है।
2. नेटवर्क वार्मअप: कनेक्शन को पहले से स्थापित करने से स्टार्टअप प्रदर्शन में सुधार हुआ।
3. मल्टी-थ्रेडिंग: प्रत्येक GPU के नेटवर्क ऑपरेशनों के लिए समर्पित थ्रेड।
4. CPU कोर पिनिंग: NUMA प्रभावों और कैश मिस से बचने के लिए थ्रेड को विशिष्ट CPU कोर से बांधना।
5. स्टेट शार्डिंग: थ्रेड के बीच विवाद को कम करना।
6. ऑपरेशन बैचिंग: कई ऑपरेशनों को एक साथ सबमिट करना।
7. लेज़ी ऑपरेशन पोस्टिंग: ऑपरेशन पहले एप्लिकेशन में कतारबद्ध किए जाते हैं। कंप्लीशन क्यू को पोल करने के बाद, हम नेटवर्क कार्ड पर लंबित ऑपरेशनों को सबमिट करने का प्रयास करते हैं, जिससे नेटवर्क संसाधनों का कुशल उपयोग सुनिश्चित होता है।
8. NUMA-जागरूक संसाधन आवंटन: मेमोरी एक्सेस विलंबता को कम करने के लिए सही NUMA नोड पर कंप्लीशन क्यू जैसे libfabric संसाधनों को आवंटित करें।

इन अनुकूलनों के माध्यम से, हमने 3,108 Gbps का अंतिम प्रदर्शन प्राप्त किया - सैद्धांतिक अधिकतम बैंडविड्थ का 97.1%।

नीचे दिया गया वीडियो हमारे कमांड-लाइन प्रोग्राम को क्रियान्वित होते हुए दिखाता है। यह GPU मेमोरी के गैर-सन्निकट चंक्स को रिमोट नोड में ट्रांसफर करता है, 3108.283 Gbps की ट्रांसमिशन गति प्राप्त करता है - जो नेटवर्क इंफ्रास्ट्रक्चर के लगभग-सैद्धांतिक बैंडविड्थ उपयोग को प्रदर्शित करता है:

निष्कर्ष

उच्च-प्रदर्शन नेटवर्किंग सिस्टम बनाने के लिए हार्डवेयर आर्किटेक्चर और सिस्टम डिज़ाइन सिद्धांतों दोनों को समझने की आवश्यकता होती है। जबकि NCCL जैसी लाइब्रेरी सामान्य पैटर्न के लिए उत्कृष्ट समाधान प्रदान करती हैं, कभी-कभी विशिष्ट आवश्यकताओं के लिए कस्टम समाधान आवश्यक होते हैं।

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

कार्यान्वयन विवरण और कोड उदाहरणों सहित इस यात्रा का पूर्ण तकनीकी विवरण, हमारे ओपन सोर्स रिपॉजिटरी में उपलब्ध है।

उद्धरण

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}