SkewAdam Mixture-of-Experts ट्रेनिंग मेमरी ६०% पेक्षा जास्त कमी करते आणि अचूकतेमध्ये (accuracy) लक्षणीय वाढ करते, ज्यामुळे डेव्हलपर्सना ६.७८ अब्ज पॅरामीटर्सचे MoE मॉडेल एका सिंगल ४० GB GPU वर चालवणे शक्य होते.

MoE ट्रेनिंगमध्ये मेमरीची मर्यादा का येते?

Mixture-of-Experts आर्किटेक्चरमध्ये, एक 'dense backbone' राखून प्रत्येक इनपुटला "expert" सब-नेटवर्क्सच्या एका लहान उपसमूहाद्वारे (subset) पाठवले जाते. याचा फायदा असा की, मॉडेलमधील कॉम्प्युटेशनमध्ये (compute) प्रमाणबद्ध वाढ न होता ते अब्जावधी पॅरामीटर्सपर्यंत स्केल होऊ शकते. प्रत्यक्षात, ऑप्टिमायझर—विशेषतः AdamW व्हेरिएंट जो बहुतेक टीम्स वापरतात—GPU RAM चा मोठा भाग वापरतो. ६.७८ अब्ज पॅरामीटर्सच्या मॉडेलसाठी, केवळ ऑप्टिमायझरचे momentum आणि variance tensors सुमारे ५० GB जागा घेतात. यामध्ये activations आणि model weights मिळवल्यास, पीक मेमरी (peak memory) ८१.४ GB पर्यंत पोहोचते, ज्यामुळे multi-GPU सेटअप किंवा संथ ट्रेनिंग शेड्यूल्स वापरावे लागतात.

SkewAdam काय वेगळे करते

SkewAdam कोणताही नवीन लर्निंग रूल शोधत नाही. ते Adam चे moments कुठे साठवायचे याचे पुनर्रचना (reshuffle) करते:

  • Dense backbone पूर्ण-अचूकतेचे (full-precision) momentum राखून ठेवते, ज्यामुळे dense layers ला आवश्यक असलेले स्मूथ अपडेट्स जपले जातात.
  • Expert bank variance चा एक 'factored approximation' साठवते, ज्यामुळे प्रत्येक expert साठी साठवलेला डेटा कमी होतो.
  • Router, जो कोणते experts सक्रिय करायचे हे ठरवतो, तो exact second-moment अंदाज (estimate) राखून ठेवतो.

या तीन घटकांना स्वतंत्र "tiers" म्हणून treating करून, ऑप्टिमायझर जिथे कमी महत्त्व आहे तिथे अनावश्यक अचूकता (redundant precision) कमी करते आणि जिथे सर्वाधिक महत्त्व आहे तिथे ती कायम ठेवते.

मोजता येण्याजोगा परिणाम

SkewAdam सह तेच ६.७८ अब्ज पॅरामीटर्सचे MoE मॉडेल चालवल्यास खालील परिणाम मिळतात:

  • Peak GPU memory: ३१.३ GB (८१.४ GB वरून कमी)
  • Optimizer-state footprint: १.२९ GB (५० GB वरून कमी)

कमी झालेली ही स्टेट (state) एका सिंगल ४० GB accelerator मध्ये सहज बसते.

केवळ बचत नाही, तर अचूकतेमध्येही वाढ

मेमरी कमी केल्यामुळे अनेकदा मॉडेलच्या गुणवत्तेवर परिणाम होतो, परंतु SkewAdam perplexity—जी लँग्वेज-मॉडेलची एक मानक मेट्रिक आहे—१२६.८ (AdamW) वरून १०८.४ पर्यंत सुधारते. हे त्याच कामासाठी Muon (१२०.२) आणि Lion (३९३.७) पेक्षाही सरस कामगिरी करते. लेखकांच्या मते, ही वाढ तिन्ही tiers मध्ये momentum जपल्यामुळे होते; Adafactor सारख्या पद्धती, ज्या पूर्णपणे momentum काढून टाकतात, त्या मागे पडतात.

उघडे प्रश्न

SkewAdam चे निकाल ६.७८ अब्ज पॅरामीटर्सच्या मॉडेलवर सिद्ध करण्यात आले आहेत. मॉडेलचा आकार कित्येक पटीने वाढल्यास किंवा लँग्वेज मॉडेलिंगच्या पलीकडील कामांसाठी हेच मेमरी-स्टेट रेशो (memory-state ratios) लागू होतील की नाही, हे अद्याप स्पष्ट नाही.

काय लक्ष ठेवायचे

  • मोठ्या MoE कॉन्फिगरेशनवरील (अब्जावधी पॅरामीटर्स) बेंचमार्क्स.
  • ओपन-सोर्स फ्रेमवर्क्सद्वारे स्वीकार आणि लोकप्रिय ट्रेनिंग स्क्रिप्ट्समध्ये समावेश.
  • लपलेले ट्रेड-ऑफ्स (trade-offs), जसे की कन्वर्जन्स स्पीड (convergence speed) किंवा वेगवेगळ्या लर्निंग-रेट शेड्यूल्स अंतर्गत स्थिरता (stability) यावर कम्युनिटीचा फीडबॅक.

स्रोत: ऑप्टिमायझरची रचना आणि प्रायोगिक निकाल तपशीलवार सांगणारी डेव्हलपर पोस्ट.