जलापेनो MLA कर्नेल की प्रस्तुति HotChips में और SemiAnalysis द्वारा उसके बाद की गई टिप्पणियों के बाद काफी चर्चा हुई है। OpenAI की हार्डवेयर टीम के रूप में, हमने इस छोटे से सुनहरे टुकड़े को बमुश्किल छुआ: तथ्य यह है कि AI हमारे कर्नेल लिख रहा है, और जब वह ऐसा करता है, तो हमें वास्तव में यह समझने की आवश्यकता नहीं है कि कर्नेल लाइन दर लाइन क्या करता है। हमने स्पष्ट रूप से यह छोड़ दिया: ऐसी चीज़ कैसे संभव है? इसके बारे में सोचने का सही तरीका क्या है, बनाम कोड जनरेशन की अधिक पारंपरिक विधि? क्या अनुकूलित कर्नेल उतना ही सुदृढ़ है जितना कि अनुकूलित न किया गया कर्नेल?
अपनी पृष्ठभूमि के लिए, मैंने एक दशक से अधिक समय तक एक्सेलेरेटर के लिए कंपाइलरों पर काम किया है। मैंने XLA शुरू किया, जो एक उत्कृष्ट कंपाइलर इंफ्रास्ट्रक्चर है जिसमें एक शानदार क्रॉस-कंपनी टीम और प्रयास काम कर रहा है। पिछले 2+ वर्षों से OpenAI में मैं यह पुनर्विचार करने का प्रयास कर रहा हूं कि AI के युग में कंपाइलरों को कैसे काम करना चाहिए। नए कंपाइलर फॉर्मूलेशन मौजूदा ताकतों पर आकर्षित होंगे, लेकिन इस बात से इनकार करना असंभव है कि टूलकिट में लाभ उठाने के लिए एक शक्तिशाली नया उपकरण है।
यह थोड़ी यात्रा होगी, लेकिन मुझे उम्मीद है कि यह स्पष्ट करेगा कि कंप्यूटर प्रोग्राम सुधार के स्वचालन के लिए AI का उपयोग कैसे किया जा रहा है; यानी अनुकूलन संकलन (optimizing compilation)। मुझे लगता है, AI के माध्यम से, हम कुछ ऐसा अनुभव कर सकते हैं जिसे हम "कंपाइलर 2.0" समझते हैं। AI मौलिक रूप से इस बात में कम सीमित है कि वह क्या प्रस्तावित कर सकता है, और वह जो प्रस्तावित करता है वह मॉडल के प्रशिक्षण और संदर्भ का परिणाम है, यह मुझे इसे "स्टोकेस्टिक ऑप्टिमाइज़र" के रूप में वर्गीकृत करने के लिए प्रेरित करता है – यह चुनौतियाँ पेश कर सकता है, लेकिन जैसा कि हम देखेंगे, महान शक्तियों का स्रोत भी है...
शैक्षणिक अनुसंधान और उद्योग अनुप्रयोग का एक बड़ा हिस्सा पहले से ही इस दिशा में अग्रसर है, और अनुकूलन कंपाइलर क्षेत्र में AI की भागीदारी की क्षमता को तेजी से उजागर कर रहा है, लेकिन हम एक ऐसे बिंदु पर हैं जहां यह एक व्यापक स्पष्टीकरण की मांग करता है।
पृष्ठभूमि
कंपाइलर प्रोग्राम लेते हैं और उन प्रोग्रामों के अनुवादित या बेहतर संस्करण आउटपुट करते हैं।
प्रोग्राम, इनपुट और आउटपुट दोनों पक्षों पर, सिमैंटिक्स (semantics) रखते हैं जो हमें बताते हैं कि प्रोग्राम का क्या अर्थ है, वे संभावित रूप से क्या कर सकते हैं, और उन चीजों के बारे में कैसे तर्क करें जो वे कर सकते हैं।
हममें से जो कंपाइलरों पर काम करते हैं, वे उन्हें शुद्ध फलनों (pure functions) की तरह सोचते हैं – वे एक डेटा संरचना लेते हैं और एक डेटा संरचना आउटपुट करते हैं जिसमें संगत सिमैंटिक्स होने चाहिए।
कभी-कभी हमारे कंपाइलर "लोअरिंग" या "अनुवाद" पर ध्यान केंद्रित करते हैं। उदाहरण के लिए, वे C ले सकते हैं और x86-64 असेंबली आउटपुट कर सकते हैं, जिसे हम अक्सर "निचला स्तर" मानते हैं। लेकिन अक्सर वे उस प्रक्रिया के एक उप-भाग के रूप में केवल अनुवाद से अधिक कर रहे होते हैं...
हमारे कंपाइलर, व्यवहार में, "अनुकूलन" पर ध्यान केंद्रित करते हैं। वे एक डेटा संरचना ले सकते हैं जो प्रोग्राम का प्रतिनिधित्व करती है – हमारी भाषा में एक "इंटरमीडिएट रिप्रेजेंटेशन" (IR) – और वे उस प्रोग्राम का एक बेहतर संस्करण तैयार करने का प्रयास करते हैं। कभी-कभी "बेहतर" का अर्थ है कि इसे चलने में कम चक्र लगते हैं, कभी-कभी इसका अर्थ है कि इसमें कम अनावश्यक कोड होगा, कभी-कभी इसका अर्थ है उन चीजों के लिए विशेषज्ञता (specializing) जिन्हें हम प्रोग्राम के बारे में "सत्य होना चाहिए" साबित कर सकते हैं (आंशिक मूल्यांकन)।
अब, संक्षेप में, विचार करें कि LLM मूल रूप से मानव पाठ का एक भाषा से दूसरी भाषा में अनुवाद करने के लिए बनाए गए थे। स्पष्ट रूप से अनुवाद उनके क्षेत्र में है। और हम दैनिक कार्यों पर LLM के अपने उपयोग के माध्यम से देख सकते हैं कि वे नए समाधान भी लिख सकते हैं और मौजूदा समाधानों में सुधार कर सकते हैं। हममें से कई कोडर्स को LLM से "इस कोड स्निपेट को ऑप्टिमाइज़ करें" पूछने का अनुभव भी है और वे उल्लेखनीय रूप से ऐसा कर सकते हैं। (हालांकि, हमें यह जानने की आवश्यकता है कि उन्होंने कोड को सही ढंग से ऑप्टिमाइज़ किया है, जिस पर हम आएंगे!) यह केवल इस बात पर प्रकाश डालने के लिए है कि LLM में वे क्षमताएं हैं जिनकी हम एक अनुकूलन कंपाइलर में तलाश करते हैं।
अनुकूलन और इष्टतमता (Optimization and Optimality)
अनुकूलन कंपाइलर, आश्चर्य की बात नहीं, किसी उद्देश्य (आमतौर पर निष्पादन समय) द्वारा, जिस प्रोग्राम पर वे काम कर रहे हैं, उसकी इष्टतमता बढ़ाने का प्रयास कर रहे हैं। सामान्य मामले में, एक मनमाने प्रोग्राम के लिए ऐसा करना इतना कठिन है कि कंपाइलर इंजीनियरों के लिए पूर्ण रोजगार प्रमेय नामक एक प्रमेय है। (मुझे यह तब पता चला जब मैंने कंपाइलर इंजीनियर बनना चुना, लेकिन इसने मुझे आराम ही दिया!)
"सुपरऑप्टिमाइज़र" अनुकूलन कंपाइलरों का एक अद्भुत छोटा उप-क्षेत्र है। कल्पना करें कि एक दिया गया प्रोग्राम है, और हम सिमैंटिक्स के माध्यम से कह सकते हैं कि यह क्या करता है। वह सबसे इष्टतम प्रोग्राम क्या है जिसमें वही सिमैंटिक्स हैं? सुपरऑप्टिमाइज़र यही हल करने का प्रयास करते हैं, और यह प्रभावी रूप से एक खोज समस्या है...
कल्पना करें कि मैं सबसे छोटा प्रोग्राम खोजने का प्रयास कर रहा हूं जिसमें वही सिमैंटिक्स थे, और मेरे पास यह पूछने का एक तरीका था कि क्या किसी उम्मीदवार प्रोग्राम में वही सिमैंटिक्स हैं। मैं, काल्पनिक रूप से, उद्देश्य क्रम में हर प्रोग्राम की गणना कर सकता हूं, और सबसे छोटा चुन सकता हूं जिसमें वही सिमैंटिक्स थे।
हालांकि, उद्देश्य क्रम में हर प्रोग्राम की गणना करना काफी दुर्गम लगता है। 2013 में "STOKE" (स्टोकेस्टिक सुपरऑप्टिमाइज़ेशन) शीर्षक से मेरे पसंदीदा शैक्षणिक पेपरों में से एक ने पूछा: "ठीक है, क्या होगा अगर हम बार-बार बेतरतीब ढंग से प्रोग्रामों को ट्वीक करें, तो क्या हम अंततः सबसे अच्छा प्रोग्राम देखेंगे?" उन्होंने प्रस्तावित किया कि एक यादृच्छिक चाल (random walk) के माध्यम से (और हमारे OG मशीन लर्निंग मित्र मार्कोव चेन मोंटे कार्लो / मेट्रोपोलिस-हेस्टिंग्स के साथ), अंततः आप वह इष्टतम प्रोग्राम देखेंगे।
मोंटे कार्लो ट्वीकिंग आमतौर पर बेवकूफी भरी होती है (आप बेतरतीब ढंग से एक ट्वीक चुनते हैं), लेकिन तेज़ भी होती है। LLM बहुत स्मार्ट होते हैं (कई रीज़निंग टोकन), लेकिन तुलनात्मक रूप से धीमे होते हैं।
क्या होगा अगर, बेवकूफ/तेज़ मोंटे कार्लो ट्वीकिंग के बजाय, हमारे पास LLM उन दिशाओं का पता लगाएं जिनमें प्रोग्रामों को ले जाना है? हमारे पास एक स्टोकेस्टिक ऑप्टिमाइज़र होगा जो बहुत बुद्धिमान था, हमारे प्रोग्राम को अनुकूलित प्रोग्राम स्पेस के माध्यम से चला रहा था।
अनुकूलन के लिए अंतर्ज्ञान (Intuitions for Optimization)
चलिए एक कदम पीछे हटते हैं। उस व्यक्ति पर विचार करें जिसे आप जानते हैं जो "कोड के स्निपेट को बहुत अधिक ऑप्टिमाइज़ करने" का सबसे अच्छा उदाहरण है। संक्षेप में उन्हें "ऑप्टिमाइज़िन' ओली" कहते हैं। ओली के पास शायद एक सहज अंतर्ज्ञान (gut intuition) है कि किस तरह के कोड ट्वीक फलदायी हो सकते हैं। ओली शायद यह देखने के लिए कुछ चीजों का प्रयास करता है कि क्या वे काम करती हैं, और यदि वे काम नहीं करती हैं, तो उन्हें वापस रोल करता है और कुछ और प्रयास करता है। लेकिन उनके पास कुछ अंतर्ज्ञान है कि किस तरह की चीजें संभव हैं, और वे कंपाइलर को कैसे हरा सकते हैं।
ओली के पास जो अंतर्ज्ञान हैं, वे अक्सर कंपाइलर जो करते हैं उससे परे होते हैं। हालांकि आधुनिक अनुकूलन कंपाइलर अपने परिणामों में काफी प्रभावशाली हैं, वे काफी सरल नियमों और अनुमानों (heuristics) पर आधारित हैं। तकनीकी शब्दजाल में, वे एक स्थानीय डेटाफ्लो ट्रांसफॉर्म के विचार पर आधारित हैं जिसे फिक्स्ड पॉइंट तक चलाया जाता है। हम चरण क्रम (phase order) पर भी विचार करते हैं; यानी हम A और फिर B पर विचार करने के लिए कंपाइलर पाइपलाइन बनाते हैं, लेकिन समग्र AB समस्या पर नहीं। शेड्यूलर और रजिस्टर आवंटनकर्ता इसका एक कुख्यात उदाहरण हैं, समग्र शेड्यूलर-रजिस्टर-आवंटनकर्ता (चरण क्रम को ध्वस्त करने के लाभ प्राप्त करने के लिए) पर कई PhD का प्रयास किया गया है, लेकिन व्यवहार में उन्हें काम करना चुनौतीपूर्ण रहा है।
यही कारण है कि ओली की विशेषज्ञता मूल्यवान है। अक्सर ओली जानता है कि स्थिति के लिए विशिष्ट अनुमानों के साथ कई NP-पूर्ण समस्याओं को कैसे संतुलित किया जाए। तो अधिक विशिष्ट संदर्भ जागरूकता और संवेदनशीलता है। ओली उन तकनीकों को भी नियोजित करने में सक्षम है जिन्हें अनुकूलन कंपाइलर लाभप्रद रूप से लागू नहीं कर सकते हैं, विशेष रूप से संयोजन में, आउटलाइनिंग या कस्टम ABI या वेक्टराइज़ेशन को सक्षम करने के लिए ट्रांसफॉर्म, या अन्य चीजों की भीड़ जो हमें बड़बड़ाने पर मजबूर करती है "काश कंपाइलर के पास ऐसा करने का कोई तरीका होता..."
अब विचार करें कि AI, अपनी जो भी तर्क क्षमताएं हैं, उनके माध्यम से, एक मिनी ओली के रूप में कार्य करने में सक्षम हो सकता है। हो सकता है कि इसके पास इस मामले में मिलान अंतर्ज्ञान न हो कि क्या फलित होगा, लेकिन इसके पास एक संकेत है कि क्या लाभदायक हो सकता है, और यह कई, कई शॉट ले सकता है।
इस दृष्टिकोण के साथ, STOKE पेपर के विपरीत, हम गारंटी नहीं दे सकते कि जैसे-जैसे समय अनंत की ओर बढ़ता है हम इष्टतम प्रोग्राम देख सकते हैं, लेकिन क्योंकि AI में "अधिक मानव जैसी" तर्क क्षमताएं हैं, यह वास्तव में प्रति इकाई समय में महत्वपूर्ण मानव-जैसी गति प्राप्त कर सकता है।
इसे वापस जोड़ना: MLA कर्नेल
मुझे यह कहकर शुरू करने दें: मैं नहीं जानता कि AI ने हमारे जलापेनो MLA कर्नेल के लिए कौन सा निम्न-स्तरीय कोड निकाला, लेकिन मैं जानता हूं कि MLA के लिए numpy कैसे टाइप किया जाता है।
XLA कंपाइलर में जिस पर मैंने पहले काम किया था, हम उन numpy ऑपरेशनों को एक साथ क्लंप में फ्यूज करते थे, और फिर इसे लूप, निर्देशों और निचले-स्तरीय प्रिमिटिव तक कम करने के लिए "एमिटर" नामक एक मेटाप्रोग्राम का उपयोग करते थे।
जब XLA कंपाइलर / एमिटर प्रोग्राम ने ऐसा किया, तो मुझे इस बात की परवाह नहीं थी कि पीछे से कौन सी असेंबली निकली। हमारे स्टोकेस्टिक ऑप्टिमाइज़र के लिए, AI अवधारणात्मक रूप से एमिटर मेटा-प्रोग्राम की जगह लेता है – यह दोनों को कम करता है और अनुकूलित करता है, और हम इसे रूफलाइन की ओर और अधिक अनुकूलित करने के लिए कह सकते हैं।
मुझे उम्मीद है कि यह स्पष्ट करता है कि AI कहाँ फिट बैठता है और यह मौजूदा अनुकूलन कंपाइलर सिस्टम में एक घटक के समान कैसे है। यह सोचना भी सहायक है: हम किस परत को "असेंबली कोड" मानते हैं, वह अब ऊपर जा रही है। जब आप सामान्य C++ टाइप करते हैं और इसे -O3 (उच्चतम विशिष्ट अनुकूलन स्तर) पर कंपाइल करते हैं, तो आप बाहर आने वाली असेंबली को समझने की उम्मीद नहीं करते हैं, भले ही आपने अपने द्वारा टाइप किए गए C++ को समझ लिया हो। हम यहां अनुरूप काम कर रहे हैं, लेकिन एक उच्च-स्तरीय और अधिक गणितीय इनपुट विनिर्देश के साथ।
अब, एक महत्वपूर्ण प्रश्न यह है कि हम कैसे जांचते हैं कि AI से हमें जो प्रोग्राम मिलता है, वह वास्तव में उच्च-स्तरीय विवरण / numpy के बराबर है। वह जाँच तंत्र AI स्टोकेस्टिक अनुकूलन प्रक्रिया की सुदृढ़ता (soundness) स्थापित करता है। मुझे उम्मीद है कि एक भविष्य की ब्लॉग पोस्ट इस पर अधिक विस्तार से जा सकती है, लेकिन अभी के लिए, यह कहना पर्याप्त है कि सिमैंटिक समानता के लिए परीक्षण संभव है और हम इसे करते हैं। एक्सेलेरेटर प्रोग्राम मजबूत, पूर्ण अनुबंधों के लिए विशेष रूप से उपयुक्त होते हैं जिन्हें हम सत्यापित कर सकते हैं "बिल्कुल वही हैं जो AI अनुकूलित प्रोग्राम करता है", क्योंकि वे अपने व्यापक संदर्भ में काफी गणितीय और डेटा प्रवाह उन्मुख होते हैं।
ध्यान दें कि इस क्षेत्र में कई प्रासंगिक तकनीकों का बीड़ा प्रोग्राम संश्लेषण (program synthesis) के उप-क्षेत्र में प्रयासों द्वारा उठाया गया था। जहां अनुकूलन कंपाइलर कहते हैं, "यहाँ सिमैंटिक्स वाला एक प्रोग्राम है, इसे बेहतर बनाएं लेकिन समतुल्य सिमैंटिक्स के साथ!", वहीं प्रोग्राम संश्लेषण कहता है, "मेरा मानना है कि इन सिमैंटिक्स वाला एक प्रोग्राम मौजूद है, कृपया सबसे अच्छा खोजने का प्रयास करें जो आप कर सकते हैं"। प्रोग्राम संश्लेषण अनुकूलन संकलन की तुलना में एक कठिन समस्या है, लेकिन यह मौलिक रूप से कम सीमित भी है। यह प्रभावी रूप से वही है जो ओली जैसे मानव तब करते हैं जब वे अनुकूलन कंपाइलर से बेहतर करते हैं, और यह कुछ ऐसा है जिसे स्वचालित करने में AI अब हमारी मदद कर सकता है। AI मूल प्रोग्राम से "प्रेरणा" ले सकता है, लेकिन उसे उस पर केवल मामूली स्थानीय ट्रांसफॉर्म करने की आवश्यकता नहीं है। क्लासिक अनुकूलन कंपाइलर यह नहीं देखेंगे कि "ओह, आपने बबल सॉर्ट लिखा है" और, अनुबंध को समझकर, इसे क्विक-सॉर्ट में बदल दें, लेकिन ओली और AI दोनों ऐसा करने में सक्षम हैं। यही बात हमें क्लासिकल अनुकूलन-कंपाइलर शासन की तुलना में स्टोकेस्टिक अनुकूलन के साथ प्रोग्राम संश्लेषण शासन में अधिक रखती है।
यह सब इस तथ्य में एक साथ आता है कि आप किसी ऐसी चीज़ से शुरू कर सकते हैं जो "numpy से बहुत दूर नहीं है", 48 घंटे प्रतीक्षा करें, और समान सिमैंटिक्स वाला एक अनुकूलित कर्नेल प्राप्त करें, जैसा कि हमने अपने HotChips टॉक में दिखाया:
जैसा कि स्लाइड भी नोट करती है, हमारी मशीन पर हम अक्सर AI को अपने मानव विशेषज्ञों से बेहतर प्रदर्शन करते हुए देखने में सक्षम होते हैं, यहां तक कि उन कर्नेल पर भी जिन्हें हम काफी अच्छी तरह से ट्यून किया हुआ मानते थे। अक्सर एक सभ्य प्राप्त करने योग्य प्रतिशत अभी भी बचा होता है, केवल उन संयोजनों / क्रमपरिवर्तनों की कई किस्मों के कारण जिन्हें खोजे जाने की आवश्यकता हो सकती है। ये अक्सर एक मानव प्रदर्शन इंजीनियर के लिए अत्यधिक कठिन होते हैं।
पुनर्कथन और निष्कर्ष
एक कंपाइलर, दिन के अंत में, सिर्फ एक फलन है। हम अपना प्रोग्राम उस फलन को देते हैं, और हमें अपने प्रोग्राम का एक बेहतर संस्करण वापस मिलता है। हम जो प्रोग्राम बाहर निकालते हैं और जो प्रोग्राम हम अंदर डालते हैं, हम उम्मीद करते हैं कि उनमें समान सिमैंटिक्स हों।
पारंपरिक अनुकूलन कंपाइलर डेटाफ्लो नियमों और अनुमानों के माध्यम से प्रोग्राम को बेहतर बनाते हैं। ये अपनी उत्पत्ति में पूरी तरह से समझने योग्य हैं, लेकिन वे जो कदम उठा सकते हैं, उनमें अधिक सीमित भी हो सकते हैं।
इसके विपरीत, AI, एक स्टोकेस्टिक ऑप्टिमाइज़र के रूप में, बस "कड़ी मेहनत से सोचना" है और कुछ निकालना है। इसके कदम मौलिक रूप से सीमित नहीं हैं, जो उन्हें हमारे मानव विशेषज्ञ ऑप्टिमाइज़र के समान बनाता है। हमें यह जांचने के तरीकों की आवश्यकता है कि यह जो प्रोग्राम निकालता है वे सुदृढ़ हैं और उन्हीं सिमैंटिक्स को लागू करते हैं जो हमने अंदर डाले हैं, और हमारे पास वे मौजूद हैं। और इस प्रकार का AI अनुकूलन विशेष रूप से गणितीय संक्रियाओं के लिए उपयुक्त है जिनमें बहुत मजबूत अनुबंध होते हैं। अनुबंध कर्नेल को लाइन दर लाइन समझने की आवश्यकता से बचाते हैं।
इस तरह हमें AI द्वारा उत्पन्न MLA कर्नेल मिला!
