DMuon: Efficient Distributed Muon Training with Near-Adam Overhead
यह शोध पत्र DMuon को प्रस्तुत करता है, जो Muon ऑप्टिमाइज़र का एक ओपन-सोर्स वितरित कार्यान्वयन है जो मौजूदा प्रशिक्षण पाइपलाइनों में निर्बाध रूप से एकीकृत होता है ताकि ऑप्टिमाइज़र-स्टेप लेटेंसी को नाटकीय रूप से कम किया जा सके और AdamW के करीब दक्षता प्राप्त की जा सके, जिससे बड़े फाउंडेशन मॉडल्स में मैट्रिक्स-ऑर्थोगोनालाइजेशन-आधारित अनुकूलन के स्केलेबल उपयोग को सक्षम बनाया जा सके।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक विशाल रोबोटिक मस्तिष्क (एक लार्ज लैंग्वेज मॉडल या एम्बॉडीड एआई) को प्रशिक्षित कर रहे हैं। इसे सिखाने के लिए, आपको एक "ऑप्टिमाइज़र" की आवश्यकता है—एक कोच जो रोबोट की गलतियों को देखता है और इसे स्मार्ट बनाने के लिए इसके आंतरिक कनेक्शनों (वेट्स) को समायोजित करता है।
वर्षों से, मानक कोच AdamW रहा है। यह एक हजारों नन्हे श्रमिकों की टीम की तरह है, जिनमें से प्रत्येक एक समय में रोबट के एक छोटे से पेंच (स्क्रू) को ठीक करता है। यह तेज़ और कुशल है, लेकिन यह हर पेंच को स्वतंत्र रूप से देखता है।
हाल ही में, Muon नामक एक नया, अधिक स्मार्ट कोच खोजा गया है। पेंचों को एक-एक करके ठीक करने के बजाय, Muon पूरे पैनलों (डेटा के पूरे मैट्रिक्स) को देखता है और उन सभी को एक साथ समायोजित करता है। यह "ग्रुप थेरेपी" वाला दृष्टिकोण रोबोट को तेज़ी से सीखने और उच्च प्रदर्शन तक पहुँचने में मदद करता है। हालाँकि, इसमें एक पेच है: यह एक डिस्ट्रीब्यूटेड सिस्टम (कई कंप्यूटरों के समूह) पर चलाने में अविश्वसनीय रूप से धीमा है।
समस्या: "ऑल-हैंड्स" मीटिंग का बॉटलनेक (अवरोध)
एक डिस्ट्रीब्यूटेड ट्रेनिंग सेटअप में, रोबोट का मस्तिष्क कई कंप्यूटरों (GPUs) में विभाजित होता है।
- AdamW एक रिमोट टीम की तरह काम करता है: कंप्यूटर A अपने पेंच ठीक करता है, कंप्यूटर B अपने पचे ठीक करता है, और उन्हें एक-दूसरे से बहुत अधिक बात करने की आवश्यकता नहीं होती है।
- Muon एक कमेटी की तरह काम करता है: एक पैनल को ठीक करने के लिए, हर कंप्यूटर को पहले पूरा पैनल देखना होगा।
Muon के एक साधारण कार्यान्वयन (naive implementation) में, हर एक कंप्यूटर को अपना काम रोकना पड़ता है, सभी से पूरा डेटा डाउनलोड करना पड़ता है, जटिल गणित करना पड़ता है, और फिर परिणाम वापस भेजना पड़ता है। यह एक विशाल बैठक आयोजित करने जैसा है जहाँ हर कोई एक निर्णय लेने से पहले पूरी 1,000 पन्नों की रिपोर्ट पढ़ता है। यह इतना लंबा समय लेता है कि "बैठक" (ऑप्टिमाइज़र स्टेप) वास्तविक कार्य (लर्निंग स्टेप्स) से भी अधिक समय ले लेती है। वास्तव में, पेपर नोट करता है कि इसमें वास्तविक काम से 2 गुना अधिक समय लग सकता है!
समाधान: DMuon (डिस्ट्रीब्यूटेड Muon)
X Square Robot Team ने DMuon बनाया, एक नया सिस्टम जो Muon को AdamW के लगभग उतना ही तेज़ बनाता है। उन्होंने यह समाधान इसलिए निकाला क्योंकि उन्होंने यह नहीं बदला कि कोच कैसे काम करता है, बल्कि यह बदला कि टीम कैसे काम करती है।
उन्होंने इन तीन मुख्य तरीकों का उपयोग करके बॉटलनेक को हल किया:
1. "नामित विशेषज्ञ" रणनीति (ओनर-सेंट्रिक एक्जीक्यूशन)
हर किसी द्वारा हर पैनल को ठीक करने के बजाय, DMuon प्रत्येक पैनल के लिए एक विशिष्ट कंप्यूटर को "ओनर" (स्वामी) के रूप में असाइन करता है।
- पहले: सबने डेटा इकट्ठा किया, सबने गणित किया, सबने समय बर्बाद किया।
- अब: केवल "ओनर" कंप्यूटर डेटा एकत्र करता है और भारी गणित करता है। अन्य कंप्यूटर केवल प्रतीक्षा करते हैं या अन्य कार्य करते हैं।
- उपमा: एक निर्माण स्थल की कल्पना करें। इसके बजाय कि हर मजदूर पूरा घर बनाने की कोशिश करे, एक विशेषज्ञ को छत के लिए, दूसरा प्लंबिंग के लिए सौंपा जाता है। अन्य लोग छत बनाने की कोशिश नहीं करते; वे बस विशेषज्ञ को वे उपकरण देते हैं जिनकी उन्हें आवश्यकता होती है। यह दूसरों को एक ही काम दो बार करने से रोकता है।
2. "कन्वेयर बेल्ट" सिस्टम (कम्युनिकेशन ओवरलैप)
जब तक "ओनर" गणित कर रहा होता है, अन्य कंप्यूटर केवल खाली नहीं बैठे रहते। DMuon एक पाइपलाइन बनाता है:
- फॉरवर्ड पास: जब रोबोट "सोच" (डेटा प्रोसेस) रहा होता है, तो सिस्टम बैकग्राउंड में विशेषज्ञों को अगला सेट टूल्स चुपचाप भेज देता है।
- बैकवर्ड पास: जब रोबोट "गलतियों से सीख" रहा होता है, तो सिस्टम विशेषज्ञों के लिए अगला डेटा बैच पहले से ही इकट्ठा कर रहा होता है।
- उपमा: यह एक रेस्टोरेंट किचन की तरह है। शेफ (ओनर) सब्जियां काट रहा है। वेटर द्वारा अगला ऑर्डर लाने का इंतज़ार करने के बजाय, वेटर शेफ के काटने खत्म करने से पहले ही अगला ऑर्डर ले आता है। प्रतीक्षा समय को काटने के समय के भीतर ही छिपा दिया जाता है।
3. "स्मार्ट असेंबली लाइन" (बैचिंग और ट्यूनिंग)
Muon जो गणित करता है वह जटिल है। कभी-कभी टुकड़े बहुत बड़े होते हैं (जैसे एक विशाल दीवार), और कभी-कभी बहुत छोटे (जैसे एक छोटा टाइल)।
- समस्या: यदि आप एक विशाल मशीन पर एक छोटा सा टाइल प्रोसेस करने की कोशिश करते हैं, तो मशीन खाली बैठी रहती है। यदि आप एक विशाल दीवार प्रोसेस करते हैं, तो इसमें बहुत समय लगता है।
- समाधान: DMuon कई छोटे टाइलों को एक एकल "बैच" में समूहबद्ध करता है ताकि मशीन व्यस्त रहे। यह एक "ट्यूनर" का भी उपयोग करता है जो स्वचालित रूप से प्रत्येक विशिष्ट आकार के टाइल को काटने का सबसे तेज़ तरीका खोज लेता है।
- उपमा: एक अकेले लिफाफे के लिए डिलीवरी ट्रक भेजने के बजाय, DMuon तब तक इंतजार करता है जब तक कि उसके पास लिफाफों का एक पूरा ट्रक लोड न हो जाए और फिर वह उन सभी को एक साथ भेजता है। यह गंतव्य शहर या खेत के आधार पर ट्रक का मार्ग भी बदल देता है।
परिणाम
पेपर ने वास्तविक दुनिया के मॉडल्स पर DMuon का परीक्षण किया, जिसमें एक रोबोट ट्रेनिंग मॉडल (Wall-OSS) और एक लैंग्वेज मॉडल (Qwen2.5) शामिल हैं।
- गति: DMuon ने पुराने, साधारण तरीके की तुलना में "ऑप्टिमाइज़र स्टेप" (बैठक) को 6 से 163 गुना तेज़ बना दिया।
- दक्षता: एक स्टेप को प्रशिक्षित करने का कुल समय अब मानक AdamW विधि से केवल 2% धीमा है।
- निष्कर्ष: DMuon टीमों को Muon की बेहतर "ग्रुप थेरेपी" सीखने की शैली का उपयोग करने की अनुमति देता है बिना भारी समय दंड चुकाए। यह एक धीमे, बोझिल प्रक्रिया को एक सुचारू, कुशल प्रक्रिया में बदल देता है, जिससे यह विशाल AI मॉडल्स के लिए उत्पादन (production) में उपयोग के लिए तैयार हो जाता है।
संक्षेप में, DMuon वह ट्रैफिक कंट्रोल सिस्टम है जो यह सुनिश्चित करता है कि सबसे स्मार्ट कोच (Muon) कंप्यूटरों के हाईवे पर बिना ट्रैफिक जाम के चल सके, जिससे अगली पीढ़ी के AI को प्रशिक्षित करना व्यावहारिक हो जाता है।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।