🤖 machine learning

rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment

यह शोध पत्र rl-triton को प्रस्तुत करता है, जो एक ओपन-सोर्स लाइब्रेरी है जो सात अलग-अलग सुदृढीकरण लर्निंग (reinforcement learning) क्रेडिट असाइनमेंट एल्गोरिदम को GPUs पर त्वरित करने के लिए ट्राइटन (Triton) में कार्यान्वित एक एकीकृत एसोसिएटिव स्कैन फ्रेमवर्क का लाभ उठाती है, जो मेमोरी ओवरहेड को कम करके और $O(\log T)$ समानांतर गणना को सक्षम करके वेक्टरकृत बेसलाइन की तुलना में 1.6–5.70× की गति वृद्धि प्राप्त करती है।

मूल लेखक: Lars Simon Zehnder

प्रकाशित 2026-08-20

📖 6 मिनट में पढ़ें🧠 गहराई से पढ़ें

CC BY 4.0

मूल लेखक: Lars Simon Zehnder

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

आर्टिफिशियल इंटेलिजेंस की दुनिया में, कंप्यूटर को अच्छे निर्णय लेना सिखाने के लिए एक निरंतर संघर्ष बना रहता है। कल्पना कीजिए कि एक रोबोट चलना सीख रहा है या एक प्रोग्राम गेम खेलना सीख रहा है। सुधार करने के लिए, सिस्टम को यह पता लगाना होगा कि किन विशिष्ट कार्यों ने सफलता दिलाई और किन कार्यों से विफलता मिली। इस प्रक्रिया को 'क्रेडिट असाइनमेंट' (credit assignment) कहा जाता है। यह घटनाओं के एक क्रम को पीछे मुड़कर देखने और यह तय करने की क्रिया है कि, "यह कदम अच्छा था," या "वह कदम बुरा था," ताकि सिस्टम अपने भविष्य के व्यवहार को समायोजित कर सके। हालाँकि रोबोट अपना अधिकांश समय दुनिया का अन्वेषण करने या अगले कदम का निर्णय लेने के लिए जटिल गणनाएँ करने में बिता सकता है, लेकिन जिस क्षण उसे अपनी गलतियों से सीखने की आवश्यकता होती है, उसे एक विशिष्ट प्रकार के गणित का प्रदर्शन करना पड़ता है। यह गणित कदमों की एक लंबी सूची को देखने और उनके बीच संबंध जोड़ने के बारे में है, जहाँ एक कदम का मूल्य उसके बाद आने वाले कदम पर निर्भर करता है। लंबे समय तक, शक्तिशाली कंप्यूटर चिप्स जिन्हें जीपीयू (GPU) कहा जाता है, पर यह गणित करना धीमा रहा क्योंकि कंप्यूटर को सूची को एक समय में एक कदम करके प्रोसेस करना पड़ता था, जैसे कि एक किताब के पन्ने दर पन्ने पढ़ना, भले ही हार्डवेयर एक साथ कई पन्ने पढ़ने में सक्षम था।

लार्स साइमन ज़ेंडर नामक एक शोधकर्ता ने 'rl-triton' नामक एक नया टूल विकसित किया है जो इस बाधा को दूर करता है। यह टूल अत्यधिक कुशल कंप्यूटर निर्देशों का एक संग्रह है जिसे विशेष रूप से सुदृढीकरण शिक्षण (reinforcement learning) में क्रेडिट असाइनमेंट के कार्य के लिए डिज़ाइन किया गया है। कंप्यूटर को चरणों की सूची को एक धीमी, क्रमिक श्रृंखला में संसाधित करने के लिए मजबूर करने के बजाय, नई विधि काम को पुनर्गठित करती है ताकि हजारों चरणों की गणना एक साथ की जा सके। मुख्य विचार पूरे घटनाओं के क्रम को एक एकल, एकीकृत गणितीय संरचना के रूप में मानना है जिसे विभाजित किया जा सकता है और समानांतर (parallel) रूप से हल किया जा सकता है। ऐसा करके, कंप्यूटर उस समय लगने वाले समय के एक अंश में गणना पूरी कर सकता है, विशेष रूप से तब जब एक ही समय में हजारों अलग-अलग परिदृश्य चल रहे हों।

शोधकर्ताओं ने वर्तमान में क्षेत्र में उपयोग किए जाने वाले मानक तरीकों के विरुद्ध इस नए दृष्टिकोण का परीक्षण किया। उन्होंने पाया कि सबसे सामान्य और कठिन परिदृश्यों के लिए—जहाँ एक साथ हजारों वातावरणों का अनुकरण (simulation) किया जा रहा है—नया टूल काफी तेज़ है। कुछ मामलों में, इसने पिछले सबसे अच्छे तरीके की तुलना में लगभग छह गुना तेज़ी से कार्य पूरा किया। यह गति सुधार डेटा के कंप्यूटर की मेमोरी के माध्यम से प्रवाहित होने के तरीके में एक चतुर परिवर्तन से आता है। पुराने तरीके में, कंप्यूटर को प्रत्येक अनुक्रम के प्रत्येक चरण के लिए अपने मुख्य मेमोरी बैंक से डेटा लाने के लिए लगातार रुकना पड़ता था, जिससे एक ट्रैफिक जाम जैसी स्थिति बन जाती थी। नई विधि डेटा को गणना इंजन के करीब रखती है, जिससे कंप्यूटर बिना उन निरंतर रुकावटों के पूरे अनुक्रम पर काम कर पाता है। यह आधुनिक एआई प्रशिक्षण के लिए विशेष रूप से महत्वपूर्ण है, जहाँ सिस्टम समानांतर में हजारों सिमुलेशन चला सकते हैं, जिनमें से प्रत्येक में सैकड़ों चरण होते हैं।

पेपर विस्तार से बताता है कि यह सात अलग-अलग प्रकार के लर्निंग एल्गोरिदम के लिए कैसे काम करता है, जिनमें से सभी एक ही अंतर्निहित गणितीय पैटर्न साझा करते हैं। नया टूल उन सभी को एक एकल, एकीकृत ढांचे के साथ संभालता है। यह वास्तविक दुनिया के डेटा की जटिल वास्तविकताओं पर भी सावधानीपूर्वक ध्यान देता है, जैसे कि जब कोई एपिसोड अचानक समाप्त हो जाता है या जब कोई सिमुलेशन बीच में ही कट जाता है। शोधकर्ताओं ने सिद्ध किया कि उनका तरीका इन सीमाओं को सही ढंग से संभालता है, यह सुनिश्चित करते हुए कि लर्निंग सिग्नल सही स्थान पर रुक जाए और एक परिदृश्य से दूसरे में गलती से न मिल जाए। उन्होंने अपने परिणामों को पुराने, पारंपरिक तरीके और एक अधिक आधुनिक, अनुकूलित संस्करण (जो मानक प्रोग्रामिंग टूल का उपयोग करता है) दोनों के विरुद्ध तुलना करके सत्यापित किया। नए टूल ने दोनों को लगातार पीछे छोड़ दिया, जिससे यह साबित हुआ कि गति में वृद्धि वास्तविक थी और केवल बेहतर कोडिंग ट्रिक्स का परिणाम नहीं थी।

सबसे दिलचस्प निष्कर्षों में से एक यह है कि समस्या के आकार के आधार पर गति का लाभ कैसे बदलता है। जब चरणों के अनुक्रम छोटे होते हैं, तो नया टूल तेज़ तो होता है, लेकिन अंतर कम होता है। हालाँकि, जैसे-जैसे अनुक्रम लंबे होते जाते हैं, यह लाभ बढ़ता जाता है। ऐसा इसलिए है क्योंकि पुराने तरीकों को सूची लंबी होने पर मेमोरी-फेचिंग (data fetching) की प्रक्रिया को कई बार दोहराना पड़ता है, जबकि नई विधि बहुत अधिक कुशलता से स्केल करती है। शोधकर्ताओं ने यह भी देखा कि यह एक एआई एजेंट के संपूर्ण प्रशिक्षण प्रक्रिया को कैसे प्रभावित करता है। उन्होंने पाया कि जबकि क्रेडिट असाइनमेंट चरण स्वयं बहुत तेज़ हो गया, समग्र प्रशिक्षण गति में सुधार कभी-कभी मामूली था। ऐसा इसलिए है क्योंकि क्रेडिट असाइनमेंट पूरे प्रशिक्षण पाइपलाइन का केवल एक हिस्सा है; यदि प्रक्रिया के बाकी हिस्से धीमे हैं, तो केवल एक हिस्से को तेज़ करने से पूरी चीज़ नाटकीय रूप से तेज़ नहीं चलेगी। हालाँकि, विशिष्ट सेटअपों में जहाँ क्रेडिट असाइनमेंट चरण कुल समय का एक बड़ा हिस्सा लेता है, वहाँ समग्र प्रशिक्षण गति में उल्लेखनीय सुधार हुआ।

यह कार्य कुछ सीमाओं को भी रेखांकित करता है। बहुत लंबे अनुक्रमों के लिए, 'रिट्रेस' (Retrace) नामक एक विशिष्ट प्रकार का एल्गोरिदम एक हार्डवेयर बाधा का सामना करता है जहाँ कंप्यूटर चिप एक विशिष्ट प्रकार के तेज़ स्टोरेज स्पेस की कमी महसूस करती है, जिससे धीमापन आता है। शोधकर्ताओं ने इस समस्या की पहचान की और नोट किया कि यह डिज़ाइन में एक ज्ञात ट्रेड-ऑफ (trade-off) है। उन्होंने यह भी उल्लेख किया कि उनका वर्तमान टूल मानक डेटा प्रारूपों के साथ सबसे अच्छा काम करता है और कुछ विशेष रूपांतरों के लिए आगे के विकास की आवश्यकता हो सकती है। इन सीमाओं के बावजूद, पेपर एआई प्रशिक्षण की एक निरंतर समस्या के लिए एक स्पष्ट और व्यावहारिक समाधान प्रस्तुत करता है। क्रमिक, चरण-दर-चरण गणना को समानांतर, एक साथ की जाने वाली गणना में बदलकर, शोधकर्ताओं ने दिखाया है कि सुदृढीकरण शिक्षण (reinforcement learning) को काफी अधिक कुशल बनाना संभव है। यह दक्षता जैसे-जैसे एआई सिस्टम बड़े और अधिक जटिल होते जा रहे हैं, उनके लिए महत्वपूर्ण है, क्योंकि उन्हें कम समय में विशाल मात्रा में डेटा से सीखने की आवश्यकता होती है। यह टूल अब दूसरों के उपयोग के लिए उपलब्ध है, जो उनके सीखने के मौलिक तरीके को बदले बिना बुद्धिमान प्रणालियों के प्रशिक्षण को गति देने का एक तरीका प्रदान करता है।

तकनीकी सारांश: rl-triton

समस्या विवरण

डीप रिइन्फोर्समेंट लर्निंग (RL) ट्रेनिंग पाइपलाइनों को प्रत्येक अपडेट स्टेप पर क्रेडिट असाइनमेंट मात्राओं—जैसे कि एडवांटेज (advantages), वैल्यू टारगेट्स (value targets), और रिटर्न एस्टिमेट्स (return estimates)—की गणना करने की आवश्यकता होती है। ये गणनाएँ $[num\_envs, seq\_len]$ आकार के टेंसरों पर संचालित होती हैं और आमतौर पर प्रथम-क्रम के रैखिक पुनरावृत्तियों (first-order linear recurrences) (जैसे, Generalized Advantage Estimation, V-Trace, Retrace) में शामिल होती हैं।

मुख्य चुनौती इन पुनरावृत्तियों में निहित अनुक्रमिक निर्भरता (sequential dependency) है। एक मानक कार्यान्वयन को अनुक्रम (sequence) की गणना करने के लिए $T$ सीरियल स्टेप्स की आवश्यकता होती है, जो एक सख्त टेम्पोरल चेन बनाता है जो हजारों कोर वाले GPUs पर भी पूर्ण समानांतरकरण (parallelization) को रोकता है। जबकि torch.compile पायथन इंटरप्रेटर ओवरहेड को हटा देता है, यह एल्गोरिदम द्वारा थोपी गई डेटा निर्भरता को नहीं तोड़ सकता, जिससे कम्प्यूटेशनल डेप्थ $O(T)$ बनी रहती है। इसके अलावा, मौजूदा वेक्टरकृत (vectorized) कार्यान्वयन जो समस्या को एक समानांतर एसोसिएटिव स्कैन (parallel associative scan) के रूप में पुनर्गठित करते हैं, अक्सर उच्च मेमोरी ट्रैफिक से ग्रस्त होते हैं: वे $O(\log T)$ रिडक्शन के प्रत्येक दोहरीकरण चरण में हाई बैंडविड्थ मेमोरी (HBM) में मध्यवर्ती परिणाम (intermediate results) को भौतिक रूप से निर्मित (materialize) करते हैं, जिससे महत्वपूर्ण लेटेंसी और बैंडविड्थ लागत आती है।

कार्यप्रणाली (Methodology)

यह शोध पत्र rl-triton प्रस्तुत करता है, जो Triton भाषा का उपयोग करके RL क्रेडिट असाइनमेंट के लिए उच्च-प्रदर्शन वाले GPU कर्नेल्स को लागू करने वाली एक ओपन-सोर्स लाइब्रेरी है। इसकी कार्यप्रणाली तीन स्तंभों पर आधारित है:

1. एकीकृत बीजगणितीय ढांचा (Unified Algebraic Framework)

लेखक यह प्रदर्शित करते हैं कि सात अलग-अलग RL एस्टीमेशन एल्गोरिदम—GAE, V-Trace, Retrace( $\lambda$ ), TD( $\lambda$ ) रिटर्न्स, डिस्काउंटेड रिटर्न्स, एलिजिबिलिटी ट्रेसेस (eligibility traces), और एपिसोडिक प्रीफिक्स sums—सभी एक ही प्रथम-क्रम रैखिक पुनरावृत्ति (first-order linear recurrence) के उदाहरण हैं:
$A_t = \alpha_t + \beta_t \cdot A_{t+1}$
(या एलिजिबिलिटी ट्रेसेस के लिए इसके फॉरवर्ड-टाइम मिरर)।
इन एल्गोरिदम को एक विशिष्ट एसोसिएटिव ऑपरेटर $\oplus$ का उपयोग करके टुपल्स $(\alpha_t, \beta_t)$ पर समानांतर एसोसिएटिव स्कैन (parallel associative scans) के रूप में पुनर्गठित किया गया है:
$(\alpha_B, \beta_B) \oplus (\alpha_A, \beta_A) = (\alpha_B + \beta_B \alpha_A, \beta_A \beta_B)$
यह रूपांतरण इस बात की अनुमति देता है कि पुनरावृत्ति को $O(T)$ अनुक्रमिक स्टेप्स के बजाय $O(\log T)$ समानांतर स्टेप्स में हल किया जा सके।

2. फ्यूज्ड ट्राइटन कर्नेल्स (Fused Triton Kernels)

मानक PyTorch कार्यान्वयन के विपरीत, जो स्कैन करने से पहले HBM में मध्यवर्ती गुणांक टेंसर ( $\alpha$ और $\beta$ ) को भौतिक रूप से निर्मित करते हैं, rl-triton फ्यूज्ड कर्नेल्स (fused kernels) का उपयोग करता है।

ऑन-चिप निर्माण (On-Chip Construction): कर्नेल्स कच्चे रोलआउट इनपुट्स (रिवॉर्ड्स, वैल्यूज, डन फ्लैग्स, आदि) को लोड करते हैं और सीधे रजिस्टरों में पुनरावृत्ति गुणांकों ( $\alpha_t, \beta_t$ ) का निर्माण करते हैं।
ऑन-चिप स्कैन (On-Chip Scan): एसोसिएटिव स्कैन tl.associative_scan के माध्यम से पूरी तरह से ऑन-चिप (रजिस्टरों और शेयर्ड मेमोरी का उपयोग करके) निष्पादित किया जाता है। अंतिम आउटपुट तक मध्यवर्ती स्कैन स्टेट्स को कभी भी HBM में नहीं लिखा जाता है।
मेमोरी दक्षता (Memory Efficiency): यह दृष्टिकोण HBM राउंड-ट्रिप्स को कम करता है। लंबाई $T$ के अनुक्रम के लिए, फ्यूज्ड कर्नेल एक एकल HBM रीड और राइट करता है, जबकि एक मटेरियलइज्ड स्कैन के लिए $O(\log T)$ HBM राउंड-ट्रिप्स की आवश्यकता होती है।

3. स्पष्ट सीमा सिमेंटिक्स (Explicit Boundary Semantics)

लाइब्रेरी जटिल एपिसोड बाउंड्री कंडीशंस को स्पष्ट रूप से संभालती है जो अक्सर मानक कार्यान्वयन में त्रुटियां पैदा करती हैं:

टर्मिनेशन बनाम ट्रंकेशन (Termination vs. Truncation): यह प्राकृतिक एपिसोड एंडिंग्स (जहाँ उत्तराधिकारी मान शून्य होता है) और ट्रंकेशन (जहाँ एक आपूर्ति किया गया बूटस्ट्रैप मान उपयोग किया जाना चाहिए) के बीच अंतर करती है।
रोलआउट विंडोज़ (Rollout Windows): यह रोलआउट बफर के अंत ( $t=T-1$ ) पर सीमा को सही ढंग से प्रबंधित करती है, यह सुनिश्चित करती है कि बूटस्ट्रैप मानों को स्थानीय अपडेट ( $\alpha_{T-1}$ ) या स्कैन कैरी ( $A_T$ ) में शामिल किया जाए, ताकि विशिष्ट एल्गोरिदम के आधार पर डबल-काउंटिंग न हो।

मुख्य योगदान

एकीकृत ढांचा: एक प्रमाण कि सात मानक RL क्रेडिट असाइनमेंट एल्गोरिदम एक ही एफाइन स्कैन ऑपरेटर के उदाहरण हैं।
फ्यूज्ड कार्यान्वयन: एल्गोरिदम-विशिष्ट Triton कर्नेल्स जो ऑन-चिप गुणांकों का निर्माण करते हैं और बिना मध्यवर्ती HBM मटेरियलइजेशन के स्कैन निष्पादित करते हैं।
मजबूत बाउंड्री हैंडलिंग: टर्मिनेटेड स्टेप्स, ट्रंकेटेड स्टेप्स और रोलआउट-विंडो सीमाओं के लिए स्पष्ट, सत्यापित लॉजिक।
परफॉरमेंस बेंचमार्क्स: मानक अनुक्रमिक लूप्स और वेक्टरकृत torch.compile बेसलाइन्स पर स्पीडअप दिखाने वाला व्यापक मूल्यांकन।
ओपन सोर्स रिलीज़: करेक्टनेस टेस्ट, परफॉरमेंस रिग्रेशन टेस्ट और डॉक्यूमेंटेशन के साथ एक पायथन पैकेज।

परिणाम

लाइब्रेरी का मूल्यांकन NVIDIA H100 (80GB HMB3) और RTX 2000 Ada GPUs पर विभिन्न एनवायरनमेंट काउंट और सीक्वेंस लेंथ केAcross किया गया।

स्टैंडर्ड लूप्स पर स्पीडअप: मासिवली पैरेलल सिमुलेशन रिजीम (जैसे, 4,096 एनवायरनमेंट्स, 128 स्टेप्स) में, rl-triton अनकंपाइल्ड अनुक्रमिक लूप्स की तुलना में 16–317× की स्पीडअप प्राप्त करता है। यह मुख्य रूप से $O(T)$ से $O(\log T)$ डिपेंडेंसी डेप्थ में कमी के कारण है।
वेक्टरकृत बेसलाइन्स पर स्पीडअप: एक वेक्टरकृत torch.compile बेसलाइन (जो $O(\log T)$ $O (lo g T)$ स्कैन का उपयोग करता है लेकिन इंटरमीडिएट्स को मटेरियल करता है) की तुलना में, rl-triton 1.6–5.70× फुल-कॉल स्पीडअप प्राप्त करता है।
- स्पीडअप उन एल्गोरिदम के लिए सबसे अधिक स्पष्ट है जिनमें कम प्रति-स्टेप कम्प्यूटेशनल कॉम्प्लेक्सिटी है (जैसे, डिस्काउंटेड रिटर्न्स, $\lambda$ -रिटर्न्स) और लंबे सीक्वेंस लेंथ, जहाँ बेसलाइन के बार-बार होने वाले HBM राउंड-ट्रिप्स बॉटलनेक बन जाते हैं।
- Retrace( $\lambda$ ) के लिए, स्पीडअप कम (1.6–1.9×) है क्योंकि इसमें भारी प्रति-स्टेप कम्प्यूटेशन (एक्शन प्रोबेबिलिटीज लोड करना) और बहुत लंबे सीक्वेंस लेंथ ( $>2048$ ) पर रजिस्टर प्रेशर की समस्या होती है।
एंड-टू-एंड प्रभाव: एक सिंथेटिक PPO अपडेट में, कर्नेल स्पीडअप ने 1.11–1.16× एंड-टू-एंड ट्रेनिंग स्पीडअप में अनुवाद किया जब GAE ने कुल अपडेट समय का 10–15% हिस्सा लिया। जब GAE का हिस्सा छोटा (~~2.2%) था, तो एंड-टू-एंड लाभ नगण्य (~~1.02×) था, जो एम्डाहल के नियम (Amdahl's Law) के अनुरूप है।
PufferLib के साथ तुलना: rl-triton 40 में से 31 टेस्टेड शेप्स के लिए मासिवली पैरेलल सिमुलेशन रिजीम (शॉर्ट रोलआउट्स, हाई एनवायरनमेंट काउंट्स) में PufferLib के हैंड-रिटन CUDA कर्नेल से बेहतर प्रदर्शन करता है, विशेष रूप से जैसे-जैसे सीक्वेंस लेंथ बढ़ती है।

महत्व और दावे

शोध पत्र का दावा है कि rl-triton RL में क्रेडिट असाइनमेंट के क्रिटिकल पाथ के लिए एक एकीकृत, उच्च-प्रदर्शन समाधान प्रदान करता है। एसोसिएटिव स्कैन फ्रेमवर्क और Triton की फ्यूजन क्षमताओं का लाभ उठाकर, यह उस मेमोरी बैंडविड्थ बॉटलनेक को समाप्त करता है जो मानक कार्यान्वयन और मानक वेक्टरकृत PyTorch दृष्टिकोणों दोनों को प्रभावित करता है।

लेखक इस बात पर जोर देते हैं कि यह लाइब्रेरी आधुनिक डिस्ट्रिब्यूटेड ट्रेनिंग और RLHF पाइपलाइनों में सामान्य मासिवली पैरेलल सिमुलेशन रिजीम (हजारों एनवायरनमेंट्स, छोटे रोलआउट्स) के लिए विशेष रूप से महत्वपूर्ण है, जहाँ HBM ट्रैफिक में कमी से वास्तविक समय (wall-clock time) की बचत होती है। हालांकि लाइब्रेरी पूरे ट्रेनिंग पाइपलाइन को तेज करने का दावा नहीं करती है (क्योंकि नेटवर्क फॉरवर्ड/बैकवर्ड पासेज अक्सर हावी होते हैं), यह क्रेडिट असाइनमेंट घटक के लिए एक महत्वपूर्ण अनुकूलन प्रदान करती है, जो प्रत्येक अपडेट स्टेप में एक सख्त डिपेंडेंसी है। यह कार्य सेगमेंटेड एफाइन रिकरेंस (segmented affine recurrences) से जुड़े अन्य डोमेन, जैसे कि लॉन्ग-कॉन्टेक्स्ट LLM पोस्ट-ट्रेनिंग और स्पाइकिंग न्यूरल नेटवर्क में भी फ्यूज्ड स्कैन की प्रयोज्यता को उजागर करता है।