🤖 machine learning

AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training

यह शोध पत्र AdaptiveLoad को प्रस्तुत करता है, जो एक अनुकूलन ढांचा (optimization framework) है जिसमें ड्यूल-कन्स्ट्रेंट लोड बैलेंसिंग और एक फ्यूज्ड LayerNorm-Modulate CUDA कर्नल शामिल है, जो परिवर्तनशील अनुक्रम लंबाई (variable sequence lengths) के कारण होने वाले कम्प्यूटेशनल असंतुलन को संबोधित करके बड़े पैमाने के वीडियो डिफ्यूजन ट्रांसफॉर्मर्स के लिए प्रशिक्षण दक्षता और GPU उपयोगिता को महत्वपूर्ण रूप से बढ़ाता है।

मूल लेखक: Yucheng Guo, Yongjian Guo, Zhong Guan, Haoran Sun, Wen Huang, Wanting Xu, Jing Long, Shuai Di, Junwu Xiong

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

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

CC BY 4.0

मूल लेखक: Yucheng Guo, Yongjian Guo, Zhong Guan, Haoran Sun, Wen Huang, Wanting Xu, Jing Long, Shuai Di, Junwu Xiong

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

यहाँ AdaptiveLoad पेपर का सरल भाषा और रचनात्मक उपमाओं (analogies) के साथ हिंदी अनुवाद दिया गया है।

बड़ी समस्या: रिले रेस में "सबसे धीमा धावक"

कल्पना कीजिए कि आप वीडियो जेनरेट करने के लिए एक सुपर-स्मार्ट AI को प्रशिक्षित (train) कर रहे हैं। यह AI एक रिले रेस में एक साथ काम करने वाली धावकों (GPUs) की एक विशाल टीम की तरह है। एक लैप (ट्रेनिंग स्टेप) को पूरा करने के लिए, टीम को अगला लैप शुरू करने से पहले हर एक धावक को फिनिश लाइन पार करनी ही होगी।

पुराना तरीका ("समान टोकन" का नियम):
पहले, टीम निष्पक्ष होने की कोशिश करती थी और हर धावक को दौड़ने के लिए कदमों की बिल्कुल समान संख्या (टokens) देती थी।

धावक A के पास छोटे, आसान कदम हैं। वह जल्दी खत्म कर लेता है।
धावक B के पास लंबे, भारी कदम हैं। गणित के तरीके (quadratic complexity) के कारण, ये लंबे कदम गणना (compute) करने में बहुत अधिक समय लेते हैं, भले ही कदमों की संख्या समान हो।

परिणाम: धावक A खत्म करता है, और फिर धावक B का इंतज़ार करने के लिए बैठ जाता है। इस प्रतीक्षा समय को "सिंक्रोनाइज़ेशन बबल" (synchronization bubble) कहा जाता है। पुराने सिस्टम में, टीम ने धीमे धावकों का इंतज़ार करने में बहुत समय बर्बाद किया, जिससे शक्तिशाली कंप्यूटर भी खाली बैठे रहे।

समाधान: AdaptiveLoad

लेखक AdaptiveLoad नामक एक नया सिस्टम प्रस्तावित करते हैं। इसे एक स्मार्ट कोच की तरह समझें जो केवल कदमों को नहीं गिनता, बल्कि वास्तव में यह देखता है कि प्रत्येक धावक का भार कितना भारी और कठिन है।

1. स्मार्ट कोच (Dual-Constraint Load Balancing)

हर किसी को समान संख्या में टोकन देने के बजाय, कोच यह तय करने के लिए कि प्रत्येक धावक को कितना काम देना है, दो नियमों का उपयोग करता है:

मेमोरी सीमा (Memory Limit): "इतना वजन मत उठाओ कि तुम उसे गिरा दो।" (यह कंप्यूटर को मेमोरी खत्म होने से रोकता है)।
समय सीमा (Time Limit): "इतना वजन मत उठाओ कि तुम्हें दौड़ने में बहुत लंबा समय लग जाए।" (यह "लॉन्ग-टेल" देरी को रोकता है)।

उपमा (Analogy): एक डिलीवरी ट्रक की कल्पना करें।

पुराना तरीका: हर ट्रक को 100 पैकेज मिलते हैं। यदि पैकेज छोटे बक्से हैं, तो ट्रक तेजी से निकल जाता है। यदि पैकेज विशाल पियानो हैं, तो ट्रक को लोड करने में 10 घंटे लगते हैं। अन्य ट्रक इंतज़ार करते हैं।
AdaptiveLoad: कोच पैकेजों को देखता है। यदि किसी ट्रक को पियानो ले जाने हैं, तो कोच उसे केवल 10 पियानो देगा ताकि वह छोटे बक्सों वाले ट्रकों के साथ ही निकल सके। यदि किसी ट्रक के पास छोटे बक्से हैं, तो उसे 100 मिलेंगे।
परिणाम: हर कोई लगभग एक ही समय पर निकलता है। कोई भी इंतज़ार नहीं करता। पेपर का दावा है कि इसने "प्रतीक्षा समय" के असंतुलन को लगभग आधा कर दिया।

2. सुपर-टूल (Fused CUDA Kernels)

जबकि कोच धावकों को व्यवस्थित करता है, पेपर ने उन उपकरणों को भी ठीक किया जिनका धावक उपयोग करते हैं।

समस्या:
पुराने सिस्टम में, AI को एक कार्य (जैसे वीडियो के फोकस को एडजस्ट करना) कई छोटे, अलग-अलग चरणों में करना पड़ता था।

चरण 1: टूल लेने के लिए गोदाम (मेमोरी) जाना।
चरण 2: दूसरा टूल लेने के लिए वापस गोदाम जाना।
चरण 3: फिर से वापस जाना।
यह एक शेफ की तरह है जो हर एक सामग्री के लिए फ्रिज तक बार-बार दौड़ता है। यह धीमा है और ऊर्जा बर्बाद करता है।

सुधार (Fused Kernel):
लेखकों ने एक "सुपर-टूल" बनाया जो सभी चरणों को एक साथ करता है।

उपमा: 10 बार फ्रिज तक दौड़ने के बजाय, शेफ एक ट्रे में सारी सामग्रियां एक साथ लेता है, पूरा व्यंजन बनाता है, और उसे प्लेट पर परोस देता है।
"D-tile" ट्रिक: पेपर में "D-tile coalesced reduction" नामक एक विशिष्ट ट्रिक का उल्लेख है। कल्पना कीजिए कि शेफ सामग्रियों को अपनी ट्रे पर इस तरह व्यवस्थित करता है कि उन्हें सबसे तेज़ तरीके से पकड़ा जा सके। यह मेमोरी एक्सेस को बेहद सुचारू और तेज़ बनाता है।

परिणाम: क्या हुआ?

जब उन्होंने एक वास्तविक दुनिया के वीडियो AI मॉडल (जिसे Wan 2.1 कहा जाता है) पर इस नए सिस्टम का परीक्षण किया:

कम इंतज़ार: "कंप्यूटेशनल असंतुलन" (कि कैसे सबसे धीमे धावक ने सभी को धीमा किया) 39% से घटकर 18.9% रह गया।
अधिक गति: पूरी टीम ट्रेनिंग में 27.2% तेज़ हो गई।
बेहतर मेमोरी उपयोग: वे बिना क्रैश हुए कंप्यूटर की मेमोरी में अधिक जटिल वीडियो फिट कर सके, जिससे समान हार्डवेयर से अधिक प्रदर्शन निकाला जा सका।
वही गुणवत्ता: AI ने पहले की तरह ही अच्छी तरह से सीखा। "स्मार्ट कोच" ने प्रशिक्षण की गुणवत्ता को नहीं बदला, केवल गति और दक्षता को बेहतर बनाया।

सारांश

AdaptiveLoad एक कठोर असेंबली लाइन (जहाँ कठिनाई की परवाह किए बिना हर कोई समान काम करता है) को एक गतिशील, स्मार्ट सिस्टम में अपग्रेड करने जैसा है। यह वर्कलोड को इस तरह संतुलित करता है कि कोई भी मशीन खाली न बैठे, और यह छोटे कार्यों को बड़े, कुशल बर्स्ट (bursts) में जोड़ देता है ताकि मशीनें पुर्जे लाने में समय बर्बाद न करें। परिणाम, AI को वीडियो बनाना सिखाने का एक बहुत तेज़ और अधिक कुशल तरीका है।

तकनीकी सारांश: AdaptiveLoad

1. समस्या विवरण

विश्व मॉडल (world models) के लिए बड़े पैमाने पर वीडियो डिफ्यूजन ट्रांसफॉर्मर (जैसे DiT और MMDiT) के प्रशिक्षण में, अत्यधिक विषम अनुक्रम लंबाई (sequence lengths) वाले मिश्रित-मोड डेटासेट (स्थिर छवियों से लेकर लंबी अवधि के वीडियो तक) को संसाधित करते समय दो महत्वपूर्ण बाधाएं आती हैं:

क्वाड्रेटिक जटिलता के कारण कम्प्यूटेशनल असंतुलन: मौजूदा वितरित प्रशिक्षण पाइपलाइनें "समान-टोकन" बकेटिंग रणनीतियों (जहाँ $B \times S = \text{constant}$ होता है) पर निर्भर करती हैं। यह दृष्टिकोण सेल्फ-अटेंशन तंत्र की $O(S^2)$ कम्प्यूटेशनल जटिलता को ध्यान में रखने में विफल रहता है। परिणामस्वरूप, लंबे अनुक्रमों वाले बकेट "कम्प्यूटेशनल स्ट्रैगलर" (धीमे चलने वाले घटक) बन जाते हैं, जिससे गंभीर लोड असंतुलन पैदा होता है। ग्लोबल सिंक्रोनाइज़ेशन का समय सबसे धीमे GPU द्वारा निर्धारित होता है, जिससे अन्य नोड्स के लिए महत्वपूर्ण 'इडल टाइम' (सिंक्रोनाइज़ेशन बबल्स) उत्पन्न होता है। प्रयोगात्मक डेटा इंगित करता है कि स्टेप लेटेंसी और कुल टोकन संख्या के बीच एक कमजोर सहसंबंध ( $R \approx 0.35$ ) है, लेकिन अनुक्रम लंबाई के वर्ग के साथ एक मजबूत सहसंबंध ( $R \approx 0.92$ ) है।
मेमोरी वॉल और ऑपरेटर ओवरहेड: मल्टीमॉडल DiT (MMDiT) आर्किटेक्चर में, प्रति इटरेशन सैकड़ों बार एडाप्टिव लेयर नॉर्मलाइज़ेशन (AdaLN) को कॉल किया जाता है। मानक कार्यान्वयन डिस्क्रीट CUDA कर्नेल का उपयोग करते हैं, जिससे स्ट्रीमिंग मल्टीप्रोसेसरों (SMs) और हाई बैंडविड्थ मेमोरी (HBM) के बीच इंटरमीडिएट टेंसरों के बार-बार और अनावश्यक राउंड-ट्रिप्स होते हैं। इसके अलावा, बैकप्रोपैगेशन के दौरान, मॉड्यूलेशन मापदंडों के लिए ग्रेडिएंट रिडक्शन, गैर-सतत (non-contiguous) एक्सेस पैटर्न के कारण अक्षम होता है, जो हार्डवेयर-स्तर के कोलेसड एक्सेस (coalesced access) को रोकता है और मेमोरी बैंडविड्थ का कम उपयोग करता है।

2. कार्यप्रणाली: AdaptiveLoad फ्रेमवर्क

यह पेपर एक पूर्ण-स्टैक अनुकूलन फ्रेमवर्क प्रस्तावित करता है जो उच्च-स्तरीय शेड्यूलिंग को निम्न-स्तरीय कर्नेल फ्यूजन के साथ एकीकृत करता है।

A. डुअल-कन्स्ट्रेंट एडेप्टिव लोड बैलेंसिंग

कम्प्यूटेशनल असंतुलन को संबोधित करने के लिए, लेखक रैखिक "समान-टोकन" धारणा के स्थान पर एक डुअल-कन्स्ट्रेंट शेड्यूलिंग लॉजिक का प्रस्ताव करते हैं जो साथ-साथ निम्नलिखित को लागू करता है:

रैखिक मेमोरी बाधा (Linear Memory Constraint): $B \times S \leq M_{mem}$ (यह सुनिश्चित करना कि बैच साइज GPU मेमोरी में फिट हो सके)।
पॉलीनोमियल कम्प्यूटेशनल बाधा (Polynomial Computational Constraint): $B \times S^p \leq M_{comp}$ ( $O(S^2)$ जटिलता को ध्यान में रखते हुए)।

एक बकेट के लिए अंतिम बैच आकार इन बाधाओं के प्रतिच्छेदन (intersection) द्वारा निर्धारित किया जाता है:
$B_{shape} = \max(1, \min(\lfloor M_{mem}/S \rfloor, \lfloor M_{comp}/S^p \rfloor))$

कॉस्ट फिटिंग सिस्टम: प्रायोगिक घातांक $p$ और कम्प्यूटेशनल बाउंड $M_{comp}$ को निर्धारित करने के लिए, सिस्टम एक डेटा-संचालित शेप बेंचमार्किंग (Shape Benchmarking) वर्कफ़्लो का उपयोग करता है। यह लाइव डिस्ट्रिब्यूटेड एनवायरनमेंट (जैसे FSDP पाथ) में सिंथेटिक पिक्सेल स्कैन निष्पादित करता है ताकि बैच साइज और अनुक्रम लंबाई के विरुद्ध स्टेप लेटेंसी को मापा जा सके। नॉन-लीनियर रिग्रेशन का उपयोग करके, यह बकेट कॉन्फ़िगरेशन को गतिशील रूप से कैलिब्रेट करने के लिए एक कॉस्ट मॉडल ( $\text{step time} \approx a + b \times B \times S^p$ ) को फिट करता है, जिससे मैन्युअल अनुभवजन्य ट्यूनिंग की आवश्यकता समाप्त हो जाती है।

B. फ्यूज्ड लेयरनॉर्म-मॉड्यूलेट CUDA कर्नेल्स

"मेमोरी वॉल" को हल करने के लिए, लेखक एक फ्यूज्ड AdaLN CUDA कर्नेल डिज़ाइन करते हैं जो नॉर्मलाइजेशन और मॉड्यूलेशन ऑपरेशन्स को एक ही कर्नेल लॉन्च में संयोजित करता है।

फॉरवर्ड पास: कर्नेल इंटरमीडिएट स्टैटिस्टिक्स (मीन और वेरिएंस) को स्टोर करने के लिए रजिस्टर्स और शेयर्ड मेमोरी का उपयोग करता है, और टू-स्टेज रिडक्शन (वार्प-लेवल शफल और ब्लॉक-लेवल रिडक्शन) करता है। फीचर डायमेंशन $D$ के तीसरे ट्रावर्सल के दौरान लीनियर ट्रांसफॉर्मेशन सीधे लागू किए जाते हैं, जिससे अनावश्यक HBM रीड/राइट्स समाप्त हो जाते हैं।
बैकवर्ड पास (D-tile कोलेसल्ड रिडक्शन): ग्रेडिएंट रिडक्शन में स्ट्राइड एक्सेस बॉटलनेक को हल करने के लिए, लेखक एक D-tile कोलेसल्ड रिडक्शन रणनीति पेश करते हैं। अनुक्रम डायमेंशन $N$ के साथ रिडक्शन करने के बजाय (जो स्ट्राइड एक्सेस का कारण बनता है), ग्रिड को $(d_{tile}, n_{tile})$ लेआउट में विभाजित किया जाता है। थ्रेड्स को एक फीचर इंडेक्स $d$ के लिए फिक्स किया जाता है और वे अनुक्रम टाइल्स के साथ वर्टिकल एक्युमुलेशन करते हैं। यह सुनिश्चित करता है कि एक वार्प के भीतर मेमोरी एक्सेस निरंतर (contiguous) है, जिससे कोलेसल्ड HBM ट्रांजेक्शन सक्षम होते हैं और बैंडविड्थ उपयोग अधिकतम होता है।
ग्राफ सरलीकरण: ऑपरेटर चेन को फ्यूज करके, कम्प्यूटेशनल ग्राफ को 5-8 डिस्क्रीट नोड्स से एक एकल एटोमिक नोड में समेटा जाता है। इंटरमीडिएट स्टैटिस्टिक्स रजिस्टर्स में रहते हैं, जिससे फ्रेमवर्क के ग्लोबल मेमोरी मैनेजमेंट को बायपास किया जाता है और एक्टिवेशन मेमोरी फुटप्रिंट काफी कम हो जाता है।

3. मुख्य योगदान

लोड मिसमैच का परिमाणीकरण और शमन: यह पेपर वीडियो डिफ्यूजन ट्रेनिंग में लोड मिसमैच को सटीक रूप से मापता है और एक डुअल-कन्स्ट्रेंट एडेप्टिव बकेट शेड्यूलिंग एल्गोरिदम प्रस्तावित करता है। यह दृष्टिकोण लंबे अनुक्रमों की $O(S^2)$ जटिलता को सक्रिय रूप से प्रबंधित करके लोड असंतुलन दर को 50% तक कम कर देता है।
डेटा-संचालित कॉस्ट मॉडलिंग: लेखक एक एंड-टू-एंड सिस्टम डिज़ाइन करते हैं जो वितरित प्रशिक्षण कॉन्फ़िगरेशन को ह्यूरिस्टिक-संचालित से डेटा-संचालित में बदल देता है। शेप बेंचमार्किंग के माध्यम से स्वचालित कॉस्ट-मॉडल फिटिंग के माध्यम से, सिस्टम बिना किसी मैन्युअल हस्तक्षेप के क्लस्टर-वाइड रिसोर्स यूटिलाइजेशन को अनुकूलित करता है।
उच्च-प्रदर्शन फ्यूज्ड कर्नेल्स: MMDiT बैकबोन के लिए एक विशेष फ्यूज्ड AdaLN CUDA कर्नेल विकसित किया गया है। D-tile कोलेसल्ड रिडक्शन पेश करके, यह बैकप्रोपैगेशन में स्ट्राइड एक्सेस बॉटलनेक को हल करता है, जिससे प्रशिक्षण थ्रूपुट में 25% से अधिक की वृद्धि होती है और साथ ही एक्टिवेशन मेमोरी उपयोग भी कम होता है।

4. प्रयोगात्मक परिणाम

प्रयोग Wan 2.1 वर्ल्ड मॉडल पर 8-GPU और 16-GPU क्लस्टरों में 10 मिलियन सैंपल्स (WebDataset और Koala-36m) के मिश्रित कॉर्पस का उपयोग करके किए गए।

थ्रूपुट सुधार:
- 8-GPU क्लस्टर: औसत प्रशिक्षण थ्रूपुट 14,383 से बढ़कर 18,069 टोकन/सेकंड (+25.6%) हो गया।
- 16-GPU क्लस्टर: औसत थ्रूपुट 30,170 से बढ़कर 38,372 टोकन/सेकंड (+27.18%) हो गया।
- AdaptiveLoad कर्व ने बेसलाइन की तुलना में उच्च थ्रूपुट फ्लोर बनाए रखा, यहाँ तक कि वर्स्ट-केस सीक्वेंस डिस्ट्रीब्यूशन में भी।
लोड बैलेंसिंग दक्षता:
- कम्प्यूटेशनल असंतुलन दर (कम्प्यूट लोड के कोएफिशिएंट ऑफ वेरिएशन द्वारा मापा गया) 39% से घटकर 18.9% हो गई।
- सिंक्रोनाइज़ेशन ओवरहेड (स्टेप टाइम का CV) काफी कम हो गया, जिसमें 16-GPU कॉन्फ़िगरेशन में AdaptiveLoad के तहत बेसलाइन मीन 18.7% से घटकर 10.4% रह गया।
मेमोरी और कर्नेल प्रदर्शन:
- पीक VRAM यूटिलाइजेशन: 22.7% सुधरा।
- कर्नेल स्पीडअप: फ्यूज्ड फॉरवर्ड पास ने 3.21× से 3.39× का स्पीडअप प्राप्त किया। बैकवर्ड पास (D-tile रिडक्शन) ने 64k सीक्वेंस लेंथ पर 1.42× तक का स्पीडअप प्राप्त किया।
- मेमोरी बचत: फ्यूज्ड कर्नेल ने एक्टिवेशन मेमोरी को लगभग 61.9% कम कर दिया (उदाहरण के लिए, $N=64,000$ पर ~4.8 GB की बचत)।
- लोड क्षमता: अनुकूलित सिस्टम ने पीक मेमोरी उपयोग बढ़ाए बिना अधिकतम अनुक्रम लंबाई 52,800 टोकन (बेसलाइन की 48,000 टोकन की सीमा से 10% अधिक) का समर्थन किया।
प्रशिक्षण स्थिरता: AdaptiveLoad के लिए लॉस कर्व्स बेसलाइन के साथ अत्यधिक सुसंगत रहे, जो यह दर्शाता है कि अनुकूलन ने संख्यात्मक शुद्धता (numerical fidelity) और अभिसरण गुणवत्ता को बनाए रखा।

5. महत्व

यह पेपर दावा करता है कि AdaptiveLoad अति-लंबे अनुक्रम वाले वीडियो जनरेशन मॉडल्स को स्केल करने के लिए एक मजबूत आधार प्रदान करता है। डुअल-कन्स्ट्रेंट शेड्यूलिंग के माध्यम से मेमोरी सुरक्षा को कम्प्यूटेशनल समानता के साथ जोड़कर और हार्डवेयर-अवेयर ऑपरेटर फ्यूजन के माध्यम से "मेमोरी वॉल" को संबोधित करके, यह फ्रेमवर्क विषम वीडियो डेटासेट में अंतर्निहित सिंक्रोनाइज़ेशन बाधाओं को प्रभावी ढंग से समाप्त करता है। परिणाम दर्शाते हैं कि रैखिक टोकन-काउंटिंग ह्यूरिस्टिक्स से आगे बढ़कर क्वाड्रेटिक जटिलता-जागरूक शेड्यूलिंग को अपनाना, वीडियो जनरेशन कार्यों के लिए बड़े पैमाने पर GPU क्लस्टरों की दक्षता को अधिकतम करने के लिए आवश्यक है।