SkewAdam Mixture-of-Experts ट्रेनिंग मेमोरी को 60% से अधिक कम कर देता है और सटीकता (accuracy) में उल्लेखनीय वृद्धि प्रदान करता है, जिससे डेवलपर्स एक 6.78 बिलियन-पैरामीटर वाले MoE मॉडल को सिंगल 40 GB GPU पर फिट कर सकते हैं।

MoE ट्रेनिंग मेमोरी की सीमा (memory wall) क्यों पार नहीं कर पाती

Mixture-of-Experts आर्किटेक्चर एक डेंस बैकबोन (dense backbone) को बनाए रखते हुए प्रत्येक इनपुट को "एक्सपर्ट" सब-नेटवर्क्स के एक छोटे सबसेट के माध्यम से रूट करता है। इसका लाभ यह है कि मॉडल कंप्यूट (compute) में आनुपातिक वृद्धि के बिना अरबों पैरामीटर्स तक स्केल हो सकता है। व्यवहार में, ऑप्टिमाइज़र—विशेष रूप से AdamW वेरिएंट जिसका अधिकांश टीमें उपयोग करती हैं—GPU RAM का अधिकांश हिस्सा ले लेता है। 6.78 B-पैरामीटर मॉडल के लिए, केवल ऑप्टिमाइज़र के मोमेंटम (momentum) और वेरिएंस (variance) टेंसर ही लगभग 50 GB की मांग करते हैं। यदि इसमें एक्टिवेशन (activations) और मॉडल वेट्स (weights) को जोड़ दिया जाए, तो पीक मेमोरी 81.4 GB तक पहुँच जाती है, जिससे मल्टी-GPU सेटअप या धीमी ट्रेनिंग शेड्यूल्स की आवश्यकता पड़ती है।

SkewAdam क्या अलग करता है

SkewAdam कोई नया लर्निंग रूल नहीं बनाता है। यह Adam के मोमेंट्स (moments) के स्थान को पुनर्गठित (reshuffle) करता है:

  • डेंस बैकबोन फुल-प्रिसिजन मोमेंटम को बनाए रखता है, जिससे डेंस लेयर्स के लिए आवश्यक स्मूथ अपडेट्स सुरक्षित रहते हैं।
  • एक्सपर्ट बैंक वेरिएंस का एक फैक्टर्ड एप्रोक्सिमेशन (factored approximation) स्टोर करता है, जिससे प्रत्येक एक्सपर्ट के लिए रखे जाने वाले डेटा में कमी आती है।
  • राउटर, जो यह तय करता है कि किन एक्सपर्ट्स को एक्टिवेट करना है, एक सटीक सेकंड-मोमेंट एस्टीमेट (second-moment estimate) रखता है।

इन तीन घटकों को अलग-अलग "टियर्स" (tiers) के रूप में मानकर, ऑप्टिमाइज़र वहां अनावश्यक प्रिसिजन (precision) को हटा देता है जहां इसकी कम आवश्यकता होती है और वहां इसे बनाए रखता है जहां इसकी सबसे अधिक आवश्यकता होती है।

मापने योग्य प्रभाव

SkewAdam के साथ उसी 6.78 B-पैरामीटर वाले MoE मॉडल को चलाने पर निम्नलिखित परिणाम मिलते हैं:

  • पीक GPU मेमोरी: 31.3 GB (81.4 GB से कम)
  • ऑप्टिमाइज़र-स्टेट फुटप्रिंट: 1.29 GB (50 GB से कम)

कम हुआ यह स्टेट (state) सिंगल 40 GB एक्सेलेरेटर के भीतर आसानी से फिट हो जाता है।

केवल बचत नहीं, सटीकता में भी वृद्धि

मेमोरी में कटौती अक्सर मॉडल की गुणवत्ता को नुकसान पहुँचाती है, लेकिन SkewAdam परप्लेक्सिटी (perplexity)—जो कि एक मानक लैंग्वेज-मॉडल मेट्रिक है—को 126.8 (AdamW) से बढ़ाकर 108.4 कर देता है। यह उसी कार्य पर Muon (120.2) और Lion (393.7) से भी बेहतर प्रदर्शन करता है। लेखकों का कहना है कि यह सुधार तीनों टियर्स में मोमेंटम को बनाए रखने से आता है; Adafactor जैसी विधियाँ जो मोमेंटम को पूरी तरह से हटा देती हैं, वे पीछे रह जाती हैं।

खुले प्रश्न

SkewAdam के परिणाम एक 6.78 B-पैरामीटर मॉडल पर प्रदर्शित किए गए हैं। यह अभी भी स्पष्ट नहीं है कि क्या यही मेमोरी-स्टेट अनुपात एक ऑर्डर ऑफ मैग्नीट्यूड (order of magnitude) बड़े मॉडलों के लिए या लैंग्वेज मॉडलिंग से परे कार्यों के लिए भी लागू होंगे।

किन बातों पर नज़र रखें

  • बड़े MoE कॉन्फ़िगरेशन (अरबों पैरामीटर्स) पर बेंचमार्क।
  • ओपन-सोर्स फ्रेमवर्क द्वारा अपनाना और लोकप्रिय ट्रेनिंग स्क्रिप्ट में शामिल होना।
  • छिपे हुए ट्रेड-ऑफ (trade-offs) पर कम्युनिटी फीडबैक, जैसे कि अलग-अलग लर्निंग-रेट शेड्यूल्स के तहत कन्वर्जेंस स्पीड (convergence speed) या स्थिरता।

स्रोत: डेवलपर पोस्ट जिसमें ऑप्टिमाइज़र के डिज़ाइन और अनुभवजन्य (empirical) परिणामों का विवरण दिया गया है।