تعرف على Open Dreamer: إعادة إنتاج JAX/Flax لخط أنابيب نموذج Dreamer 4 World، مع نشر وصفة التدريب الكاملة

أصدرت مجموعة صغيرة من الباحثين في مجال الذكاء الاصطناعي (Reactor) Open Dreamer، وهو تطبيق مفتوح لخط أنابيب النموذج العالمي Dreamer 4 المكتوب بلغة JAX وFlax NNX.

ما شحنت فعلا

تم الافراج عن اثنين من المستودعات. يحمل Next-state/open-dreamer خط أنابيب التدريب: رمز فيديو سببي، ونموذج ديناميكيات كامنة مشروطة بالحركة، وتوليد الطرح، وتسجيل FVD. يحتوي فريق Reactor-team/open-dreamer على حد أدنى من أدوات الطرح المحلية التي تولد إطارات من ملف MP4 وملف إجراء مطابق.

الأداة الثالثة هي العرض التوضيحي للمتصفح المستضاف في وقت تشغيل Reactor. إنه يبث عالم Minecraft الذي تم إنشاؤه في الوقت الفعلي ويكشف عن تبديل Game ⟷ Dream الذي ينقل البث من اللعبة الحقيقية إلى النموذج العالمي إطارًا بإطار.

كان الهدف المعلن هو إعادة إنتاج بحث Dreamer 4. لقد تجنب فريق البحث عمدا الأساليب خارج تلك الورقة البحثية لإبقاء مساحة البحث ضيقة. لقد بدأوا باستخدام CoinRun، وهي منصة ثنائية الأبعاد تم إنشاؤها إجرائيًا ويمكن تدريبها على وحدة معالجة رسومات واحدة، ثم قاموا بتوسيع مسار العمل إلى فيديو أسلوب اللعب على طراز Minecraft/VPT.

العمارة: عمود فقري واحد ونموذجان

يستخدم كل من الرمز المميز والنموذج الديناميكي نفس العمود الفقري للمحول السببي. يتناوب هذا العمود الفقري بين نوعين من الاهتمام. تقوم طبقات الفضاء بنشر المعلومات بين عناصر إطار واحد. تقوم طبقات الوقت السببية بنشر المعلومات بين الإطارات.

إن رمز الرمز المميز عبارة عن أداة تشفير تلقائية مقنعة تعتمد على المحولات بدلاً من VAE. أبلغ الفريق عن ضغط يصل إلى 100× تقريبًا ولاحظ أن التصميم لا يحتاج إلى KL أو خسارة عدائية. ويقولون إن الإخفاء يجعل الفضاء الكامن أكثر قابلية للانتشار.

ينفذ النموذج الديناميكي التنبؤ بالإطار التالي ويتم تدريبه باستخدام نماذج الانتشار ومطابقة التدفق ونماذج الاختصار. كما يتنبأ بالإجراء التالي. بدلاً من التبديل بين وحدة انتقالية وسياسة منفصلة، ​​يتم دمج عملية الطرح في مجموعات لكل خطوة زمنية (الإجراء السابق، والحالة، والسياسة). يتم تشغيل الاهتمام المكاني داخل كل كتلة. الاهتمام الزمني السببي يربط الكتل عبر الزمن.

والأهم من ذلك، أن الرموز المميزة للنموذج العالمي لا يمكنها قراءة الرمز المميز للوكيل. وبالتالي، لا يمكن أن تؤثر معلومات المهمة والسياسة على الحالات المستقبلية إلا من خلال الإجراء التالي.

وصفة التدريب، كما تم تكوينها

تكوينات Minecraft التي تم شحنها تجعل الوصفة ملموسة.

النموذج الديناميكي هو 1.6B معلمات: 30 طبقة سببية، d_model 1920، 30 رأس انتباه، و3 رؤوس KV للاهتمام بالاستعلام المجمع. كل طبقة رابعة هي طبقة الاهتمام بالوقت. تحمل كل خطوة زمنية 32 رمزًا مميزًا للتسجيل، و packing_factor: 2 يقوم بتعبئة الرموز المميزة المجاورة في كل رمز مكاني ديناميكي. يستخدم الاهتمام بالوقت نافذة منزلقة مكونة من 192 خطوة.

يمتد التدريب لـ 200000 خطوة باستخدام Muon، وجدول WSD، ومعدل التعلم الأقصى 3e-4. يتم تشغيل عينات الاختصار/bootstrap عند الخطوة 100000 عند جزء دفعة يبلغ 0.25. اضمحلال EMA هو 0.999.

يُصدر تكوين أداة الرمز المميز 512 رمزًا كامنًا لكل إطار بعرض عنق الزجاجة يبلغ 16. يتم تعبئة الإطارات الأولية مقاس 360 × 640 إلى 368 × 640، لذا ينقسم كلا البعدين إلى تصحيحات 16 × 16. عمق التشفير هو 12 في d_model 1536؛ عمق فك التشفير هو 8 في d_model 1024. يصل احتمال إخفاء MAE إلى 0.9، ويتم تطبيق LPIPS بوزن 0.2 في نصف الخطوات الزمنية.

يتم تحليل إجراءات VPT إلى 27 قناة عمل ثنائية بالإضافة إلى 121 فئة ماوس فئوية، بدون قنوات مستمرة.

Computemaxxing وجدار الذاكرة

أبلغ فريق البحث عن استخدام 57-58% من نموذج FLOPs، مقابل معيار 60% المعلن للتدريب الصحي على المحولات. المنطق هو حجة خط السقف. في جهاز B200، يكون التقاطع بين النطاق الترددي المرتبط بالحساب والحساب هو 292 FLOP/بايت. تؤدي تغذية 256 إطارًا لكل وحدة معالجة رسومات إلى دفع عبء العمل إلى ما بعد نقطة التلال تلك.

ذهب Sharding في الاتجاه الآخر من التوقعات. عند 1.6B من المعلمات، تشغل حالة النموذج – المعلمات، والتدرجات، وحالة المُحسِّن، وEMA – ما يقرب من 24 جيجا بايت، وهو ما يتناسب مع B200. وكانت التنشيط التكلفة الحقيقية. قام فريق البحث بتجربة توازي البيانات، FSDP، توازي الموتر، وتوازي التسلسل، ثم استقر على توازي البيانات البسيطة بالإضافة إلى فحص التنشيط.

تم حل عملية تحميل البيانات عن طريق الترميز المسبق لمجموعة البيانات بأكملها .arrayrecord الملفات، ثم استخدام Grain مع مخزن مؤقت للجلب المسبق من جانب GPU. لم يكن فك تشفير ffmpeg سريعًا بما يكفي للحفاظ على تغذية وحدات معالجة الرسومات.

قسم الاستقرار هو الحمولة الحقيقية

ويوضح فريق البحث أن الاستقرار استهلك النصيب الأكبر من وقتهم. ملاحظتهم الرئيسية: تحدث معظم مشاكل الاستقرار بالرغم من الخسارة تنخفض. تتحسن MSE بسلاسة بينما تتدهور جودة التوليد.

تم توثيق ستة إصلاحات. حل Muon محل LaProp، الذي ارتفع بشكل عشوائي ومتزايد في كثير من الأحيان، عبر جولتين تبلغ كل منهما حوالي 400 B200 ساعة. يتم التعامل مع أوزان EMA على أنها إلزامية لاستدلال الانتشار. الدقة المختلطة حساسة للحدود: تظل المعلمات عائمة 32، ويغطي BF16 معظم عمليات تنشيط matmul ومدخلات الانتباه، ويتم الاحتفاظ بـ float32 للتطبيع ورأس إخراج التدفق الديناميكي.

عند فقدان الوزن، يستخدمون التنبؤ x مع فقدان مساحة v، والذي يتم تقليله إلى مصطلح ترجيح مشابه لمصطلح Dreamer 4 ولكن بمقام مربع. لقد أبلغوا عن تحسن صغير ولكنه ملحوظ. جعل النقل الأمثل من Minibatch barycentric بين الضوضاء والتسلسلات الكامنة توليد الطرح أكثر استقرارًا. تم اختبار المعلمات μ واعتبر أنها غير ضرورية، ويرجع ذلك جزئيًا إلى أن Muon يحافظ على ثبات المعلمات الفائقة عبر أحجام النماذج.

إحدى النتائج الإضافية من مرحلة CoinRun: عملية مسح iso-FLOPs تضع مقياس الحوسبة الأمثل عند مستوى تقريبًا $ن\propto {ج}^{0.56}$ و $د\propto {ج}^{0.44}$.

ما هو ليس في المربع

لا يتضمن المستودع حلقة تدريب استنساخ السلوك أو التعلم عن بعد؛ تم إدراج حلقة وكيل Dreamer 4 BC/RL الكاملة كعنصر مفتوح في خريطة الطريق. لم يتم استخدام سياسة CoinRun الموضحة في المنشور في لعبة Minecraft ولم يتم إصدارها.

ومع ذلك، لا ينشر المنشور أيضًا نتائج FVD scripts/eval_fvd.py يأتي مزودًا بأداة تعتمد على I3D تم تكوينها لأربعة إطارات سياقية وأفق 240 إطارًا.

الوجبات السريعة الرئيسية

• يقوم Open Dreamer بإعادة إنتاج مسار Dreamer 4 في JAX/Flax NNX، مع رمز تدريب وعرض توضيحي لـ Minecraft.
• نموذج الديناميكيات هو 1.6B معلمات، 30 طبقة، d_model في عام 1920، تدرب 200 ألف خطوة مع مون.
• الأرقام الهندسية المُبلغ عنها: 57-58% MFU على B200، 256 إطارًا لكل وحدة معالجة رسومات، ~24 جيجا بايت لحالة النموذج.
• وكان الاستقرار، وليس الإنتاجية، هو عنق الزجاجة؛ أخفت منحنيات الخسارة معظم انحدارات جودة الجيل.

قم بالاطلاع على منشور المدونة والعرض التوضيحي، ومستودع التدريب، ومستودع الاستدلال، وReactor on X. كل الفضل في هذا البحث يعود إلى الباحثين في هذا المشروع.