Triton for MTIA: Bridging the Programming Model Gaps for Custom AI Accelerators
यह शोध पत्र मेटा के कस्टम MTIA-2i एक्सेलेरेटर पर ट्राइटन प्रोग्रामिंग भाषा के पहले प्रोडक्शन-स्केल पर तैनाती को प्रस्तुत करता है, जो यह प्रदर्शित करता है कि एक नया कंपाइलर बैकएंड और न्यूनतम भाषा विस्तार उच्च-प्रदर्शन, कुशल कर्नेल विकास को सक्षम करते हैं जो एमएल फ्रेमवर्क और कस्टम हार्डवेयर के बीच के अंतर को पाटता है।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक महाद्वीप के पार यात्रियों (डेटा) को ले जाने के लिए एक विशाल, हाई-स्पीड ट्रेन सिस्टम बनाने की कोशिश कर रहे हैं। वर्षों तक, केवल मानक, हाई-टेक इलेक्ट्रिक ट्रेनें ही उपलब्ध थीं (GPUs) जिन्हें चलाना हर कोई जानता था। उनमें स्वचालित दरवाजे, बिल्ट-इन जीपीएस और एक सुचारू सवारी थी। लेकिन हाल ही में, एक नए प्रकार के ट्रेन इंजन का आविष्कार किया गया है (कस्टम AI एक्सेलेरेटर्स)। ये नए इंजन अविश्वसनीय रूप से शक्तिशाली हैं और विशेष रूप से भारी कार्गो के लिए बनाए गए हैं, लेकिन वे अजीब हैं। उनमें स्वचालित दरवाजे नहीं हैं; आपको उन्हें मैन्युअल रूपel रूप से खोलना और बंद करना पड़ता है। उनमें जीपीएस नहीं है; आपको अपना नक्शा खुद बनाना पड़ता है। और वे पुराने ट्रेनों वाले ट्रैक पर नहीं चलते हैं।
समस्या यह है कि जो इंजीनियर ट्रेन के डिब्बे बनाते हैं (सॉफ्टवेयर डेवलपर्स), वे पुराने, आसानी से चलने वाले इलेक्ट्रिक ट्रेनों के आदी हैं। यदि वे नए, सुपर-फास्ट इंजन का उपयोग करना चाहते हैं, तो उन्हें या तो ड्राइविंग का एक पूरी तरह से नया, कठिन तरीका सीखना होगा, या उन्हें हर एक नए रूट के लिए एक बिल्कुल नया ट्रेन डिब्बा बनाना होगा। इससे सब कुछ धीमा हो जाता है। बड़ा सवाल यह है: क्या हम पुराने, आसान-से-चलाने वाले ड्राइविंग स्टाइल को इन अजीब, नए इंजनों पर काम करने के लिए सिखा सकते हैं बिना गति का लाभ खोए? यह पेपर ठीक यही तलाशता है, जो मानक AI सॉफ्टवेयर और उन कस्टम हार्डवेयर (जैसे मेटा द्वारा बनाए गए) के बीच के अंतर को पाटने की कोशिश करता है।
कहानी: एक नए इंजन को पुरानी भाषा बोलना सिखाना
मेटा, जो फेसबुक और इंस्टाग्राम के पीछे की कंपनी है, अपने AI मॉडल को तेज़ी से और सस्ते में चलाने के लिए MTIA-2i नामक अपने विशेष कंप्यूटर चिप्स बना रहा है। MTIA-2i को एक कस्टम-बिल्ट रेस कार इंजन की तरह समझें। यह जो यह करता है उसमें अद्भुत है, लेकिन यह मानक इंजनों (GPUs) से अलग बनाया गया है जिनके आदी अधिकांश AI प्रोग्रामर हैं।
AI की दुनिया में, एक लोकप्रिय प्रोग्रामिंग भाषा है जिसे Triton कहा जाता है। आप Triton को AI के लिए एक "यूनिवर्सल रिमोट कंट्रोल" के रूप में देख सकते हैं। एक मानक GPU को डेटा कैसे हिलाना है यह बताने के लिए जटिल, बिखरा हुआ कोड लिखने के बजाय, डेवलपर्स Triton में सरल, साफ कोड लिखते हैं, और एक स्मार्ट कंपाइलर इसे उस मशीन भाषा में अनुवादित करता है जिसे GPU समझता है। यह एक अनुवादक से अंग्रेजी बोलने जैसा है जो मशीन से बात करना जानता है।
हालाँकि, Triton मूल रूप से केवल मानक GPU इंजनों के लिए डिज़ाइन किया गया था। जब मेटा ने अपने कस्टम MTIA-2i रेस कार इंजन पर Triton का उपयोग करने की कोशिश की, तो यह काम नहीं किया। "रिमोट कंट्रोल" नए इंजन के बटनों में फिट नहीं बैठा। MTIA-2i इंजन एसिंक्रोनसली (asynchronously) काम करता है (यह आदेश लेता है और उन्हें बाद में पूरा करता है), यह अपने मेमोरी को एक अजीब "सर्कुलर बफ़र" तरीके से मैनेज करता है, और इसे सब कुछ सुचारू रूप से चलाने के लिए बहुत विशिष्ट निर्देशों की आवश्यकता होती है। मानक Triton कोड इन बारीकियों को नहीं संभाल सका।
यह पेपर क्या करता है: एक नया अडैप्टर बनाना
मेटा की टीम ने इसे ठीक करने का फैसला किया। उन्होंने केवल पुराने रिमोट को जबरदस्ती काम कराने की कोशिश नहीं की; उन्होंने विशेष रूप से MTIA-2i के लिए एक नया कंपाइलर बैकएंड बनाया। यह नया सिस्टम एक सुपर-स्मार्ट अडैप्टर की तरह काम करता है। यह उस सरल, आसानी से पढ़े जाने वाले Triton कोड को लेता है जिसे डेवलपर्स लिखते हैं और इसे पूरी तरह से उन जटिल, लो-लेवल निर्देशों में अनुवादित करता है जिनकी MTIA-2i इंजन को आवश्यकता होती है।
उन्होंने इसे कैसे काम करने के योग्य बनाया, इसके लिए कुछ मजेदार उपमाओं का उपयोग किया है:
- "FIFO" बुफे लाइन: मानक GPUs पर, कंप्यूटर मेमोरी को एक स्मार्ट बुफे की तरह प्रबंधित करता है जहाँ प्लेटें तुरंत ली और बदली जाती हैं। MTIA-2i पर, यह एक मैनुअल असेंबली लाइन की तरह है जिसमें "फर्स्ट-इन, फर्स्ट-आउट" (FIFO) बफर है। नए कंपाइलर ने इस लाइन को पूरी तरह से प्रबंधित करना सीखा, यह सुनिश्चित करते हुए कि डेटा ठीक उसी समय पहुंचे जब मशीन को इसकी आवश्यकता हो, बिना प्रोग्रामर को हर एक प्लेट को मैन्युअल रूप से धकेलने पड़े।
- दो शेफ वाली रसोई: MTIA-2i चिप में दो "कोर्स" (रसोई में दो शेफ की तरह) होते हैं जो एक साथ काम कर सकते हैं। कंपाइलर ने यह पता लगाया कि खाना पकाने के कार्यों को उनके बीच कैसे विभाजित किया जाए ताकि वे एक-दूसरे से न टकराएं या खाली न बैठें। यह विशेषज्ञों को यह भी करने देता है कि वे चाहें तो फैंसी काम करने के लिए शेफ को सटीक रूप से बता सकें कि प्याज कौन काटेगा और सूप कौन चलाएगा।
- "जिगज़ैग" डिलीवरी रूट: कल्पना कीजिए कि आपको 64 घरों में पैकेज पहुँचाने हैं। यदि आप बस सड़क पर एक सीधी रेखा में (लीनियर) जाते हैं, तो पहले कुछ घर जल्दी पूरे हो जाते हैं, लेकिन आखिरी वाले बहुत देर तक इंतजार करते हैं। टीम ने पाया कि एक "जिगज़ैग" डिलीवरी रूट (आगे-पीछे जाना) काम को बहुत बेहतर तरीके से संतुलित करता है। उन्होंने Triton में एक फीचर जोड़ा ताकि प्रोग्रामर आसानी से इस स्मार्ट रूट को चुन सकें।
परिणाम: तेज़, लचीला और वास्तविक दुनिया के लिए तैयार
टीम ने इसे केवल लैब में नहीं बनाया; उन्होंने इसे वास्तविक दुनिया में लागू किया। उन्होंने सफलतापूर्वक इन नए Triton कर्नेल को 60 अलग-अलग प्रकार के AI मॉडल में प्रोडक्शन में तैनात किया।
यहाँ उन्हें क्या मिला:
- गति: सरल Triton भाषा में लिखा गया कोड, जटिल, लो-लेवल C++ (चिप की "नेटिव" भाषा) में लिखे गए कोड के समान ही तेज़ चला। वास्तव में, जनरल मैट्रिक्स मल्टीप्लिकेशन (GEMM)—एक भारी गणितीय कार्य—के लिए, Triton कोड ने चिप की सैद्धांतिक अधिकतम गति के 80% से अधिक को प्राप्त किया।
- अपनाने की दर: केवल एक तिमाही के भीतर, Triton का उपयोग करने वाले मॉडल प्रकारों की संख्या तीन गुना बढ़ गई।
- प्रभाव: Triton अब इन मॉडलों के 50% लेयर्स और 47% नॉन-GEMM निष्पादन समय को संभालता है। यह लगभग आधा काम इस नए, आसान-से-उपयोग वाले सिस्टम द्वारा किया जा रहा है!
यह क्यों महत्वपूर्ण है
यह पेपर तर्क देता है कि आपको "आसान लिखना" और "सुपर फास्ट" के बीच चुनाव करने की आवश्यकता नहीं है। इससे पहले, यदि आप MTIA जैसे कस्टम चिप का उपयोग करना चाहते थे, तो आपको एक लो-लेवल कोडिंग जादूगर होना पड़ता था और सब कुछ शून्य से लिखना पड़ता था। यदि आप Triton जैसी हाई-लेवल भाषा का उपयोग करना चाहते थे, तो आप मानक GPUs तक ही सीमित थे।
इस अंतर को पाटकर, मेटा ने दिखाया कि हाई-लेवल भाषाओं को लगभग किसी भी कस्टम हार्डवेयर पर काम करने के लिए अनुकूलित किया जा सकता है। इसका मतलब है कि भविष्य में, जैसे-जैसे कंपनियां और भी अजीब और विशिष्ट AI चिप्स बनाती हैं, डेवलपर्स को हर एक के लिए एक नई, कठिन भाषा सीखने की आवश्यकता नहीं होगी। वे उन्हीं परिचित, शक्तिशाली उपकरणों का उपयोग करना जारी रख सकते हैं जिन्हें वे पहले से जानते हैं, जबकि कंपाइलर नए हार्डवेयर के जटिल विवरणों को संभाल लेगा।
पेपर यह भी नोट करता है कि उन्होंने पहले KNYFE नामक एक अलग, बहुत ही लो-लेवल भाषा का परीक्षण किया था, लेकिन उन्होंने इसे छोड़ दिया क्योंकि इसे सीखना अधिकांश लोगों के लिए बहुत कठिन था। उन्होंने साबित किया कि एक लचीली, हाई-लेवल भाषा जैसे Triton पर टिके रहना, भले ही उसमें कुछ छोटे कस्टम बदलाव हों, विकास को तेज़ और कुशल बनाए रखने के लिए जीतने वाली रणनीति है।
संक्षेप में, यह पेपर प्रदर्शित करता है कि सही कंपाइलर जादू के साथ, हम "यूनिवर्सल रिमोट" को लगभग किसी भी "टीवी" पर काम करने के लायक बना सकते हैं, यहाँ तक कि उन अजीब, कस्टम-बिल्ट टीवीों पर भी, बिना पिक्चर क्वालिटी से समझौता किए।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।