🤖 machine learning

Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis Rotation

यह शोध पत्र यह पहचानता है कि एसिंक्रोनस पाइपलाइन पैरेललिज्म (asynchronous pipeline parallelism) में ग्रेडिएंट स्टेलेनेस (gradient staleness), हेसियन आइगेनबेसिस (Hessian eigenbasis) और कोऑर्डिनेट सिस्टम के बीच मिसअलाइनमेंट के कारण बढ़ जाती है, और एक "बेसिस रोटेशन" (basis rotation) फ्रेमवर्क का प्रस्ताव करता है ताकि इन बेसेस को पुनः संरेखित किया जा सके, जिससे बड़े पैमाने पर वितरित प्रशिक्षण (distributed training) के लिए प्रशिक्षण इटरेशन में काफी कमी आती है और स्केलेबिलिटी बहाल होती है।

मूल लेखक: Hyunji Jung, Sungbin Shin, Namhoon Lee

प्रकाशित 2026-05-28

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

CC BY 4.0

मूल लेखक: Hyunji Jung, Sungbin Shin, Namhoon Lee

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

मुख्य विचार: असेंबली लाइन की समस्या

कल्पना कीजिए कि आप एक विशाल, जटिल रोबोट (एक लार्ज लैंग्वेज मॉडल) बना रहे हैं और इसके लिए एक बहुत बड़ी फैक्ट्री है। इसे तेज़ी से बनाने के लिए, आपने इंजीनियरों की एक टीम (कंप्यूटर) रखी है और काम को एक असेंबली लाइन में बाँट दिया है।

श्रमिक 1 पैर बनाता है।
श्रमिक 2 धड़ (torso) बनाता है।
श्रमिक 3 सिर बनाता है।
श्रमिक 4 इन सबको जोड़कर पूरा करता है।

पुराना तरीका (Synchronous): हर कोई इंतज़ार करता है। श्रमिक 2 तब तक धड़ बनाना शुरू नहीं कर सकता जब तक श्रमिक 1 पैर बनाना खत्म न कर दे। श्रमिक 3, श्रमिक 2 का इंतज़ार करता है। इससे "बबल्स" (bubbles) या खाली समय पैदा होता है जहाँ श्रमिक अपने से पहले वाले व्यक्ति का इंतज़ार करते हुए खाली खड़े रहते हैं। यह सुरक्षित है, लेकिन धीमा है।

नया तरीका (Asynchronous): इस खाली समय को ठीक करने के लिए, फैक्ट्री मैनेजर कहता है, "इंतज़ार मत करो! जैसे ही तुम एक हिस्सा पूरा करो, तुरंत अगला हिस्सा बनाना शुरू कर दो।" श्रमिक 2 वर्तमान रोबोट के पैर बनने के दौरान ही अगले रोबोट के लिए धड़ बनाना शुरू कर देता है। इससे हर कोई व्यस्त रहता है और काम की गति बहुत बढ़ जाती है।

समस्या: "पुराने" (Stale) निर्देश
यहाँ एक पेंच है: क्योंकि हर कोई इतनी तेज़ी से काम कर रहा है, इसलिए श्रमिक 2 को अपने काम को सुधारने के लिए जो निर्देश (gradients) मिलते हैं, वे उस रोबोट पर आधारित होते हैं जिसे श्रमिक 1 ने मिनटों पहले बनाया था। जब तक श्रमिक 2 को निर्देश मिलता है, तब तक रोबोट का आकार बदल चुका होता है। वह निर्देश "स्टेल" (stale/पुराना) है।

यदि निर्देश बहुत पुराने हैं, तो श्रमिक 2 उस हिस्से को ठीक करने की कोशिश कर सकता है जो अब अस्तित्व में ही नहीं है, या ऐसा बदलाव कर सकता है जो वास्तव में रोबोट को खराब कर दे। शोध में पाया गया है कि जैसे-जैसे फैक्ट्री बड़ी होती जाती है (अधिक श्रमिक/स्टेज), ये "स्टेल" निर्देश और भी खराब होते जाते हैं, जिससे पूरा प्रोजेक्ट धीमा हो जाता है या विफल हो जाता है।

छिपा हुआ अपराधी: "गलत दिशा दिखाने वाला कंपास"

लेखकों ने खोजा कि ये "स्टेल" निर्देश इतनी तबाही क्यों मचाते हैं। समस्या सिर्फ यह नहीं है कि निर्देश पुराने हैं; समस्या यह है कि श्रमिक एक खराब कंपास का उपयोग कर रहे हैं।

परिदृश्य (The Landscape): कल्पना करें कि रोबोट बनाने की प्रक्रिया एक पहाड़ से नीचे उतरकर सबसे निचली घाटी (सबसे अच्छा मॉडल) खोजने जैसी है। ज़मीन समतल नहीं है; इसमें ऐसे पहाड़ और घाटियाँ हैं जो मुड़ती और घूमती हैं।
कंपास (The Optimizer): फैक्ट्री एक स्मार्ट कंपास (जिसे Adam कहा जाता है) का उपयोग करती है, जो श्रमिकों को बताता है कि हर दिशा (उत्तर, दक्षिण, पूर्व, पश्चिम) में ढलान कितनी है ताकि वे कुशलता से चल सकें।
गलत संरेखण (The Misalignment): समस्या यह है कि कंपास का "उत्तर" वास्तविक पहाड़ के "उत्तर" से मेल नहीं खाता। कंपास घूम गया है (rotated)।
- जब कंपास सही संरेखित होता है, तो एक पुराना निर्देश बस थोड़ा सा गलत होता है।
- जब कंपास गलत संरेखित (misaligned) होता है, तो एक पुराना निर्देश श्रमिक को पूरी तरह से गलत दिशा में भेज देता है। वे घाटी की ओर चलने के बजाय गोल-गोल घूमने (oscillating) लगते हैं।

पेपर का दावा है कि बड़ी फैक्ट्रियों में जिनमें कई श्रमिक होते हैं, यह गलत संरेखण और भी बुरा हो जाता है, जिससे "स्टेल" निर्देश खतरनाक बन जाते हैं। श्रमिक भ्रमित हो जाते हैं, गोल-गोल घूमते हैं, और रोबमाल कभी बन ही नहीं पाता।

समाधान: "बेसिस रोटेशन" (कंपास को फिर से संरेखित करना)

लेखक बेसिस रोटेशन (Basis Rotation) नामक एक चतुर समाधान प्रस्तावित करते हैं।

श्रमिकों को ऐसा कंपास इस्तेमाल करने के लिए मजबूर करने के बजाय जो पहाड़ के सापेक्ष घूम गया है, वे स्वयं कंपास को घुमा देते हैं ताकि वह पहाड़ के आकार से पूरी तरह मेल खा सके।

रोटेशन (घुमाव): वे पहाड़ के वास्तविक आकार (Hessian eigenbasis) की गणना करते हैं और भौतिक रूप से कंपास को घुमाते हैं ताकि उसका "उत्तर" ठीक वहीं इशारा करे जहाँ पहाड़ की सबसे तीव्र ढलान है।
परिणाम: अब, भले ही निर्देश थोड़े पुराने (stale) हों, फिर भी वे सही सामान्य दिशा में इशारा करते हैं। श्रमिक अब गोल-गोल नहीं घूमते। वे देरी के बावजूद, सीधे घाटी की ओर नीचे उतरते हैं।

उपमा (Analogy): कल्पना कीजिए कि आप एक घुमावदार सड़क पर कार चला रहे हैं और एक ऐसे नक्शे को देख रहे हैं जो उल्टा है। यदि आपको 5 मिनट पुराना ट्रैफिक रिपोर्ट मिलता है, तो आप गलत दिशा में मुड़ सकते हैं और दुर्घटनाग्रस्त हो सकते हैं। लेकिन यदि आप नक्शे को घुमाकर सड़क के अनुरूप कर देते हैं, तो एक पुराना ट्रैफिक रिपोर्ट भी आपको सही रास्ते पर रहने में मदद करेगा।

उन्होंने क्या पाया (परिणाम)

टीम ने बहुत बड़े AI मॉडल (3 बिलियन पैरामीटर्स तक) और कई श्रमिकों (32 स्टेज तक) पर इसका परीक्षण किया।

इस सुधार के बिना: जैसे-जैसे उन्होंने अधिक श्रमिक जोड़े, प्रशिक्षण धीमा होता गया। 32 श्रमिकों के साथ, यह केवल एक श्रमिक के मुकाबले लगभग 6 गुना धीमा था क्योंकि "स्टेल" निर्देश अराजकता पैदा कर रहे थे।
बेसिस रोटेशन के साथ: अराजकता समाप्त हो गई। प्रशिक्षण की गति 32 श्रमिकों के साथ भी तेज़ बनी रही।
- उन्होंने पिछले सर्वोत्तम तरीकों की तुलना में समान गुणवत्ता का रोबोट 81.7% तेज़ी से (कदमों के मामले में) प्राप्त किया।
- यह इतना प्रभावी रहा कि वे बिना किसी दक्षता खोए एक विशाल फैक्ट्री (32 स्टेज) का उपयोग कर सके।

सारांश

समस्या: एसिंक्रोनस ट्रेनिंग (बिना इंतज़ार किए काम करना) "स्टेल" निर्देश पैदा करती है जो AI को भ्रमित करते हैं, खासकर जब फैक्ट्री बहुत बड़ी हो।
कारण: AI का आंतरिक "कंपास" समस्या के आकार के साथ गलत संरेखित है, जिससे पुराने निर्देश खतरनाक हो जाते हैं।
समाधान: बेसिस रोटेशन कंपास को समस्या के बिल्कुल अनुरूप घुमा देता है।
परिणाम: अब AI विशाल, मल्टी-स्टेज फैक्ट्रियों पर बिना धीमे हुए प्रशिक्षित हो सकता है, जिससे यह विशाल AI मॉडल बनाने में बहुत अधिक कुशल हो जाता है।

यह पेपर चिकित्सा उपयोगों या इन विशिष्ट AI मॉडलों के प्रशिक्षण के अलावा भविष्य के अनुप्रयोगों के बारे में चर्चा नहीं करता है; यह पूरी तरह से प्रशिक्षण प्रक्रिया में इस विशिष्ट गणितीय बाधा (bottleneck) को ठीक करने पर केंद्रित है।

तकनीकी सारांश: बेसिस रोटेशन के माध्यम से एसिंक्रोनस पाइपलाइन पैरेललिज्म में स्टेलेनेस (Staleness) को कम करना

1. समस्या विवरण (Problem Statement)

एसिंक्रोनस पाइपलाइन पैरेललिज्म कई उपकरणों (devices) पर बड़े पैमाने के भाषा मॉडल (LLMs) को प्रशिक्षित करने के लिए एक महत्वपूर्ण तकनीक है, जिसका उद्देश्य सिंक्रोनस निष्पादन (synchronous execution) में निहित "पाइपलाइन बबल्स" (idle periods) को समाप्त करना है। हालाँकि, यह दक्षता लाभ ग्रेडिएंट स्टेलेनेस (gradient staleness) द्वारा बाधित होता है: ग्रेडिएंट गणना और उसके अनुप्रयोग के बीच का लौकिक अंतराल (temporal gap) यह सुनिश्चित करता है कि अपडेट पुराने मॉडल स्टेट्स का उपयोग करके लागू किए जाते हैं।

लेखक एक महत्वपूर्ण, अक्सर अनदेखी की जाने वाली विकृति (pathology) की पहचान करते हैं: ग्रेडिएंट विलंब (gradient delay) पाइपलाइन गहराई के साथ रैखिक रूप से बढ़ता है। बड़े पैमाने के कॉन्फ़िगरेशन में जहाँ चरणों (stages) की संख्या दस या सैकड़ों तक पहुँच सकती है, यह विलंब एसिंक्रोनस प्रशिक्षण की स्केलेबिलिटी को मौलिक रूप से कमजोर कर देता है। अनुभवजन्य अवलोकन (Empirical observations) दर्शाते हैं कि निश्चित मॉडल के लिए चरणों की संख्या बढ़ाने से अभिसरण (convergence) में भारी गिरावट आ सकती है (उदाहरण के लिए, 1 से 32 चरणों पर जाने पर 5.81-गुना धीमा होना)।

मुख्य मुद्दा केवल विलंब नहीं है, बल्कि कोऑर्डिनेट-वाइज एडेप्टिव ऑप्टिमाइज़र्स (coordinate-wise adaptive optimizers) जैसे Adam के साथ इसका इंटरेक्शन है, जो LLM प्री-ट्रेनिंग के लिए मानक हैं। लेखक तर्क देते हैं कि Adam की प्रभावशीलता हेसियन आइजनबेसिस (Hessian eigenbasis) और मानक कोऑर्डिनेट बेसिस के बीच संरेखण (alignment) पर निर्भर करती है। जब ये मिसअलाइन्ड होते हैं, तो Adam की कोऑर्डिनेट-वाइज एडेप्टिविटी विफल हो जाती है, जिससे अपडेट ट्रैजेक्टरी में दोलन (oscillations) उत्पन्न होते हैं। विलंब के तहत, ये दोलन विलंबित ग्रेडिएंट्स को पुराने या प्रतिकूल दिशाओं की ओर मोड़ देते हैं, जिससे अभिसरण गंभीर रूप प्रकार से खराब हो जाता है।

2. कार्यप्रणाली: बेसिस रोटेशन (Methodology: Basis Rotation)

इसे संबोधित करने के लिए, पेपर बेसिस रोटेशन का प्रस्ताव करता है, जो एक ऐसा फ्रेमवर्क है जो ऑप्टिमाइज़ेशन स्पेस को रूपांतरित करता है ताकि हेसियन आइजनबेसिस को मानक कोऑर्डिनेट बेसिस के साथ संरेखित किया जा सके, जिससे विलंब की उपस्थिति में एडेप्टिव ऑप्टिमाइज़र्स की प्रभावकारिता को बहाल किया जा सके।

मुख्य तंत्र (Core Mechanism)

मानक कोऑर्डिनेट सिस्टम में सीधे वेट्स $w_t$ को अपडेट करने के बजाय, यह विधि एक रोटेशन मैट्रिक्स $U$ (जिसके कॉलम हेसियन के अनुमानित आइजनवेक्टर्स हैं) का उपयोग करके ऑप्टिमाइज़ेशन स्पेस को रोटेट करती है। अपडेट रोटेटेड स्पेस $\tilde{w} = U^\top w$ में किया जाता है:
$\tilde{w}_t = \tilde{w}_{t-1} - \eta_{t-1} \frac{\text{EMA}(U^\top \nabla f(U \tilde{w}_{t-1}))}{\sqrt{\text{EMA}((U^\top \nabla f(U \tilde{w}_{t-1}))^2) + \epsilon}}$
इसके बाद अपडेट को मूल स्पेस में प्रोजेक्ट किया जाता है:
$w_t = w_{t-1} - \eta_{t-1} U \left( \frac{\text{EMA}(U^\top \nabla f(w_{t-1}))}{\sqrt{\text{EMA}((U^\top \nabla f(w_{t-1}))^2) + \epsilon}} \right)$
बेसिस को संरेखित करके, यह विधि सुनिश्चित करती है कि विलंबित अपडेट दिशाएं वास्तविक ऑप्टिमाइज़ेशन पथ के साथ सुसंगत बनी रहें, जिससे स्टेलेनेस से उत्पन्न होने वाले शोर (noise) को कम किया जा सके।

व्यावहारिक कार्यान्वयन (Practical Implementation)

बड़े मॉडलों के लिए इसे कम्प्यूटेशनल रूप से व्यवहार्य बनाने के लिए, लेखक दो संरचनात्मक धारणाओं को अपनाते हैं:

ब्लॉक-डायगोनैलिटी (Block-diagonality): यह माना जाता है कि हेसियन ब्लॉक-डायगोनल है, जिससे रोटेशन को पूरे पैरामीटर स्पेस के बजाय प्रति वेट मैट्रिक्स लागू किया जा सकता है।
क्रोनेकर फैक्टराइजेशन (Kronecker Factorization): एक वेट मैट्रिक्स $W \in \mathbb{R}^{m \times n}$ के लिए रोटेशन मैट्रिक्स $U_W$ को $U \in \mathbb{R}^{m \times m}$ और $V \in \mathbb{R}^{n \times n}$ में फैक्टराइज़ किया जाता है, जिससे ऑपरेशन सुलभ हो जाता है।

आइजनबेसिस अनुमान रणनीतियाँ (Eigenbasis Estimation Strategies)

पेपर रोटेशन मैट्रिसेस के अनुमान के लिए दो अक्षों पर आधारित एक वर्गीकरण (taxonomy) पेश करता है:

अनुमान स्रोत ( $S$ ):
- 2nd-order: एम्पेरिकल फिशर (empirical Fisher) को अनुमानित करने के लिए सेकंड मोमेंट्स ( $L = E[GG^\top]$ ) का उपयोग करता है।
- 1st-order: हेसियन को अनुमानित करने के लिए फर्स्ट मोमेंट्स (मोमेंटम बफर्स $M_t$ ) का उपयोग करता है, जिससे मेमोरी ओवरहेड कम होता है।
रोटेशन ज्योमेट्री ( $G$ ):
- Bilateral: पूर्ण क्रोनेकर संरचना को कैप्चर करने के लिए टू-साइडेड रोटेशन ( $U^\top W V$ ) लागू करता है।
- Unilateral: छोटे आयाम (dimension) पर वन-साइडेड रोटेशन लागू करता है।

लेखक आइजनवेक्टर्स को कुशलतापूर्वक कंप्यूट करने के लिए पावर इटरेशन के एक सिंगल स्टेप के बाद QR डिकंपोजिशन का उपयोग करते हैं, और फिडेलिटी और लागत के बीच संतुलन बनाने के लिए बेसिस को अंतराल पर (जैसे, हर 10 इटरेशन में) अपडेट करते हैं।

3. प्रमुख योगदान (Key Contributions)

एक स्केलिंग बॉटलनेक की पहचान: पेपर प्रदर्शित करता है कि एसिंक्रोनस पाइपलाइन पैरेललिज्म चरणों की संख्या बढ़ने के साथ गंभीर अभिसरण गिरावट (convergence degradation) का सामना करता है, जो एक ऐसी घटना है जिसे औपचारिक साहित्य में पहले कम अध्ययन किया गया है।
विफलता तंत्र का स्पष्टीकरण: लेखक बेसिस मिसअलाइनमेंट को प्राथमिक कारण के रूप में पहचानते हैं कि क्यों Adam-टाइप ऑप्टिमाइज़र्स विलंब के प्रति संवेदनशील हैं। वे बेसिस मिसअलाइनमेंट (जिसे हेसियन के $(1,1)$ -नॉर्म द्वारा मापा जाता है) के कारण विलंब दंड (delay penalties) मल्टीप्लिकेटिव रूप से बढ़ जाते हैं, यह सिद्ध करने के लिए अनुभवजन्य साक्ष्य और सैद्धांतिक अभिसरण विश्लेषण प्रदान करते हैं।
बेसिस रोटेशन फ्रेमवर्क: ऑप्टिमाइज़ेशन ट्रैजेक्टरी को पुन: संरेखित करने का एक अभिनव तरीका, जो सैद्धांतिक रूप से बेसिस मिसअलाइनमेंट को कम करने और विलंब-प्रेरित अस्थिरता का मुकाबला करने के लिए प्रमाणित है।
स्केल पर अनुभवजन्य सत्यापन: 3B पैरामीटर्स तक के LLM प्री-ट्रेनिंग पर व्यापक प्रयोग जो इस विधि की प्रभावकारिता को प्रदर्शित करते हैं।

4. प्रयोगात्मक परिणाम (Experimental Results)

लेखकों ने 95M से 3B पैरामीटर्स वाले डिकोडर-ओनली ट्रांसफॉर्मर्स पर स्टैंडर्ड बेसलाइन्स (PipeDream, PipeDream-LR, Nesterov) के विरुद्ध बेसिस रोटेशन का मूल्यांकन किया।

अभिसरण गति (Convergence Speed): 32 पाइपलाइन स्टेज के साथ 3B-पैरामीटर मॉडल के प्रशिक्षण में, बेसिस रोटेशन ने सर्वश्रेष्ठ प्रदर्शन करने वाले एसिंक्रोनस बेसलाइन की तुलना में 81.7% कम इटरेशन के साथ समान ट्रेनिंग लॉस प्राप्त किया।
विलंब के प्रति मजबूती (Robustness to Delay): जबकि स्टैंडर्ड बेसलाइन्स ने 1 से 32 स्टेज तक बढ़ने पर 4.24x से 5.81x का धीमापन दिखाया, बेसिस रोटेशन ने इस धीमेपन को घटाकर 1.27x कर दिया।
स्केलेबिलिटी (Scalability): बेसलाइन्स के विपरीत, जो मॉडल साइज बढ़ने के साथ बढ़ता हुआ ट्रेनिंग लॉस दिखाते थे (जो मानक स्केलिंग लॉ का उल्लंघन करता है), बेसिस रोटेशन ने सफलतापूर्वक स्केलेबिलिटी को बहाल किया, जो मॉडल साइज बढ़ने के साथ प्रदर्शन में सुधार दिखाता है।
दक्षता (Efficiency): बेसिस अपडेट के ओवरहेड के बावजूद, बेसिस रोटेशन ने प्रतिस्पर्धी बेसलाइन्स की तुलना में लक्षित लॉस तक पहुँचने के लिए 54.3% कम GPU घंटे लिए।
स्टेज-अवेयर ऑप्टिमाइज़ेशन: प्रत्येक स्टेज पर विलंब के अनुपात में कम्प्यूटेशनल बजट आवंटित करने से अतिरिक्त 29.2% की स्पीडअप मिली।
सामान्यीकरण (Generalization): यह विधि वेट स्टैशिंग (एक मेमोरी-सेविंग तकनीक जो ग्रेडिएंट एरर पेश करती है) के बिना भी प्रभावी रही और Mixture-of-Experts (MoE) आर्किटेक्चर में भी सामान्यीकृत हुई।

5. महत्व और दावे (Significance and Claims)

पेपर का दावा है कि एसिंक्रोनस ट्रेनिंग में देखी गई प्रदर्शन गिरावट स्टेलेनेस की एक अपरिहार्य लागत नहीं है, बल्कि बेसिस मिसअलाइनमेंट के कारण विलंब के प्रति ऑप्टिमाइज़र की संवेदनशीलता का एक बायप्रोडक्ट है। बेसिस रोटेशन का प्रस्ताव करके, लेखक तर्क देते हैं कि उन्होंने "स्केल पर बबल-फ्री निष्पादन" के एक नए युग को अनलॉक कर दिया है।

यह कार्य सुझाव देता है कि बड़े पैमाने के वितरित प्रशिक्षण के लिए, ऑप्टिमाइज़ेशन लैंडस्केप के ज्यामितीय मिसअलाइनमेंट को संबोधित किए बिना केवल पाइपलाइन डेप्थ बढ़ाना घटते प्रतिफल (diminishing returns) की ओर ले जाता है। बेसिस रोटेशन एक सिद्धांतबद्ध, सैद्धांतिक रूप से ठोस समाधान प्रदान करता है जो एसिंक्रोनस पाइपलाइन पैरेललिज्म को ट्रिलियन-पैरामीटर मॉडल्स के प्रशिक्षण के लिए आवश्यक सैकड़ों चरणों तक कुशलतापूर्वक स्केल करने की अनुमति देता है, जिससे ग्रेडिएंट विलंब के नकारात्मक प्रभाव को प्रभावी ढंग से बेअसर किया जा सकता है।