TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference
TokenWeave एक नवीन प्रणाली है जो विशेष GPU फीचर्स का उपयोग करके RMSNorm ऑपरेशन को AllReduce कम्युनिकेशन के साथ फ्यूज करके, छोटे बैच साइज पर वितरित (distributed) LLM इन्फरेंस के लिए कुशल कंप्यूट-कम्युनिकेशन ओवरलैप प्राप्त करती है, जिससे 1024 टोकन प्रति इटरेशन जैसे कम स्तर पर भी लेटेंसी को कम करती है और थ्रूपुट को बढ़ाती है।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक विशाल, हाई-स्पीड फैक्ट्री (एक लार्ज लैंग्वेज मॉडल) चला रहे हैं जो ग्राहकों के अनुरोधों को प्रोसेस करती है। इस फैक्ट्री को पर्याप्त तेज़ बनाने के लिए, आपने काम करने के लिए 8 विशेषज्ञ श्रमिकों (GPUs) की एक टीम रखी है। वे एक साथ काम करने के लिए एक सुपर-फास्ट कन्वेयर बेल्ट सिस्टम (NVLink) से जुड़े हुए हैं ताकि वे अपना काम तुरंत साझा कर सकें।
हालाँकि, एक समस्या है: श्रमिक एक-दूसरे से बात करने के लिए बहुत अधिक समय बिताते हैं।
इस सुपर-फास्ट कन्वेयर बेल्ट के बावजूद, श्रमिकों को आपस में नोट्स (संदेश) भेजने के लिए अपना वास्तविक निर्माण कार्य (कंप्यूटेशन) रोकना पड़ता है। पेपर में, लेखकों ने पाया कि बड़े मॉडल्स के लिए, यह "बात करने का समय" कुल समय का लगभग 20% खा जाता है। यह एक शेफ की तरह है जो अगला सब्ज़ी काटने से पहले दूसरे शेफ को कॉल करके पूछने के लिए रुक जाता है, "क्या तुम्हारे पास नमक है?"
पुराना तरीका: चीजों को तोड़ना
इसे ठीक करने के पिछले प्रयासों ने काम को छोटे-छोटे टुकड़ों में तोड़ने की कोशिश की। विचार यह था: "जब श्रमिक A, श्रमिक B को एक नोट भेज रहा हो, तो श्रमिक A अगला सब्ज़ी काट सकता है।"
लेकिन लेखकों ने पाया कि यह छोटे ऑर्डर्स (जैसा कि तब होता है जब आप AI से छोटा सवाल पूछते हैं) के लिए अच्छी तरह काम नहीं करता था। काम को बहुत छोटे टुकड़ों में तोड़ने से वास्तव में श्रमिक धीमे हो जाते थे क्योंकि उन्हें बहुत बार रुकना और शुरू करना पड़ता था। यह एक रिले रेस की तरह है जहाँ आप हर 10 फीट पर बैटन (बैटन) पास करते हैं; दौड़ने में लगने वाला समय बैटन हाथ में देने के लिए रुकने में लगने वाले समय से कम है!
नया समाधान: TokenWeave
लेखकों ने TokenWeave नामक एक नया सिस्टम बनाया है। इसे एक स्मार्ट मैनेजर के रूप में सोचें जो प्रतीक्षा समय को खत्म करने के लिए फैक्ट्री के फर्श को पुनर्गठित करता है। उन्होंने इसे तीन सरल तरीकों से किया है:
1. "स्मार्ट स्प्लिट" (दो-लेन वाला हाईवे)
काम को लाखों छोटे टुकड़ों में तोड़ने के बजाय, TokenWeave ऑर्डर को केवल दो बड़े हिस्सों में विभाजित करता है।
- चंक A (Chunk A) काम के पहले आधे हिस्से पर काम करना शुरू करता है।
- चंक B (Chunk B) काम के दूसरे आधे हिस्से पर काम करना शुरू करता है।
- जादू: जब चंक A अपने गणित (math) में व्यस्त होता है, तब चंक B अपने नोट्स पास करने में व्यस्त होता है। फिर, वे स्विच करते हैं। चंक A नोट्स पास करता है जबकि चंक B गणित करता है।
- यह क्यों काम करता है: लेखकों ने ठीक से पता लगाया कि काम को कैसे विभाजित किया जाए ताकि श्रमिक कन्वेयर बेल्ट के लिए "अटक" न जाएं। वे इसे "वेव-अवेयर" (wave-aware) कहते हैं, जिसका अर्थ है कि वे सुनिश्चित करते हैं कि श्रमिक हमेशा व्यस्त रहें, बिल्कुल एक अच्छी तरह से टाइम किए गए ट्रैफिक लाइट सिस्टम की तरह जो कारों को बिना रुके चलते रहने में मदद करता है।
2. "फ्यूज्ड कर्नेल" (ऑल-इन-वन टूल)
पुराने फैक्ट्री में, श्रमिकों को दो अलग-अलग काम करने पड़ते थे:
- नोट्स पास करना (कम्युनिकेशन)।
- डेटा को नॉर्मलाइज़ करना (एक गणितीय चरण जिसे RMSNorm कहा जाता है)।
लेखकों ने महसूस किया कि इन दोनों चरणों को अलग-अलग करना बर्बादी थी। यह ऐसा है जैसे हथौड़ा चलाने के लिए सप्लाई क्लोजेट तक जाना, फिर वापस वर्कबेंच पर आना, फिर पेचकस लेने के लिए वापस क्लोजेट तक जाना।
- समाधान: उन्होंने एक नया "सुपर-टूल" (एक फ्यूज्ड कर्नल) बनाया जो नोट-पासिंग और गणित के चरण को एक साथ करता है।
- बोनस: यह सुपर-टूल इतना कुशल है कि इसे चलाने के लिए फैक्ट्री की कुल शक्ति का एक बहुत छोटा हिस्सा (केवल 2-8 श्रमिक 132 में से) चाहिए। इससे बाकी श्रमिक भारी काम (कंप्यूटेशन) पर पूरी तरह ध्यान केंद्रित करने के लिए स्वतंत्र रहते हैं।
3. "स्मार्ट रीऑर्डरिंग" (सही क्रम में करना)
आमतौर पर, फैक्ट्री पहले सभी नोट्स पास करती है, फिर गणित करती है। लेकिन लेखकों ने महसूस किया कि यदि वे चरणों को पुनर्व्यवस्थित करते हैं, तो नोट-पासिंग प्रक्रिया के दौरान ही गणित का चरण (RMSNorm) किया जा सकता है।
- एनालॉजी (उपमा): पूरे ट्रक के आने का इंतज़ार करने के बजाय कि आप अनलोडिंग शुरू करें, आप पहला बॉक्स देखते ही अनलोडिंग शुरू कर देते हैं जैसे ही ट्रक आता है। TokenWeave चरणों को इस तरह से पुनर्व्यवस्थित करता है कि डेटा अभी भी मूव हो रहा होता है तभी गणित होता है, जिससे बहुत सारा समय बचता है।
परिणाम
लेखकों ने इस नए सिस्टम का परीक्षण शक्तिशाली कंप्यूटरों (8x H100 GPUs) पर Llama और Qwen जैसे वास्तविक दुनिया के मॉडल्स के साथ किया।
- गति (Speed): उन्होंने पाया कि TokenWeave ने मौजूदा सर्वोत्तम सिस्टमों की तुलना में फैक्ट्री को 1.28 गुना तेज़ (28% की स्पीडअप) बना दिया।
- छोटे ऑर्डर्स: बहुत छोटे सवालों (केवल 1,000 शब्द) के लिए भी, यह 1.2 गुना तेज़ था। पिछले सिस्टम छोटे ऑर्डर्स के साथ वास्तव में धीमे हो जाते थे।
- थ्रूपुट (Throughput): फैक्ट्री प्रति घंटे 19% अधिक ग्राहकों को संभाल सकती है।
- "मैजिक" दावा: कुछ मामलों में, TokenWeave इतना कुशल था कि इसने उस सैद्धांतिक फैक्ट्री के मुकाबले बेहतर प्रदर्शन किया जिसमें शून्य कम्युनिकेशन (Zero Communication) था। ऐसा इसलिए क्योंकि उनके नए "सुपर-टूल" ने गणित के चरण को इतना बेहतर बना दिया कि उसने बातचीत में लगने वाले समय की भरपाई कर दी।
सारांश
TokenWeave एक ऑर्केस्ट्रा के मास्टर कंडक्टर की तरह है। संगीतकारों को एक-दूसरे से बात करने (जिससे संगीत धीमा हो जाता है) देने के बजाय, यह उन्हें सिखाता है कि कंडक्टर द्वारा शीट म्यूजिक पास किए जाने के दौरान ही अपना हिस्सा कैसे बजाना है। काम को केवल दो स्मार्ट चंक्स में विभाजित करके और एक नए "ऑल-इन-वन" टूल का उपयोग करके, उन्होंने प्रतीक्षा समय को समाप्त कर दिया, जिससे AI इन्फरेंस काफी तेज़ और अधिक कुशल हो गया।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।