🤖 machine learning

DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training

DASH (Deterministic Attention Scheduling for High-Throughput), बैकवर्ड पास को एक DAG शेड्यूलिंग समस्या के रूप में स्वरूपित करके और डिसेंडिंग Q-टाइल इटरेशन (Descending Q-Tile Iteration) तथा शिफ्ट शेड्यूलिंग (Shift Scheduling) जैसी नवीन रणनीतियों को पेश करके, LLM ट्रेनिंग में डिटर्मिनिस्टिक अटेंशन के महत्वपूर्ण प्रदर्शन ओवरहेड को संबोधित करता है, जो NVIDIA H800 GPUs पर पाइपलाइन स्टॉल्स को कम करता है और थ्रूपुट को 1.28 $\times$ तक सुधारता है।

मूल लेखक: Xinwei Qiang, Hongmin Chen, Shixuan Sun, Jingwen Leng, Xin Liu, Minyi Guo

प्रकाशित 2026-06-11

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

CC BY 4.0

मूल लेखक: Xinwei Qiang, Hongmin Chen, Shixuan Sun, Jingwen Leng, Xin Liu, Minyi Guo

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

यहाँ "DASH: High-Throughput Reproducible LLM Training के लिए Deterministic Attention Scheduling" पेपर का सरल भाषा और रचनात्मक उपमाओं (analogies) के साथ अनुवाद दिया गया है।

बड़ी समस्या: "Reproducible" (पुनरुत्पादनीयता) की बाधा

कल्पना कीजिए कि आप एक बहुत बड़ी रसोई (GPU) चला रहे हैं जहाँ सैकड़ों शेफ (प्रोसेसिंग यूनिट्स) मिलकर एक विशाल भोजन (Large Language Model को ट्रेन करना) बनाने के लिए काम कर रहे हैं।

AI की दुनिया में, वैज्ञानिकों को यह सक्षम होना चाहिए कि वे ठीक उसी भोजन को दोबारा बना सकें और उन्हें बिल्कुल वैसा ही परिणाम मिले। इसे reproducibility (पुनरुत्पादनीयता) कहा जाता है। यदि वे रेसिपी में थोड़ा सा भी बदलाव करते हैं, तो उन्हें ठीक से पता होना चाहिए कि स्वाद में कैसे बदलाव आया।

हालाँकि, कंप्यूटरों की एक अजीब आदत है: जब वे संख्याओं को आपस में जोड़ते हैं, तो क्रम मायने रखता है। यदि शेफ A बर्तन में नमक डालता है, और फिर शेफ B काली मिर्च डालता है, तो परिणाम थोड़ा अलग होगा बजाय इसके कि शेफ B पहले काली मिर्च डाले और फिर शेफ A नमक डाले। एक अराजक रसोई में जहाँ शेफ बेतरतीब ढंग से ऑर्डर चिल्लाते हैं, हर बार खाना बनाने पर अंतिम स्वाद थोड़ा बदल जाता है। इसे non-determinism (अनिश्चितता) कहते हैं।

इसे ठीक करने के लिए, वर्तमान मानक (FlashAttention-3) शेफ को एक सख्त, पूर्व-निर्धारित क्रम में लाइन में खड़े होने के लिए मजबूर करता है। शेफ 1 जाता है, फिर शेफ 2, फिर शेफ 3। यह गारंटी देता है कि हर बार स्वाद बिल्कुल एक जैसा होगा।

पेंच (The Catch): यह सख्त लाइन-अप धीमा है। जब शेफ 1 नमक डाल रहा होता है, तो शेफ 2 को खाली खड़ा रहना पड़ता है और इंतजार करना पड़ता है। शेफ 3 को और भी लंबा इंतजार करना पड़ता है। रसोई में शेफ खाली खड़े होकर कुछ न करने का नाटक कर रहे हैं, बस अपनी बारी का इंतजार कर रहे हैं। पेपर कहता है कि यह "इंतजार" पूरी ट्रेनिंग प्रक्रिया को लगभग 38% धीमा कर देता है। यह समय और पैसे की बहुत बड़ी बर्बादी है।

समाधान: DASH (Deterministic Attention Scheduling)

लेखकों ने DASH नामक एक नया सिस्टम बनाया है। केवल सबको एक उबाऊ लाइन में खड़ा करने के बजाय, उन्होंने रसोई के वर्कफ़्लो को इस तरह से फिर से डिज़ाइन किया कि शेफ काम जारी रख सकें और साथ ही उस सख्त क्रम का पालन भी कर सकें जो रेसिपी को 'reproducible' बनाने के लिए आवश्यक है।

उन्होंने इस समस्या को एक ट्रैफिक पहेली की तरह माना। कल्पना कीजिए कि शेफ हाईवे पर मर्ज होने की कोशिश कर रहे कारें हैं। पुराना तरीका यह था कि उन्हें एक-एक करके मर्ज कराया जाए, जिससे भारी ट्रैफिक जाम लग जाता था। DASH यह तय करता है कि कारें बिना रुके सुचारू रूप से मर्ज हो सकें।

उन्होंने इसे हल करने के लिए दो मुख्य तरकीबों का उपयोग किया:

तरकीब 1: "रिवर्स लाइन" (Descending Q-Tile Iteration)

कल्पना कीजिए कि एक कमरे में प्रवेश करने के लिए लोगों की लाइन लगी है। आमतौर पर, आप पहले व्यक्ति को अंदर जाने देते हैं, फिर दूसरे को, फिर तीसरे को। लेकिन इस विशेष प्रकार के खाना पकाने (जिसे "Causal Attention" कहा जाता है) में, लाइन के पहले व्यक्ति को वास्तव में पीछे के सभी लोगों द्वारा एक छोटा सा काम पूरा करने तक इंतजार करना पड़ता है। इससे रसोई में एक लंबा, खाली अंतराल बन जाता है।

DASH का समाधान: लाइन को क्रम में (1, 2, 3...) बुलाने के बजाय, वे इसे उल्टे क्रम (3, 2, 1...) में बुलाते हैं।

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

तरकीब 2: "स्टैगर्ड शिफ्ट" (Shift Scheduling)

दूसरे प्रकार के खाना पकाने (जिसे "Full Attention" कहा जाता है) के लिए, समस्या यह है कि हर कोई एक ही समय में एक ही काउंटर का उपयोग करना चाहता है। यदि वे सभी एक साथ अपने मसाले डालने की कोशिश करेंगे, तो वे आपस में टकरा जाएंगे।

DASH का समाधान: वे एक चक्रीय शिफ्ट (cyclic shift) का उपयोग करते हैं। कल्पना कीजिए कि एक रिले रेस (relay race) है जहाँ धावक एक ही समय पर शुरू नहीं होते हैं।

शेफ 1 सामग्री A के साथ शुरू करता है।
शेफ 2 सामग्री B के साथ शुरू करता है (जिसे शेफ 1 बाद में उपयोग करेगा)।
शेफ 3 सामग्री C के साथ शुरू करता है।
जब तक शेफ 1 A के साथ अपना काम पूरा करता है, शेफ 2 उसे सौंपने के लिए तैयार होता है।

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

परिणाम: तेज़, लेकिन जादुई नहीं

लेखकों ने शक्तिशाली NVIDIA H800 GPUs (सुपरकंप्यूटर जिनका उपयोग AI के लिए किया जाता है) पर इनका परीक्षण किया।

जीत: उनके नए सिस्टम ने पुराने, धीमे तरीके की तुलना में "सख्त क्रम" वाले खाना पकाने को 1.28 गुना तेज़ बना दिया। इसने "तेज़ लेकिन अव्यवस्थित" और "धीमा लेकिन सटीक" के बीच के अंतर को कम कर दिया।
वास्तविकता की जाँच: पेपर ने यह भी पाया कि वास्तविक दुनिया में "परफेक्ट" होना हमेशा "सर्वश्रेष्ठ" नहीं होता है।
- कुछ बहुत बड़े, जटिल कार्यों के लिए, "स्टैगर्ड शिफ्ट" (तरकीब 2) वास्तव में पुराने तरीके से थोड़ा धीमा हो गया।
- क्यों? नया तरीका इतना जटिल था कि शेफ (GPU cores) इतने सारे चरणों को याद रखने में घबरा गए। उनके पास "स्क्रैचपैड स्पेस" (registers) कम पड़ गया और उन्हें नोट्स ज़मीन पर (memory में) गिराने पड़े, जिससे उनकी गति धीमी हो गई।
- सबक: कभी-कभी, एक सरल तरकीब (जैसे रिवर्स लाइन) गणितीय रूप से पूर्ण लेकिन जटिल तरकीब से बेहतर होती है, यह इस पर निर्भर करता है कि आपकी रसोई कितनी बड़ी है।

सारांश

यह पेपर DASH पेश करता है, जो AI कंप्यूटर में "शेफ" को व्यवस्थित करने का एक स्मार्ट तरीका है। यह सुनिश्चित करता है कि AI ट्रेनिंग पूरी तरह से 'reproducible' (बिट-दर-बिट समान) हो, बिना कंप्यूटर को खाली बैठने या इंतजार करने के लिए मजबूर किए। ऑपरेशन्स के क्रम को बदलकर—कभी लाइन को उल्टा करके, तो कभी शुरुआत के समय को अलग-अलग करके—उन्होंने प्रक्रिया को काफी तेज कर दिया, जिससे विश्वसनीय AI मॉडल को ट्रेन करना सस्ता और तेज़ हो गया।

तकनीकी सारांश: DASH – उच्च-थ्रूपुट प्रतिलिपि योग्य (Reproducible) LLM प्रशिक्षण के लिए नियतात्मक अटेंशन शेड्यूलिंग (Deterministic Attention Scheduling)

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

लार्ज लैंग्वेज मॉडल्स (LLMs) के प्रशिक्षण में, अस्थिरताओं के निदान और आर्किटेक्चरल परिवर्तनों को सत्यापित करने के लिए प्रतिलिपि योग्यता (reproducibility) आवश्यक है। हालांकि, अटेंशन मैकेनिज्म में बिटवाइज़ समान परिणाम (determinism) प्राप्त करना अक्सर एक गंभीर प्रदर्शन दंड (performance penalty) का कारण बनता है।

मानक अटेंशन कार्यान्वयन (जैसे, FlashAttention-3) में गैर-नियतता (non-determinism) का मूल कारण फ्लोटिंग-पॉइंट अंकगणित की गैर-सहचर्यता (non-associativity) है। शुद्धता सुनिश्चित करने के लिए, बैकवर्ड पास को ग्रेडिएंट टेंसरों (विशेष रूप से $dQ$) के लिए एक निश्चित संचय क्रम (accumulation order) को लागू करना चाहिए। वर्तमान नियतात्मक कार्यान्वयन, जैसे कि FlashAttention-3, ग्रेडिएंट रिडक्शन ऑपरेशन्स को सिंक्रोनाइज़ेशन बैरियर्स का उपयोग करके सीरियलाइज़ करके इसे प्राप्त करते हैं। जबकि यह निरंतरता सुनिश्चित करता है, यह महत्वपूर्ण पाइपलाइन स्टॉल्स (बबल्स) बनाता है जहाँ स्ट्रीमिंग मल्टीप्रोसेसर (SMs) पिछले रिडक्शन के पूरा होने की प्रतीक्षा करते समय खाली रहते हैं। अनुभवजन्य डेटा इंगित करता है कि यह गैर-नियतात्मक समकक्षों की तुलना में थ्रूपुट को 37.9% तक कम कर सकता है, जो मुख्य रूप से टाइल निष्पादन शेड्यूल्स और कठोर संचय क्रम के बीच गलत संरेखण के कारण होता है।

2. कार्यप्रणाली

लेखक DASH (Deterministic Attention Scheduling for High-Throughput) प्रस्तावित करते हैं, जो एक निर्देशित अचक्रीय ग्राफ (Directed Acyclic Graph - DAG) पर एक शेड्यूलिंग अनुकूलन समस्या के रूप में नियतात्मक बैकवर्ड पास को पुनर्गठित करता है।

2.1 समस्या सूत्रीकरण

बैकवर्ड पास को एक DAG के रूप में मॉडल किया गया है जहाँ:

नोड्स (Nodes) गणना चरणों (टाइल प्रोसेसिंग) और रिडक्शन चरणों का प्रतिनिधित्व करते हैं।
एज (Edges) डेटा निर्भरता और अनिवार्य संचय क्रम का प्रतिनिधित्व करते हैं।
उद्देश्य: एंड-टू-एंड विलंबता को कम करने के लिए DAG के क्रिटिकल पाथ (critical path) की लंबाई को न्यूनतम करना।
बाधाएं (Constraints): एक विशिष्ट की-वैल्यू (KV) टाइल के लिए ऑपरेशन्स को फास्ट रजिस्टर-रेसिडेंट संचय का लाभ उठाने के लिए एक ही SM पर निरंतर निष्पादित होना चाहिए, जिससे ग्राफ में अटूट श्रृंखलाएं बनती हैं।

2.2 शेड्यूलिंग रणनीतियाँ

DASH नियतात्मक क्रम का पालन करते हुए पाइपलाइन बबल्स को कम करने के लिए दो पूरक रणनीतियाँ पेश करता है:

A. डिसेंडिंग Q-टाइल इटरेशन (Descending Q-Tile Iteration - ह्यूरिस्टिक)

तंत्र: क्वेरी (Q) ब्लॉक्स के प्रोसेसिंग ऑर्डर को उलट देता है। Q-टाइल्स को $0 $से$ N $तक प्रोसेस करने के बजाय, शेड्यूलर उन्हें$ N $से$ 0$ तक प्रोसेस करता है।
तर्क: कॉज़ल अटेंशन (causal attention) में, निर्भरताएं त्रिकोणीय होती हैं। Q-टाइल्स को रिवर्स ऑर्डर में प्रोसेस करने से SMs निर्भरताओं को जल्दी हल कर सकते हैं, जिससे बाद के अटेंशन हेड्स के लिए संसाधन जल्दी मुक्त हो जाते हैं। यह एक मजबूती से जुड़ा हुआ पाइपलाइन बनाता है जो हेड्स के बीच खाली अंतराल को काफी कम कर देता है।
उपयोगिता: यह फुल और कॉज़ल दोनों मास्क के लिए प्रभावी है, हालांकि इसे मुख्य रूप से कॉज़ल स्टॉल्स को कम करने के लिए डिज़ाइन किया गया है।

B. शिफ्ट शेड्यूलिंग (Shift Scheduling - सैद्धांतिक रूप से इष्टतम)

तंत्र: कंप्यूटेशनल कार्यों को SMs को सौंपने के लिए एक चक्रीय शिफ्ट असाइनमेंट (cyclic shift assignment) का उपयोग करता है।
- फुल मास्क (Full Masks): SMs KV ब्लॉक्स को एक चक्रीय क्रम में प्रोसेस करते हैं (जैसे, SM $i$ प्रोसेस करता है $i, i+1, \dots, N-1, 0, \dots, i-1$ )। यह सुनिश्चित करता है कि किसी भी दिए गए $dQ$ ब्लॉक के लिए, रिडक्शन कार्य स्वाभाविक रूप से बिना किसी संघर्ष के क्रमबद्ध होते हैं, जो इस शर्त को पूरा करता है कि कोई भी निर्भरता एज DAG की डेप्थ मोनोटोनिसिटी (depth monotonicity) का उल्लंघन नहीं करती है (लेम्मा 1)।
- कॉज़ल मास्क (Symmetric Shift): कॉज़ल मास्क के असंतुलित वर्कलोड (जहाँ शुरुआती KV ब्लॉक्स के अधिक इंटरैक्शन होते हैं) को संबोधित करता है। यह SMs को सममित (symmetric) KV ब्लॉक्स ( $i$ और $N-1-i$ ) को संभालने के लिए जोड़ता है, प्रभावी रूप से त्रिकोणीय वर्कलोड को एक संतुलित वर्गाकार में फोल्ड करता है। इसे दो चरणों में निष्पादित किया जाता है: एक सघन लोअर-लेफ्ट रेक्टेंगल और एक फोल्डेड अपर-राइट ट्रायंगल।
सैद्धांतिक आधार: यह लेम्मा पर आधारित है जो कहता है कि एक DAG में शून्य-भार वाली निर्भरता एजेस को जोड़ने से क्रिटिकल पाथ की लंबाई तब सुरक्षित रहती है जब और केवल जब स्रोत नोड की डेप्थ, डेस्टिनेशन नोड की डेप्थ के कम या उसके बराबर हो ( $depth(u) \le depth(v)$ )। शिफ्ट शेड्यूलिंग एक संघर्ष-मुक्त रिडक्शन ऑर्डर का निर्माण करती है जो इस बाधा का पालन करती है।

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

मूल कारण की पहचान: पेपर पहचानता है कि नियतात्मक अटेंशन में प्रदर्शन गिरावट स्वयं सीरियलाइजेशन के कारण नहीं है, बल्कि टाइल शेड्यूलिंग और पूर्व-निर्धारित संचय क्रम के बीच गलत संरेखण के कारण है।
DAG औपचारिकता: यह नियतात्मक अटेंशन बैकवर्ड शेड्यूलिंग समस्या के लिए पहला DAG-आधारित औपचारिकता प्रदान करता है, जो क्रिटिकल पाथ की लंबाई के सिद्धांतपूर्ण अनुकूलन को सक्षम बनाता है।
नवीन शेड्यूलिंग एल्गोरिदम:
- डिसेंडिंग Q-टाइल इटरेशन: क्वेरी ब्लॉक ट्रैवर्सल को रिवर्स करके कॉज़ल अटेंशन प्रदर्शन में सुधार करने वाला एक मजबूत ह्यूरिस्टिक।
- शिफ्ट शेड्यूलिंग: फुल मास्क के लिए वर्कलोड को संतुलित करने और पाइपलाइन बबल्स को समाप्त करने वाला एक सैद्धांतिक रूप से इष्टतम एल्गोरिदम।
- सिमेट्रिक शिफ्ट शेड्यूलिंग: सममित पेयरिंग और टू-फेज फोल्डिंग के माध्यम से असंतुलित वर्कलोड को संतुलित करने के लिए कॉज़ल मास्क के लिए एक विस्तार।
अनुभवजन्य सत्यापन: यह NVIDIA H800 GPUs पर बेसलाइन की तुलना में नियतात्मक अटेंशन के प्रदर्शन अंतर को काफी कम करने की क्षमता प्रदर्शित करता है।

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

प्रयोग NVIDIA H800 GPUs पर FlashAttention-3 को बेसलाइन के रूप में उपयोग करके किए गए थे।

थ्रूपुट सुधार:
- फुल मास्क: शिफ्ट शेड्यूलिंग लगातार बेसलाइन से बेहतर प्रदर्शन करती है, जो अटेंशन बैकवर्ड पास में 1.28× तक की स्पीडअप प्राप्त करती है।
- कॉज़ल मास्क: डिसेंडिंग Q-टाइल इटरेशन और सिमेट्रिक शिफ्ट शेड्यूलिंग सभी कॉन्फ़िगरेशन में सुधार दिखाते हैं।
एंड-टू-एंड प्रदर्शन:
- LLaMA3-8B, Qwen2.5-7B, और Mistral-8x7B जैसे मॉडल्स के लिए पूरे ट्रांसफार्मर ब्लॉक्स (फॉरवर्ड और बैकवर्ड पास सहित) पर लागू होने पर, DASH औसतन ~5% की स्पीडअप प्राप्त करता है (सीक्वेंस लेंथ और मॉडल के आधार पर 2% से 10% के बीच)।
- फुल-मास्क मॉडल्स (जैसे, Stable Diffusion 3.5, SAM) भी ~4% की स्पीडअप देखते हैं।
हार्डवेयर सीमाएँ और ट्रेड-ऑफ:
- पेपर कुछ विशिष्ट परिदृश्यों में सैद्धांतिक इष्टतमता और व्यावहारिक प्रदर्शन के बीच अंतर को नोट करता है। बड़े हेड डायमेंशन (जैसे, 128) वाले कॉज़ल मास्क के लिए, सिमेट्रिक शिफ्ट शेड्यूलिंग, सरल डिसेंडिंग Q-टाइल इटरेशन से कम प्रदर्शन करती है।
- कारण: सिमेट्रिक शिफ्ट शेड्यूलिंग का जटिल स्टेट मैनेजमेंट रजिस्टर प्रेशर (register pressure) बढ़ाता है। वर्तमान आर्किटेक्चर (H800) पर, यह रजिस्टर स्पिलिंग को धीमे लोकल मेमोरी में मजबूर करता है, जिससे एल्गोरिथम के लाभ निष्फल हो जाते हैं। सरल डिसेंडिंग रणनीति इस रजिस्टर प्रेशर से बचती है।
- फुल मास्क के साथ अत्यधिक सीक्वेंस लेंथ (16k) पर, हाई-पैरेललिज्म रेजिम्स में कम्प्यूटेशनल लाभों की तुलना में रिमोट L2 कैश सेगमेंट के माध्यम से बढ़े हुए इंटर-SM कम्युनिकेशन लेटेंसी के कारण शिफ्ट शेड्यूलिंग कभी-कभी थोड़ा घट जाता है।

5. महत्व और दावे

पेपर का दावा है कि DASH, प्रदर्शन से समझौता किए बिना, कुशल और प्रतिलिपि योग्य LLM प्रशिक्षण को महत्वपूर्ण रूप से आगे बढ़ाता है। इंटेलिजेंट DAG-आधारित प्लानिंग के माध्यम से निष्पादन शेड्यूल से संचय क्रम को अलग करके, DASH यह प्रदर्शित करता है कि नियतता (determinism) का प्रदर्शन दंड एक निश्चित लागत नहीं है, बल्कि एक समाधान योग्य शेड्यूलिंग समस्या है।

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