Axon: A Synthesizing Superoptimizer for Tensor Programs
Axon एक सिंथेसाइजिंग सुपरऑप्टिमाइज़र है जो प्रोग्राम सिंथेसिस और SMT-आधारित सिमेंटिक वेरिफिकेशन का उपयोग करके AI एक्सेलेरेटर्स के लिए उच्च-प्रदर्शन वाले टेंसर प्रोग्राम स्वचालित रूप से उत्पन्न करता है ताकि हैंड-क्राफ्टेड रीराइट रूल्स पर निर्भर रहे बिना बीजगणितीय रूपांतरणों (algebraic transformations) को खोजने, टाइलिंग और इंस्ट्रक्शन सिलेक्शन को अनुकूलित करने और ऑपरेटर्स को फ्यूज करने के लिए किया जा सके।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक बेहतरीन केक बनाने की कोशिश कर रहे हैं, लेकिन आप एक बहुत ही हाई-टेक किचन में काम कर रहे हैं जहाँ ओवन और मिक्सर बहुत ही विशिष्ट और नखरे वाले हैं। आप जानते हैं कि आप क्या बनाना चाहते हैं (रेसिपी), लेकिन मशीन को इसे जितनी जल्दी हो सके करने के लिए आपको यह जानना होगा कि सामग्री को कैसे लोड करना है, कब चलाना है, और ओवन के कौन से विशिष्ट बटन दबाने हैं।
लंबे समय तक, AI चिप्स (ओवन) के लिए कोड लिखना (बनाने का तरीका) हाथ से केक बनाने जैसा था, एक बार में एक छोटा सा टुकड़ा करके। आपको एक मास्टर शेफ होना पड़ता था जो ओवन की सटीक कार्यप्रणाली को समझता हो ताकि अच्छा परिणाम मिल सके। यदि आप सामग्री लोड करने के तरीके में एक छोटी सी गलती भी करते हैं, तो केक धीमा और सूखा बनकर निकलता है।
Axon एक नया टूल है जो एक स्वचालित, सुपर-स्मार्ट sous-chef (सहायक शेफ) की तरह काम करता है जो न केवल आपकी रेसिपी का पालन करता है—बल्कि उसे बेहतर तरीके से बनाने के नए तरीके भी खोज निकालता है।
यहाँ बताया गया है कि Axon कैसे काम करता है, जिसे सरल चरणों में विभाजित किया गया है:
1. "क्या" बनाम "कैसे" (The "What" vs. The "How")
आमतौर पर, एक प्रोग्रामर एक रेसिपी लिखता है जिसमें कहा जाता है, "मैदा और अंडे मिलाएं, फिर बेक करें।" वे इस बात की चिंता नहीं करते कि ओवन उन्हें कैसे मिलाता है।
- समस्या: ओवन (AI चिप) बहुत नखरेबाज है। इसके अलग-अलग हिस्से (compartments) होते हैं (जैसे एक बड़ा मिक्सिंग बाउल, एक छोटा व्हिस्क, और एक हीटिंग एलिमेंट)। यदि आप इसे बड़े बाउल में सब कुछ मिलाने के लिए कहते हैं जबकि यह व्हिस्क का उपयोग कर सकता था, तो यह समय बर्बाद करता है।
- Axon का काम: Axon आपकी साधारण रेसिपी लेता है और कहता है, "रुको, मैं इसे और तेजी से कर सकता हूँ।" यह केवल रेसिपी का पालन नहीं करता; यह निर्देशों को ओवन की विशिष्ट खूबियों के अनुरूप फिर से लिखता है।
2. "जादुई पुनर्व्यवस्था" (एल्जेब्रिक ट्रांसफॉर्मेशन - Algebraic Transformations)
कल्पना कीजिए कि आपकी रेसिपी कहती है: "पहले, सभी गाजरें काटें, फिर उन्हें सूप में मिलाएं, फिर नमक डालें।"
- पुराना तरीका: आप काटते हैं, फिर मिलाते हैं, फिर नमक डालते हैं। आपको कटिंग बोर्ड और स्टोव के बीच बार-बार कटोरा लेकर आना-जाना पड़ता है।
- Axon का तरीका: Axon गणित को देखता है और महसूस करता है, "हे, मैं नमक तभी डाल सकता हूँ जब गाजरें काटी जा रही हों, या मैं गाजरों को काटने से पहले ही सूप में मिला सकता हूँ, और इससे स्वाद बिल्कुल वैसा ही रहेगा!"
- परिणाम: चरणों के क्रम को पुनर्व्यवस्थित करके, Axon ओवन को एक के बाद एक करने के बजाय एक ही समय में दो चीजें करने की अनुमति देता है (जैसे काटने और गर्म करने का काम एक साथ करना)। यह गणितीय रूप से सिद्ध करता है कि केक का स्वाद अभी भी वैसा ही रहेगा, भले ही चरण अलग हों।
3. "पैकिंग पहेली" (टाइलिंग - Tiling)
कल्पना कीजिए कि आपके पास सामग्री का एक बड़ा ढेर है, लेकिन आपका मिक्सिंग बाउल छोटा है। आपको सामग्री को पेंट्री (धीमी) से बाउल (तेज़) में ले जाना पड़ता है, थोड़ा मिलाना पड़ता है, वापस ले जाना पड़ता है, और यही प्रक्रिया दोहरानी पड़ती है।
- समस्या: यदि आप एक बार में बहुत अधिक सामग्री ले जाते हैं, तो बाउल भर जाएगा। यदि आप बहुत कम ले जाते हैं, तो आप अपना सारा समय पेंट्री से आने-जाने में बिता देंगे।
- Axon का तरीका: Axon एक मास्टर पैकर है। यह सामग्री को "बैच" (टाइल्स) में समूहबद्ध करने के लाखों अलग-अलग तरीकों को आजमाता है। यह पूछता है, "क्या होगा अगर मैं एक बार में 100 गाजरें ले जाऊं? क्या होगा अगर मैं 50 ले जाऊं?" यह इन विभिन्न ग्रुपिंग को टेस्ट करता है ताकि वह एकदम सही आकार पा सके जो बाउल को भरने और ओवरफ्लो होने के बीच का संतुलन बनाए रखे, जिससे पेंट्री के चक्कर कम से कम लगें।
4. "निर्देश पुस्तिका" (ISA सिंथेसिस - ISA Synthesis)
ओवन के पास एक विशिष्ट कंट्रोल पैनल है जिसमें "हाई हीट," "व्हिस्क," और "फोल्ड" जैसे बटन हैं।
- समस्या: कभी-कभी रेसिपी कहती है "मिक्स करें," लेकिन ओवन में "मिक्स" बटन नहीं होता। इसमें केवल "व्हिस्क" और "फोल्ड" बटन होते हैं।
- Axon का तरीका: Axon "मिक्स" को ठीक उसी क्रम में अनुवादित करता है जिसे ओवन समझ सके, जैसे "व्हिस्क फिर फोल्ड"। यह शॉर्टकट भी ढूंढ लेता है, जैसे यह पहचानना कि इस विशिष्ट ओवन पर "स्टोव तक ले जाते समय फोल्ड करना" एक सिंगल बटन प्रेस है, जिससे समय बचता है।
5. "स्वाद परीक्षण" (एम्पिरिकल सिलेक्शन - Empirical Selection)
यह सबसे महत्वपूर्ण हिस्सा है। Axon केवल अंदाज़ा नहीं लगाता कि कौन सा तरीका सबसे अच्छा है।
- प्रक्रिया: यह रेसिपी के हजारों अलग-अलग संस्करण बनाता है (कुछ अलग क्रम के साथ, कुछ अलग बैच साइज के साथ, कुछ अलग बटन प्रेस के साथ)।
- परीक्षण: यह वास्तव में वास्तविक ओवन (AI चिप) पर इन सभी संस्करणों को चलाता है और उनका समय नोट करता है। यह उस संस्करण को चुनता है जो सबसे तेज़ काम करता है। यह 1,000 छोटे केक बनाने जैसा है यह देखने के लिए कि कौन सा केक 10 मिनट के बजाय 5 मिनट में बनकर तैयार होता है।
परिणाम: कितनी तेज़?
पेपर में Amazon के Trainium चिप्स (एक प्रकार का AI ओवन) पर 20 अलग-अलग प्रकार के "केक्स" (AI टास्क) के साथ Axon का परीक्षण किया गया।
- मानक ओवन के मुकाबले: स्टैंडर्ड तरीके से बेक करने (Amazon के Neuron कंपाइलर) की तुलना में, Axon ने सिंगल टास्क को 3.7 गुना तक तेज़ और जटिल टास्क को 19 गुना तक तेज़ बना दिया।
- विशेषज्ञ शेफ के मुकाबले: यहाँ तक कि उन मानव विशेषज्ञों द्वारा लिखी गई रेसिपी की तुलना में जो ओवन को अंदर से जानते हैं, Axon कभी-कभी 1.35 गुना तेज़ था।
- अन्य रोबोट के मुकाबले: इसने दूसरे ऑटोमेटेड टूल "Mirage" को औसतन लगभग 10% से मात दी, और यह 16 प्रकार के केक भी बना सका जिन्हें वह दूसरा टूल बिल्कुल भी हैंडल नहीं कर सका।
सारांश
Axon एक ऐसा टूल है जो यह लेता है कि एक AI को क्या करने की आवश्यकता है और स्वचालित रूप से यह पता लगाता है कि एक विशिष्ट कंप्यूटर चिप पर इसे करने का सबसे तेज़ तरीका क्या है। यह यह सब करके करता है कि वह गणितीय रूप से सिद्ध करता है कि वह चरणों को पुनर्व्यवस्थित कर सकता है, डेटा को पैक करने के लाखों तरीकों का परीक्षण करता है, और फिर विजेता चुनने के लिए कोड को वास्तव में चलाता है। यह मनुष्यों को हार्डवेयर के सूक्ष्म, उबाऊ विवरणों का विशेषज्ञ बनने की आवश्यकता को समाप्त कर देता है, जिससे वे बस रेसिपी लिख सकते हैं जबकि Axon बेकिंग का काम संभालता है।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।