Open Dreamer: JAX/Flax में Dreamer 4 पाइपलाइन का ओपन-सोर्स कार्यान्वयन, पूर्ण प्रशिक्षण विधि सहित
AI शोधकर्ताओं की एक टीम ने Open Dreamer जारी किया, जो JAX और Flax NNX के साथ निर्मित Dreamer 4 वर्ल्ड-मॉडल पाइपलाइन का एक ओपन-सोर्स कार्यान्वयन है। इस प्रोजेक्ट में एक प्रशिक्षण पाइपलाइन, एक स्थानीय इन्फ़रेंस हार्नेस, और एक ब्राउज़र डेमो शामिल है जो Minecraft दुनिया को रीयल-टाइम में उत्पन्न करता है। मॉडल ब्लॉक-कॉज़ल ट्रांसफ़ॉर्मर बैकबोन का उपयोग करता है, 57-58% मॉडल FLOPs उपयोग प्राप्त करता है, और टीम इस बात पर जोर देती है कि स्थिरता मुख्य चुनौती थी, थ्रूपुट नहीं।
रिएक्टर द्वारा प्रायोजित शोधकर्ताओं फ्रांसेस्को सैको, डिएगो मार्टी और एडवर्ड हू ने ओपन ड्रीमर जारी किया, जो JAX/Flax NNX में ड्रीमर 4 का एक ओपन इम्प्लीमेंटेशन (open implementation) है। दो रिपॉजिटरी जारी की गईं: एक प्रशिक्षण के लिए (टोकनाइज़र, डायनामिक्स मॉडल, रोलआउट, FVD स्कोरिंग) और एक स्थानीय इन्फ़रेंस के लिए। एक ब्राउज़र डेमो एक जनरेटेड Minecraft दुनिया को स्ट्रीम करता है और वास्तविक गेम और वर्ल्ड मॉडल के बीच टॉगल करने की अनुमति देता है। आर्किटेक्चर टोकनाइज़र और डायनामिक्स मॉडल दोनों के लिए ब्लॉक-कॉज़ल ट्रांसफ़ॉर्मर बैकबोन का उपयोग करता है। टोकनाइज़र एक ट्रांसफ़ॉर्मर-आधारित मास्क्ड ऑटोएनकोडर है जो KL या विरोधी लॉस के बिना लगभग 100× संपीड़न प्राप्त करता है। डायनामिक्स मॉडल (1.6 बिलियन पैरामीटर) डिफ्यूज़न फ़ोर्सिंग, फ़्लो मैचिंग और शॉर्टकट मॉडल का उपयोग करके अगले-फ्रेम की भविष्यवाणी करता है, और अगली कार्रवाई की भी भविष्यवाणी करता है। प्रशिक्षण Muon ऑप्टिमाइज़र, WSD शेड्यूल, 3e-4 की पीक लर्निंग रेट और 0.999 के EMA डेके के साथ 200,000 स्टेप तक चलता है। टीम ने B200 GPU पर 57-58% मॉडल FLOPs उपयोगिता हासिल की, जिसमें मॉडल स्टेट लगभग 24 GiB में फ़िट होता है। एक्टिवेशन मुख्य मेमोरी लागत थी, जिसे डेटा पैरेललिज़्म और एक्टिवेशन चेकपॉइंटिंग द्वारा हल किया गया। स्थिरता संबंधी समस्याएँ सबसे बड़ी चुनौती थीं; छह फ़िक्स में Muon ऑप्टिमाइज़र, अनिवार्य EMA वेट, सावधान मिक्स्ड प्रिसिज़न, v-स्पेस लॉस वेटिंग, मिनीबैच बैरीसेंट्रिक ऑप्टिमल ट्रांसपोर्ट और μ-पैरामीटराइज़ेशन को कम महत्व देना शामिल है। रिपॉजिटरी में बिहेवियर-क्लोनिंग या RL ट्रेनिंग लूप शामिल नहीं है, और FVD स्कोर प्रकाशित नहीं किए गए हैं।
स्रोत: MarkTechPost —
मूल