AWS Sagemaker Hyperpod वरील उच्च-कार्यक्षमता GPU मेमरी ट्रान्सफर

आधुनिक डीप लर्निंग इन्फ्रास्ट्रक्चरसाठी अनेकदा मशिन्सवर जीपीयू दरम्यान मोठ्या प्रमाणावर डेटा ट्रान्सफर करणे आवश्यक असते. Perplexity येथे, आम्हाला एका अनन्य तांत्रिक आव्हानाचा सामना करावा लागला: संभाव्य कमाल वेगाने मशिन्स दरम्यान नॉन-कॉन्टीग्युअस जीपीयू मेमरी क्षेत्रांचे कार्यक्षमतेने ट्रान्सफर करणे. आमचे लक्ष्यित प्लॅटफॉर्म, AWS p5 इन्स्टन्सेस, ३२ नेटवर्क कार्ड्सद्वारे ३२०० जीबीपीएस इतका प्रभावी नेटवर्क बँडविथ देतात. हा लेख आमच्या कस्टम उच्च-कार्यक्षमता नेटवर्किंग सोल्यूशन तयार करण्याच्या प्रवासाची माहिती देतो, ज्याने या सैद्धांतिक बँडविथच्या ९७.१% भाग साध्य केला आहे.

तांत्रिक आव्हान

आमच्या युझ केसमध्ये अनेक महत्त्वाच्या तांत्रिक गरजा समोर आल्या:

• रिमोट जीपीयूच्या नॉन-कॉन्टीग्युअस मेमरी चंक्स दरम्यान हाय-बँडविथ ट्रान्सफर
• Kubernetes उपयोजनांमधून चालू ऑपरेशन्समध्ये अडथळा न आणता नोड्स गतिशीलरीत्या जोडण्याची किंवा काढून टाकण्याची क्षमता
• पीअर-टू-पीअर कम्युनिकेशन पॅटर्न्ससाठी समर्थन

एनव्हियाची (NVIDIA's) NCCL लायब्ररी ही डिस्ट्रिब्युटेड डीप लर्निंगसाठी डी फॅक्टो स्टँडर्ड असली, तरी ती आमच्या युझ केससाठी आदर्श नव्हती:

• NCCL हे सामूहिक संवादात उत्कृष्ट आहे परंतु यासाठी एक स्थिर "जगत" (world) स्थापित करणे आवश्यक आहे, ज्यामध्ये सहभागी नोड्स समायोजित करताना संपूर्ण क्लस्टर रीस्टार्ट करणे आवश्यक असते.
• NCCL चा सिंक्रोनस कम्युनिकेशन मॉडेल आमच्या असिंक्रोनस वर्कलोडसाठी गुंतागुंत वाढवतो
• ऑप्टिमायझेशनसाठी आम्हाला आमच्या मेमरी ट्रान्सफर पॅटर्न्सवर थेट नियंत्रण हवे होते
• आमचे स्वतःचे सोल्यूशन तयार केल्यामुळे मूल्यवान शिकण्याच्या संधी मिळाल्या

आधुनिक उच्च-कार्यक्षमता नेटवर्क्स

आमचे सोल्यूशन समजून घेण्यासाठी, आधुनिक उच्च-कार्यक्षमता नेटवर्क्स पारंपारिक नेटवर्किंगपेक्षा कसे भिन्न आहेत हे आपण आधी पाहूया.

आपण दररोज वापरत असलेली बहुतेक नेटवर्क्स TCP/IP प्रोटोकॉलवर अवलंबून असतात, जिथे ॲप्लिकेशन्स सॉकेट्स वापरून ऑपरेटिंग सिस्टम कर्नलद्वारे नेटवर्क कार्डशी संवाद साधतात. तथापि, उच्च-कार्यक्षमता नेटवर्क्स RDMA (रिमोट डायरेक्ट मेमरी ॲक्सेस) वापरतात - एक संपूर्णपणे भिन्न हार्डवेक आणि सॉफ्टवेअर स्टॅक जो सीपीयूचा समावेश न करता मशिन्स दरम्यान थेट मेमरी प्रवेश सक्षम करतो.

AWS इलास्टिक फॅब्रिक अडॅप्टर (EFA) पुरवते, जे एक कस्टम नेटवर्क इंटरफेस आहे आणि Amazon चा स्केलेबल रिलायबल डेटाग्राम (SRD) नावाचा कस्टम प्रोटोकॉल लागू करतो. पारंपारिक TCP/IP नेटवर्किंगच्या विपरीत जिथे युझर स्पेस, कर्नल स्पेस आणि नेटवर्क बफर्स दरम्यान डेटा अनेक वेळा कॉपी केला जावा लागतो, RDMA सह EFA मुळे जीपीयू मेमरी आणि नेटवर्क कार्ड दरम्यान थेट डेटा ट्रान्सफर शक्य होते, ज्यामुळे सीपीयू पूर्णपणे बायपास होतो.

उच्च-कार्यक्षमता नेटवर्क डिझाइनचे तत्त्वज्ञान

उच्च-कार्यक्षमता नेटवर्किंग सिस्टम तयार करण्यासाठी अनेक मूलभूत गृहीतकांचा फेरविचार करावा लागतो:

1. बफर मालकी: पारंपारिक सॉकेट्सच्या विपरीत जिथे कर्नल नेटवर्क बफर्स व्यवस्थापित करतो आणि युझर स्पेस व कर्नल स्पेस दरम्यान कॉपी करणे आवश्यक असते, RDMA मध्ये ॲप्लिकेशन्सना त्यांचे स्वतःचे बफर्स व्यवस्थापित करणे आवश्यक असते. जेव्हा एखादे ॲप्लिकेशन नेटवर्क ऑपरेशन सुरू करते, तेव्हा ते ऑपरेशन पूर्ण होईपर्यंत बफर मालकी नेटवर्क कार्डकडे हस्तांतरित करते, ज्यामुळे डेटा कॉपी करण्याची गरज नाहीशी होते.
2. मेमरी नोंदणी: ॲप्लिकेशन्सना ऑपरेटिंग सिस्टम कर्नलसह मेमरी क्षेत्रे नोंदणीकृत करणे आवश्यक आहे. कर्नल व्हर्च्युअल ॲड्रेस मॅपिंग सेट करतो ज्यामुळे CPU, GPUs आणि नेटवर्क कार्ड्सना एकच व्हर्च्युअल ॲड्रेस समजणे शक्य होते. ही नोंदणी एक वेळचे ऑपरेशन आहे जे त्यानंतरच्या झिरो-कॉपी डेटा ट्रान्सफर सक्षम करते.
3. कंट्रोल प्लेन विरुद्ध डेटा प्लेन: उच्च-कार्यक्षमता नेटवर्क्स ऑपरेशन्सना दोन श्रेणींमध्ये विभागतात:
   1. सुरक्षितता सुनिश्चित करण्यासाठी नियंत्रण विमान ऑपरेशन्स (जसे की कनेक्शन सेटअप आणि मेमरी नोंदणी) कर्नल मार्गे जातात
   2. कमाल कार्यक्षमतेसाठी डेटा प्लेन ऑपरेशन्स (वास्तविक डेटा ट्रान्सफर) कर्नलला बायपास करतात
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 सायकल आणि मेमरी बँडविथ वापरला जातो. ॲप्लिकेशनला प्रत्येक नेटवर्क ऑपरेशनसाठी युझर स्पेस आणि कर्नल स्पेस दरम्यान कॉंटेक्स्ट स्विच देखील करावे लागते.

याउलट, RDMA (उजवी बाजू) खरी झिरो-कॉपी डेटा ट्रान्सफर प्रदान करते:

1. नेटवर्क कार्ड थेट GPU मेमरीवरून वाचते
2. डेटा थेट रिमोट नेटवर्क कार्डकडे प्रवास करतो
3. रिमोट नेटवर्क कार्ड थेट गंतव्यस्थान GPU मेमरीमध्ये लिहितो
4. ट्रान्सफर पूर्ण झाल्याचे जाणून घेण्यासाठी ॲप्लिकेशनला फक्त युझर स्पेस मधील पूर्णता रांगे (CQ) तपासावी लागते

RDMA आणि योग्य हार्डवेअर पेअरिंगसह, दोन GPUs दरम्यान डेटा ट्रान्सफर करण्यासाठी फक्त स्थानिक 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. मजबूतीसाठी ऑपरेशन क्यूईंग (queuing) सुरू केले
6. एकाच नेटवर्क कार्डवर ९७.४% बँडविथ वापर साध्य केला

32 नेटवर्क कार्ड्सपर्यंत स्केल करताना, आम्ही अनेक महत्त्वपूर्ण ऑप्टिमायझेशन लागू केले:

1. ऑपरेशन क्यूईंग: नेटवर्क कार्ड्सवर थेट ऑपरेशन्स सबमिट करण्याऐवजी, आम्ही ॲप्लिकेशन-स्तरीय रांग राखतो. हे नेटवर्कच्या गर्दीपासून मजबूती प्रदान करते आणि प्रोग्रामिंग मॉडेल सुलभ करते.
2. नेटवर्क वॉर्मअप: कनेक्शन पूर्व-स्थापित केल्याने स्टार्टअप कार्यक्षमता सुधारली.
3. मल्टी-थ्रेडिंग: प्रत्येक GPU च्या नेटवर्क ऑपरेशन्ससाठी समर्पित थ्रेड्स.
4. CPU कोर पिनिंग: NUMA प्रभावांचे आणि कॅशे मिस टाळण्यासाठी विशिष्ट CPU कोरांना थ्रेड्स बाईंड करणे.
5. स्टेट शार्डिंग: थ्रेड्स मधील स्पर्धा कमी करणे.
6. ऑपरेशन बॅचिंग: एकाधिक ऑपरेशन्स एकत्र सबमिट करणे.
7. लेझी ऑपरेशन पोस्टिंग: ऑपरेशन्स प्रथम ॲप्लिकेशनमध्ये रांगेत ठेवल्या जातात. कम्प्लिशन रांग पोल केल्यानंतर, आम्ही नेटवर्क संसाधनांचा कार्यक्षम वापर सुनिश्चित करून, प्रलंबित ऑपरेशन्स नेटवर्क कार्डवर सबमिट करण्याचा प्रयत्न करतो.
8. NUMA-जागरूक संसाधन वाटप: मेमरी ॲक्सेस लेटन्सी कमी करण्यासाठी योग्य NUMA नोडवर कम्प्लिशन रांगांसारख्या libfabric संसाधनांचे वाटप करा.

या ऑप्टिमायझेशन्सद्वारे, आम्ही ३,१०८ जीबीपीएस ची अंतिम कार्यक्षमता गाठली - जी सैद्धांतिक कमाल बँडविथच्या ९७.१% आहे.

खालील व्हिडिओ आमचा कमांड-लाइन प्रोग्राम कृतीत दाखवतो. हे रिमोट नोडला GPU मेमरीचे नॉन-कॉन्टीग्युअस चंक्स ट्रान्सफर करते, ३१०८.२८३ जीबीपीएस चा ट्रान्समिशन स्पीड साध्य करते - नेटवर्क इन्फ्रास्ट्रक्चरचा जवळपास-सैद्धांतिक बँडविथ वापर प्रदर्शित करते:

निष्कर्ष

उच्च-कार्यक्षमता नेटवर्किंग सिस्टम तयार करण्यासाठी हार्डवेअर आर्किटेक्चर आणि सिस्टम डिझाइन तत्त्वे दोन्ही समजून घेणे आवश्यक आहे. 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}
}