ग्रेडिएंट चेकपॉइंटिंग एक तकनीक है जिसका उपयोग Deep learning में तंत्रिका नेटवर्क के प्रशिक्षण की मेमोरी खपत को कम करने के लिए किया जाता है। मानक बैकप्रोपेगेशन के दौरान, एक नेटवर्क को बैकवर्ड पास में ग्रेडिएंट की गणना करने के लिए फॉरवर्ड पास में गणना किए गए सभी मध्यवर्ती सक्रियणों को संग्रहीत करना होता है। बहुत गहरे मॉडलों के लिए, जैसे कि बड़े भाषा मॉडल और ट्रांसफॉर्मर, यह भंडारण उपलब्ध हार्डवेयर की मेमोरी क्षमता से अधिक हो सकता है। ग्रेडिएंट चेकपॉइंटिंग हर सक्रियण को संग्रहीत न करके इसका समाधान करती है; इसके बजाय, यह केवल एक उपसमुच्चय रखती है और बैकवर्ड पास के दौरान मांग पर छोड़े गए सक्रियणों की पुनर्गणना करती है। यह बढ़ी हुई कम्प्यूटेशनल लागत को काफी कम मेमोरी उपयोग के साथ बदलता है, जिससे बड़े मॉडलों का प्रशिक्षण या समान हार्डवेयर पर बड़े बैच आकारों का उपयोग संभव होता है।
यह तकनीक 2016 में कार्नेगी मेलन विश्वविद्यालय और ओपनएआई के शोधकर्ताओं द्वारा "Training Deep Nets with Sublinear Memory Cost" शीर्षक वाले एक पेपर में पेश की गई थी। लेखकों, जिनमें Tianqi Chen, Bing Xu, Chiyuan Zhang, और Carlos Guestrin शामिल हैं, ने प्रदर्शित किया कि केवल कुछ चेकपॉइंट्स पर सक्रियणों को संग्रहीत करके (उदाहरण के लिए, हर कुछ परतों पर) और बाकी की पुनर्गणना करके, n परतों वाले एक गहरे नेटवर्क के प्रशिक्षण की मेमोरी लागत को O(n) से O(sqrt(n)) तक कम किया जा सकता है, जिसकी लागत लगभग एक अतिरिक्त फॉरवर्ड पास होती है। यह मौलिक कार्य तब से मशीन लर्निंग समुदाय में एक मानक उपकरण बन गया है, विशेष रूप से जब मॉडल आकार नाटकीय रूप से बढ़े हैं।
मानक बैकप्रोपेगेशन मेमोरी का उपयोग कैसे करता है
एक पारंपरिक प्रशिक्षण लूप में, फॉरवर्ड पास नेटवर्क की हर परत के लिए सक्रियणों की गणना करता है। ये सक्रियण मेमोरी में संग्रहीत होते हैं क्योंकि बैकवर्ड पास को चेन नियम के माध्यम से ग्रेडिएंट की गणना करने के लिए उनकी आवश्यकता होती है। L परतों वाले एक नेटवर्क के लिए, इसके लिए L सेट सक्रियणों को संग्रहीत करने की आवश्यकता होती है, जिनमें से प्रत्येक बड़ा हो सकता है। उदाहरण के लिए, सैकड़ों परतों वाला एक अवशिष्ट नेटवर्क या दर्जनों अटेंशन ब्लॉक वाला एक ट्रांसफॉर्मर एक एकल प्रशिक्षण उदाहरण के लिए गीगाबाइट्स सक्रियण डेटा जमा कर सकता है। बड़े बैच आकारों के साथ प्रशिक्षण करते समय, मेमोरी आवश्यकता बैच आकार के साथ रैखिक रूप से बढ़ती है, जो अक्सर इसे प्राथमिक बाधा बनाती है।
चेकपॉइंटिंग रणनीति
ग्रेडिएंट चेकपॉइंटिंग नेटवर्क को खंडों में विभाजित करती है, प्रत्येक खंड की सीमा पर एक चेकपॉइंट होता है। फॉरवर्ड पास के दौरान, केवल इन चेकपॉइंट सीमाओं पर सक्रियणों को मेमोरी में सहेजा जाता है। एक खंड के भीतर अन्य सभी मध्यवर्ती सक्रियणों को छोड़ दिया जाता है। जब बैकवर्ड पास एक खंड तक पहुंचता है, तो यह सहेजे गए चेकपॉइंट सक्रियण का उपयोग करके उस खंड के लिए फॉरवर्ड पास की पुनर्गणना करता है, जिससे ग्रेडिएंट गणना के लिए आवश्यक मध्यवर्ती सक्रियणों को पुनः उत्पन्न किया जाता है। यह पुनर्गणना कम्प्यूटेशनल ओवरहेड जोड़ती है, जो आमतौर पर प्रति प्रशिक्षण चरण एक अतिरिक्त फॉरवर्ड पास के बराबर होती है, लेकिन यह पीक मेमोरी उपयोग को नाटकीय रूप से कम करती है।
चेकपॉइंट प्लेसमेंट का चुनाव एक व्यापार-बंद है। अधिक चेकपॉइंट्स का मतलब कम पुनर्गणना लेकिन उच्च मेमोरी उपयोग है; कम चेकपॉइंट्स का मतलब कम मेमोरी लेकिन अधिक गणना है। n परतों वाले एक नेटवर्क के लिए चेकपॉइंट्स की इष्टतम संख्या लगभग sqrt(n) है, जो मेमोरी और कंप्यूट को संतुलित करती है। व्यवहार में, PyTorch और TensorFlow जैसे फ्रेमवर्क उपयोगकर्ताओं को चेकपॉइंट अंतराल निर्दिष्ट करने या स्वचालित ह्यूरिस्टिक्स का उपयोग करने की अनुमति देते हैं।
विविधताएं और सुधार
मूल तकनीक के कई शोधन विकसित किए गए हैं। एक सामान्य भिन्नता चयनात्मक चेकपॉइंटिंग है, जहां केवल कुछ परत प्रकारों (जैसे अटेंशन ब्लॉक या कन्वोल्यूशनल परतें) को चेकपॉइंट किया जाता है, जबकि अन्य को सामान्य रूप से संग्रहीत किया जाता है। एक अन्य दृष्टिकोण, जिसे मेमोरी-कुशल ग्रेडिएंट चेकपॉइंटिंग कहा जाता है, एक अधिक परिष्कृत शेड्यूल का उपयोग करता है जो कई स्तरों की ग्रैन्युलैरिटी पर सक्रियणों को संग्रहीत करता है, जिससे अतिरिक्त पुनर्गणना की लागत पर मेमोरी और कम हो जाती है। कुछ फ्रेमवर्क "ऑफलोडिंग" भी लागू करते हैं, जहां चेकपॉइंट्स को CPU मेमोरी या डिस्क में स्थानांतरित किया जाता है, हालांकि यह डेटा ट्रांसफर ओवरहेड पेश करता है।
ट्रांसफॉर्मर मॉडलों के संदर्भ में, ग्रेडिएंट चेकपॉइंटिंग को अक्सर मिश्रित-परिशुद्धता प्रशिक्षण और ग्रेडिएंट क्लिपिंग जैसी अन्य मेमोरी-बचत तकनीकों के साथ जोड़ा जाता है। उदाहरण के लिए, GPT-3 जैसे मॉडल का प्रशिक्षण, जिसमें 175 बिलियन पैरामीटर हैं, ऐसे अनुकूलन के बिना असंभव होगा। यह तकनीक बड़े मॉडलों को फाइन-ट्यून करने में भी उपयोग की जाती है, जहां मेमोरी बचत चिकित्सकों को क्लस्टर के बजाय एक एकल GPU पर चलाने की अनुमति देती है।
व्यावहारिक कार्यान्वयन
आधुनिक डीप लर्निंग फ्रेमवर्क में, ग्रेडिएंट चेकपॉइंटिंग आमतौर पर एक सरल API के रूप में उजागर होती है। उदाहरण के लिए, PyTorch में, torch.utils.checkpoint मॉड्यूल एक checkpoint फ़ंक्शन प्रदान करता है जो एक मॉड्यूल या संचालन के अनुक्रम को लपेटता है। जब लपेटा गया मॉड्यूल निष्पादित होता है, तो इसके सक्रियण सहेजे नहीं जाते हैं; इसके बजाय, वे बैकवर्ड पास के दौरान पुनर्गणना किए जाते हैं। TensorFlow tf.recompute_grad के माध्यम से एक समान सुविधा प्रदान करता है। ये कार्यान्वयन बुककीपिंग को स्वचालित रूप से संभालते हैं, जिससे शोधकर्ताओं के लिए अपने मॉडल आर्किटेक्चर को संशोधित किए बिना तकनीक को अपनाना आसान हो जाता है।
ग्रेडिएंट चेकपॉइंटिंग का कम्प्यूटेशनल ओवरहेड नगण्य नहीं है। sqrt(n) चेकपॉइंट्स वाले एक नेटवर्क के लिए, प्रशिक्षण के दौरान कुल फॉरवर्ड गणना मानक प्रशिक्षण की तुलना में लगभग 30-40% बढ़ जाती है। हालांकि, यह लागत अक्सर स्वीकार्य होती है क्योंकि विकल्प - बैच आकार या मॉडल आकार को कम करना - अभिसरण या मॉडल गुणवत्ता को नुकसान पहुंचा सकता है। कई मामलों में, बड़े बैच आकार का उपयोग करने से प्राप्त गति पुनर्गणना ओवरहेड से अधिक होती है।
बड़े मॉडल प्रशिक्षण पर प्रभाव
ग्रेडिएंट चेकपॉइंटिंग बहुत बड़े मॉडलों के प्रशिक्षण की आधारशिला बन गई है। ओपनएआई, एंथ्रोपिक, और गूगल डीपमाइंड जैसी कंपनियां सैकड़ों बिलियन पैरामीटर वाले मॉडलों को प्रशिक्षित करने के लिए इस पर निर्भर हैं। उदाहरण के लिए, 8 GPUs वाले एक एकल नोड पर 70-बिलियन-पैरामीटर मॉडल का प्रशिक्षण, चेकपॉइंटिंग के बिना उन GPUs की संयुक्त मेमोरी से अधिक सक्रियणों को संग्रहीत करने की आवश्यकता होगी। ग्रेडिएंट चेकपॉइंटिंग का उपयोग करके, ये संगठन प्रशिक्षण कार्य को उपलब्ध हार्डवेयर में फिट कर सकते हैं, हालांकि लंबे प्रशिक्षण समय के साथ।
यह तकनीक जनरेटिव एआई अनुप्रयोगों के लिए भी आवश्यक है जिनमें लंबे अनुक्रम शामिल होते हैं, जैसे दस्तावेज़ सारांशीकरण या कोड जनरेशन। इन मामलों में, सक्रियण मेमोरी अनुक्रम लंबाई के साथ बढ़ती है, और चेकपॉइंटिंग मेमोरी सीमा से अधिक हुए बिना लंबे संदर्भों की अनुमति देती है। इसने सीधे 100,000 टोकन या उससे अधिक के संदर्भ विंडो वाले मॉडलों के विकास को सक्षम किया है।
अन्य मेमोरी अनुकूलन से संबंध
ग्रेडिएंट चेकपॉइंटिंग अक्सर अन्य तकनीकों के साथ उपयोग की जाती है। बैच-सामान्यीकरण और परत-सामान्यीकरण सीधे मेमोरी को कम नहीं करते हैं, लेकिन वे प्रशिक्षण स्थिरता में सुधार कर सकते हैं, जो चेकपॉइंटिंग का पूरक है। मॉडल प्रूनिंग पैरामीटरों की संख्या को कम करता है, लेकिन सक्रियण एक बाधा बने रहते हैं, इसलिए चेकपॉइंटिंग अभी भी आवश्यक है। डेटा वृद्धि प्रभावी डेटासेट आकार बढ़ाती है लेकिन सक्रियण मेमोरी को प्रभावित नहीं करती है। वितरित प्रशिक्षण में, ग्रेडिएंट चेकपॉइंटिंग को पाइपलाइन समानांतरता के साथ जोड़ा जा सकता है, जहां विभिन्न परतों को विभिन्न उपकरणों को सौंपा जाता है, ताकि प्रति-डिवाइस मेमोरी दबाव को और कम किया जा सके।
एक उल्लेखनीय विकल्प ग्रेडिएंट संचय है, जो कई छोटे बैचों पर ग्रेडिएंट जमा करके एक बड़े बैच आकार का अनुकरण करता है। यह अनुकूलक स्थितियों के लिए मेमोरी को कम करता है लेकिन सक्रियण मेमोरी को कम नहीं करता है, इसलिए यह चेकपॉइंटिंग का विकल्प नहीं है। एक और संबंधित विचार प्रतिवर्ती परतें हैं, जैसा कि कुछ अवशिष्ट-नेटवर्क वेरिएंट में उपयोग किया जाता है, जहां सक्रियणों को आउटपुट से पुनर्निर्मित किया जा सकता है, लेकिन इसके लिए आर्किटेक्चरल परिवर्तनों की आवश्यकता होती है और यह चेकपॉइंटिंग की तुलना में कम सामान्य है।
सीमाएं और व्यापार-बंद
ग्रेडिएंट चेकपॉइंटिंग की प्राथमिक सीमा प्रति प्रशिक्षण चरण बढ़ा हुआ वॉल-क्लॉक समय है। उन मॉडलों के लिए जो पहले से ही कंप्यूट-बाउंड हैं, अतिरिक्त फॉरवर्ड पास प्रशिक्षण को 20-40% धीमा कर सकता है। इसके अलावा, यह तकनीक मॉडल पैरामीटर या अनुकूलक स्थितियों के लिए मेमोरी को कम नहीं करती है, जो बड़े मॉडलों के लिए भी पर्याप्त हो सकती है। अत्यंत बड़े मॉडलों के लिए, चिकित्सकों को पैरामीटर ऑफलोडिंग के साथ चेकपॉइंटिंग को संयोजित करने या एडब्ल्यूएस ट्रेनियम या ग्रोक जैसे विशेष हार्डवेयर का उपयोग करने की आवश्यकता हो सकती है जिनमें बड़ी ऑन-चिप मेमोरी होती है।
एक और सूक्ष्म मुद्दा यह है कि पुनर्गणना संख्यात्मक अंतर पेश कर सकती है, हालांकि ये व्यवहार में आमतौर पर नगण्य होते हैं। तकनीक को समान सक्रियणों को कई बार पुनर्गणना करने से बचने के लिए सावधानीपूर्वक कार्यान्वयन की भी आवश्यकता होती है, जो ओवरहेड बढ़ाएगा। इन चुनौतियों के बावजूद, ग्रेडिएंट चेकपॉइंटिंग एक व्यापक रूप से उपयोग की जाने वाली और विश्वसनीय विधि बनी हुई है।
भविष्य की दिशाएं
जैसे-जैसे मॉडल बढ़ते रहते हैं, शोधकर्ता अधिक कुशल चेकपॉइंटिंग रणनीतियों की खोज कर रहे हैं। कुछ हालिया कार्य मॉडल की संरचना और हार्डवेयर की मेमोरी प्रोफ़ाइल के आधार पर कौन से सक्रियणों को संग्रहीत करना है, यह तय करने के लिए सीखे गए ह्यूरिस्टिक्स का उपयोग करते हैं। अन्य व्यक्तिगत संचालन के स्तर पर चेकपॉइंटिंग की जांच कर रहे हैं, जिससे बारीक-दानेदार नियंत्रण की अनुमति मिलती है। मेमोरी को और कम करने के लिए चेकपॉइंटिंग को कम-परिशुद्धता प्रशिक्षण के साथ संयोजित करने में भी रुचि है, हालांकि यह सटीकता व्यापार-बंद पेश करता है।
यह तकनीक तब तक प्रासंगिक बने रहने की संभावना है जब तक मेमोरी डीप लर्निंग में एक बाधा है। बड़े भाषा मॉडल और मल्टी-हेड अटेंशन आर्किटेक्चर के उदय के साथ, मेमोरी-कुशल प्रशिक्षण विधियों की मांग केवल बढ़ेगी। ग्रेडिएंट चेकपॉइंटिंग, अन्य अनुकूलन के साथ, कृत्रिम बुद्धिमत्ता में प्रगति का एक प्रमुख सक्षमकर्ता बना रहेगा।
निष्कर्ष
ग्रेडिएंट चेकपॉइंटिंग एक सरल लेकिन शक्तिशाली विचार है: चयनात्मक रूप से सक्रियणों को संग्रहीत करके और बैकप्रोपेगेशन के दौरान उन्हें पुनर्गणना करके, यह गहरे नेटवर्कों को बहुत कम मेमोरी के साथ प्रशिक्षित करने की अनुमति देता है। 2016 में पेश किया गया, यह क्षेत्र में एक मानक अभ्यास बन गया है, जो उन मॉडलों के प्रशिक्षण को सक्षम बनाता है जो अन्यथा उपलब्ध हार्डवेयर पर असंभव होते। जबकि यह कम्प्यूटेशनल ओवरहेड जोड़ता है, व्यापार-बंद अक्सर सार्थक होता है, विशेष रूप से बड़े पैमाने के मॉडलों के लिए। जैसे-जैसे क्षेत्र और भी बड़े आर्किटेक्चर की ओर बढ़ता है, ग्रेडिएंट चेकपॉइंटिंग मशीन लर्निंग टूलबॉक्स में एक आवश्यक उपकरण बना रहेगा।