इनफ्रेंस ऑप्टिमाइजेशन - लेटेंसी के लिए मॉडल को कैसे ऑप्टिमाइज करें?

गति अच्छी है, लेकिन निरंतर प्रवाह (फ्लो) उससे भी बेहतर है। जानें कि torch.compile() से लेकर CUDA Graphs और डायनेमिक शेप्स तक—PyTorch के साथ अपने AI मॉडल इन्फरेंस को कैसे अनुकूलित (ऑप्टिमाइज़) करें।
कल्पना कीजिए कि अगर आपको घर से दफ्तर जाने के लिए हर दिन अपना मैप देखना पड़े। हर मोड़, हर लेन, हर सिग्नल फैसलों का एक निरंतर सिलसिला है, जिसमें से हर एक आपको थोड़ा धीमा कर देता है। अब कल्पना कीजिए कि आपने पहले ही दिन रास्ता सीख लिया। कोई मैप नहीं, कोई दोबारा सोचना नहीं, बस सहज, आसान गति। आप न केवल तेज़ हैं, बल्कि अधिक स्वतंत्र हैं। आपका मन हल्का है, जो यात्रा पर केंद्रित है, न कि निर्देशों पर।
यही है लेटेंसी (latency) के लिए ऑप्टिमाइज़ करना। यह केवल गति के बारे में नहीं है। यह हर अनावश्यक जांच, हर बेकार ठहराव, हर झिझक को दूर करने के बारे में है। यह फ्लो (प्रवाह) बनाने के बारे में है। जब आपके सिस्टम को ठीक-ठीक पता होता है कि उसे कहाँ जाना है, तो वह उद्देश्य और स्पष्टता के साथ आगे बढ़ता है। ऑप्टिमाइज़ेशन का मतलब काम को गायब कर देना है।
चलिए थोड़ा प्रैक्टिकल करते हैं!
आइए एक पूरी तरह से काम करने वाले ट्रांसफॉर्मर मॉडल के साथ शुरुआत करें जो हमारे बेसलाइन के रूप में काम करेगा और काम को पूरा करेगा। इसे एक भरोसेमंद वाहन की तरह समझें जो आपको पॉइंट A से पॉइंट B तक सफलतापूर्वक पहुँचाता है। इसमें कोई खराबी नहीं है, यह बिल्कुल डिज़ाइन के अनुसार काम करता है। लेकिन किसी भी अच्छे इंजीनियर की तरह, हम इसे और भी बेहतर बनाने के अवसर देख सकते हैं।
यहाँ transformers का उपयोग करके BERT इन्फ्रेंस (inference) का एक न्यूनतम कार्यान्वयन दिया गया है।
torch.compile() - रास्ता सीखना
अब आइए torch.compile() के साथ अपने बेसलाइन को बेहतर बनाएं - PyTorch का इन-बिल्ट ऑप्टिमाइज़ेशन जो हमारे डायनेमिक निष्पादन को एक अधिक कुशल स्टैटिक ग्राफ में बदल देता है। इसे उस क्षण की तरह समझें जब आप अपने रास्ते से इतने परिचित हो जाते हैं कि आप बिना बार-बार दिशा-निर्देशों की मदद लिए आसानी से गाड़ी चला सकते हैं।
यह एक बार हमारे मॉडल के निष्पादन को ट्रैक करता है, हर ऑपरेशन और डिसीजन पॉइंट को रिकॉर्ड करता है। फिर यह एक अनुकूलित कंप्यूटेशन पाथ (computation path) बनाता है। डायनेमिक पायथन निष्पादन कंपाइल किए गए C कोड के समान स्टैटिक ऑपरेशन्स में बदल जाता है। PyTorch आपके मॉडल के निष्पादन को एक बार ट्रैक करता है, जिसमें हर ऑपरेशन, हर टेंसर आकार, हर डिसीजन पॉइंट रिकॉर्ड होता है। फिर यह इस ट्रैक को ऑप्टिमाइज़ करता है, अनावश्यक जांचों को हटाता है और एक सुव्यवस्थित निष्पादन पाथ बनाता है जैसा कि चित्र 1 में दिखाया गया है। अनगिनत if-else स्टेटमेंट जिन्हें पायथन आम तौर पर इवैल्यूएट करता है, टेंसर आकारों की जांच करता है, यह निर्धारित करता है कि किन CUDA कर्नेल्स का उपयोग करना है, और मेमोरी लेआउट तय करता है - वे सभी ऑपरेशन्स के एक पूर्वनिर्धारित क्रम में सिमट जाते हैं।

CUDA Graphs - आपके कंप्यूटेशन्स के लिए हाईवे सिस्टम
torch.compile() के पास एक और भी शक्तिशाली तरीका है: CUDA Graphs। यदि नियमित CUDA कर्नेल लॉन्च शहर में ड्राइविंग की तरह हैं जहाँ आप लाइटों पर रुकते हैं, ट्रैफ़िक को रास्ता देते हैं, अप्रत्याशित परिस्थितियों से निपटते हैं, तो CUDA Graphs आपकी खुद की प्राइवेट हाईवे की तरह हैं जहाँ सभी हरी लाइटें हैं।
CUDA Graphs व्यक्तिगत कर्नेल लॉन्च करने के ओवरहेड को समाप्त करते हैं। आम तौर पर, प्रत्येक ऑपरेशन के लिए CPU और GPU के बीच एक राउंड ट्रिप की आवश्यकता होती है - CPU, GPU को निर्देश देता है कि क्या करना है, पावती (acknowledgment) की प्रतीक्षा करता है, और फिर अगला निर्देश भेजता है। CUDA Graphs के साथ, GPU ऑपरेशन्स के पूरे क्रम को एक बार रिकॉर्ड किया जाता है और फिर एक एकल इकाई (single unit) के रूप में फिर से चलाया जाता है जैसा कि चित्र 2 में दिखाया गया है।

max-autotune - क्या यह केवल CUDA ग्राफ ही जनरेट करता है? 🤔
max-autotune मोड केवल आपके मौजूदा ऑपरेशन्स को ऑप्टिमाइज़ नहीं करता, यह पूरी तरह से नए ऑपरेशन्स जनरेट करता है। इसे एक विशेषज्ञ मैकेनिक की तरह समझें जो केवल आपके वर्तमान इंजन को ट्यून नहीं करता बल्कि आपके लिए विशेष रूप से आपके रूट के हिसाब से डिज़ाइन किया गया पूरी तरह से कस्टम इंजन बनाता है। यह विशेष रूप से आपके मॉडल के ऑपरेशन्स के लिए कस्टम Triton कर्नेल्स जनरेट कर रहा है।

यह प्रक्रिया इस तरह काम करती है -> PyTorch प्रत्येक ऑपरेशन के लिए कई Triton कर्नेल कार्यान्वयन जनरेट करता है, उन सभी का बेंचमार्क लेता है, और सबसे तेज़ को चुनता है। यह ऐसा है जैसे विशेषज्ञ ड्राइवरों की एक टीम आपके गंतव्य के लिए विभिन्न मार्गों का परीक्षण करती है और फिर आपको उनके द्वारा खोजा गया इष्टतम मार्ग सिखाती है। ये कस्टम कर्नेल्स अक्सर हाथ से लिखे गए CUDA कार्यान्वयनों से बेहतर प्रदर्शन करते हैं क्योंकि वे आपके विशिष्ट उपयोग के मामले की पूरी जानकारी के साथ जनरेट किए जाते हैं। उन्हें ठीक-ठीक पता होता है कि किस टेंसर आकार की उम्मीद करनी है, कौन से मेमोरी एक्सेस पैटर्न होंगे, और GPU मेमोरी पदानुक्रमों के बीच डेटा मूवमेंट को कैसे कम किया जाए।
लेकिन रुकिए, क्या होगा अगर मेरा रूट रोज़ बदल जाए?
क्या होता है जब आपके मॉडल को अलग-अलग इनपुट आकारों को संभालने की आवश्यकता होती है? कल्पना कीजिए कि यदि मौसम, ट्रैफ़िक या आपके मूड के आधार पर आपका दैनिक मार्ग बदल जाए। यह डायनेमिक शेप की समस्या है, जब आपके ट्रांसफॉर्मर को अलग-अलग लंबाई के वाक्यों को प्रोसेस करने की आवश्यकता होती है, या आपका कंप्यूटर विज़न मॉडल विभिन्न रिज़ॉल्यूशन की छवियों को संभालता है। चुनौती यह है कि CUDA Graphs मूल रूप से स्टैटिक होते हैं - वे निश्चित टेंसर आकारों के साथ ऑपरेशन्स के एक विशिष्ट क्रम को रिकॉर्ड करते हैं। यदि आपके इनपुट आकार बदलते हैं, तो ग्राफ अमान्य हो जाता है, जैसे कि आपने ऐसे रास्ते याद किए हों जो केवल तभी काम करते हैं जब ट्रैफ़िक एक विशिष्ट पैटर्न में बहता है।
यहीं पर max-autotune-no-cudagraphs का महत्व सामने आता है। यह एक अत्यधिक कुशल ड्राइवर की तरह है जो कई इष्टतम मार्गों को जानता है और अनुभव तथा विशेषज्ञता के लाभों को खोए बिना वास्तविक समय में खुद को ढाल सकता है।
max-autotune-no-cudagraphs मोड आपको दोनों दुनिया के सर्वश्रेष्ठ लाभ देता है। आपको अभी भी अपने विशिष्ट ऑपरेशन्स के लिए जनरेट किए गए कस्टम Triton कर्नेल्स, ग्राफ ऑप्टिमाइज़ेशन जो अनावश्यक गणनाओं को हटाते हैं, और फ़्यूज़्ड ऑपरेशन्स जो मेमोरी बैंडविड्थ को कम करते हैं, मिलते हैं। जो आप खोते हैं वह है CUDA Graph रीप्ले मैकेनिज्म, लेकिन जो आप पाते हैं वह है बिना री-कंपाइलेशन के विभिन्न इनपुट आकारों को संभालने की क्षमता।
यह दृष्टिकोण विशेष रूप से उन प्रोडक्शन सिस्टम्स के लिए शक्तिशाली है जहाँ आप पहले से इनपुट आकारों की भविष्यवाणी नहीं कर सकते। यह एक विशेषज्ञ टैक्सी ड्राइवर की तरह है जिसे किसी एक विशिष्ट मार्ग को याद रखने की आवश्यकता नहीं है, बल्कि उसने कुशल नेविगेशन के सिद्धांतों को आत्मसात कर लिया है और उन्हें किसी भी गंतव्य पर लागू कर सकता है।
ऑपरेशन दक्षता में सुधार करके मॉडल ऑप्टिमाइज़ेशन के बारे में बहुत सारी बातें हो गईं। और क्या है?
लंबी ड्राइव का लाभ -> कैसे टेंसर को प्री-एलोकेट करना आपके फुल टैंक की तरह है
कल्पना कीजिए कि आप अपने टैंक को फुल करके लंबी यात्रा पर निकले हैं, आपको रीफिल के लिए रुकना नहीं पड़ेगा और आप एक बार में यात्रा पूरी कर सकते हैं। इसी तरह, टेंसर को प्री-एलोकेट (पहले से आवंटित) करने से आपके मॉडल को मेमोरी एलोकेशन के बार-बार रुकने वाले झटकों के बिना इन्फ्रेंस में सुचारू रूप से आगे बढ़ने में मदद मिलती है। हर बार जब आपका मॉडल इन्फ्रेंस के बीच में सिस्टम से मेमोरी मांगता है, तो यह गैस स्टेशन खोजने के लिए गाड़ी को किनारे लगाने जैसा होता है, यात्रा रुक जाती है, कीमती समय बीत जाता है, और सहज लय टूट जाती है।

हर बार जब PyTorch इन्फ्रेंस के दौरान एक नया टेंसर बनाता है, तो यह आपके हाईवे पर एक अप्रत्याशित टोल बूथ का सामना करने जैसा होता है। अनुरोध तुरंत लगता है, लेकिन सतह के नीचे एक जटिल प्रक्रिया चलती है: उपलब्ध मेमोरी खोजना, अलाइनमेंट आवश्यकताओं की जांच करना, एलोकेशन तालिकाओं को अपडेट करना, और कभी-कभी मौजूदा मेमोरी ब्लॉक्स को डीफ़्रेग्मेंट भी करना।
पारंपरिक टेंसर एलोकेशन बड़े बैच आकारों पर औसतन 21.3% के फ्रेगमेंटेशन से ग्रस्त होता है। मेमोरी ऑप्टिमाइज़ेशन तकनीकें न्यूरल नेटवर्क आर्किटेक्चर के आधार पर समग्र मेमोरी उपयोग को 20-70% तक कम करते हुए इस फ्रेगमेंटेशन को पूरी तरह से समाप्त कर देती हैं। लेकिन वास्तविक बदलाव तब होता है जब आप प्री-एलोकेशन को इंटेलिजेंट मेमोरी पूलिंग रणनीतियों के साथ जोड़ते हैं।
आइए इसे एक व्यावहारिक उदाहरण से समझते हैं। यहाँ एक उदाहरण दिया गया है जो इनके बीच का अंतर प्रदर्शित करता है
एक नया आवंटित टेंसर और एक प्री-एलोकेटेड टेंसर।
>>> Time without pre-allocation: 0.1738 seconds
>>> Time with pre-allocation: 0.0353 seconds
प्रत्येक मेमोरी एलोकेशन के लिए सिस्टम कॉल, अनुमति जांच और प्रशासनिक रिकॉर्ड-कीपिंग की आवश्यकता होती है। प्री-एलोकेशन इसे प्रत्येक इन्फ्रेंस में हजारों बार करने के बजाय शुरू में एक बार में ही संभाल लेता है, जैसा कि चित्र में दिखाया गया है।
आधुनिक GPUs और CPUs में परिष्कृत भविष्यवाणी तंत्र होते हैं जो सुसंगत पैटर्न के साथ सबसे अच्छा काम करते हैं। यह क्रूज़ कंट्रोल की तरह है, एक बार जब सिस्टम जान जाता है कि आगे क्या आने वाला है, तो वह उसी के अनुसार ऑप्टिमाइज़ कर सकता है।

अब तक सब ठीक है, क्या हम कुछ और सुधार कर सकते हैं?
हाँ! आपका GPU कई ऑपरेशन्स को एक सहज गति में जोड़ना सीख सकता है। यह kernel fusion (कर्नेल फ्यूजन) का जादू है। यह आपके सिस्टम को मेमोरी की अनावश्यक यात्राएं बंद करने और इसके बजाय सब कुछ सुचारू रूप से चलाने की कला है। आपके GPU में दो प्रकार की मेमोरी होती है: कंप्यूट कोर के ठीक बगल में स्थित बिजली जैसी तेज़ रजिस्टर्स, और धीमी ग्लोबल मेमोरी जिसके लिए लंबी यात्रा की आवश्यकता होती है। पारंपरिक ऑपरेशन्स प्रत्येक सामग्री के लिए स्टोर की अलग-अलग यात्राएं करने जैसे हैं। फ्यूज्ड ऑपरेशन्स? यह आपकी स्मार्ट शॉपिंग लिस्ट की तरह है जो एक ही सहज यात्रा में सब कुछ ले आती है।
आइए फ़्यूज़्ड कर्नेल्स का उपयोग करके सरल जोड़ और गुणा ऑपरेशन्स को ऑप्टिमाइज़ करने का प्रयास करें।
तीन अलग-अलग GPU कर्नेल लॉन्च करने के बजाय, जिनमें से प्रत्येक का अपना मेमोरी ओवरहेड होता है, addcmul यह सब एक ही ऑपरेशन में करता है। हमारे इंटरमीडिएट परिणाम कभी भी धीमी ग्लोबल मेमोरी को नहीं छूते हैं। वे सीधे GPU के तेज़ रजिस्टर्स से होकर बहते हैं जैसे किसी पूरी तरह से इंजीनियर की गई पाइपलाइन से पानी बहता है, जैसा कि चित्र 6 में दिखाया गया है।

हमारी यात्रा का अंत
शहर की सड़कों पर घूमने से लेकर कस्टम-निर्मित हाईवे तक, ईंधन के लिए पिट स्टॉप से लेकर निर्बाध, बिना रुकावट के क्रूज़ तक, इन्फ्रेंस ऑप्टिमाइज़ेशन की यह यात्रा केवल गति से कहीं अधिक रही है। यह क्षेत्र को समझने, इंजन को ट्यून करने और अंततः, सिस्टम को उद्देश्य और सटीकता के साथ प्रवाहित करने के बारे में रहा है।
हमने एक विश्वसनीय मॉडल के साथ शुरुआत की और इसे प्रदर्शन के मामले में एक बेहतरीन मशीन में बदल दिया, रूट को ट्रैक करने के लिए torch.compile(), ट्रैफिक लाइट हटाने के लिए CUDA Graphs, सही इंजन तैयार करने के लिए max-autotune, और ईंधन की हर बूंद का उपयोग करने के लिए मेमोरी ऑप्टिमाइज़ेशन का उपयोग किया।
लेकिन याद रखें, ऑप्टिमाइज़ेशन कभी वास्तव में खत्म नहीं होता है। किसी भी सड़क की तरह, नई परिस्थितियाँ उत्पन्न होती हैं, आकार बदलते हैं, मॉडल विकसित होते हैं, वर्कलोड शिफ्ट होते हैं। मुख्य बात केवल एक रास्ता जानने में नहीं है, बल्कि तेजी से और बुद्धिमानी से अनुकूलन करने के कौशल में महारत हासिल करने में है।
तो जैसे ही हम इस ड्राइव को समाप्त करते हैं, अपने साथ न केवल उपकरण ले जाएं, बल्कि स्पष्टता, दक्षता और ऐसे सिस्टम बनाने की मानसिकता भी ले जाएं जो केवल काम न करें, बल्कि प्रवाहित हों।
हैप्पी ऑप्टिमाइज़िंग! 🙂


