mesh transformer jax
1.0.0
مكتبة هايكو تستخدم عوامل xmap / pjit في JAX لنموذج التوازي للمحولات.
يشبه نظام التوازي نظام Megatron-LM الأصلي، والذي يتميز بالكفاءة في استخدام وحدات TPU نظرًا للشبكة الشبكية ثنائية الأبعاد عالية السرعة. هناك أيضًا إصدار نموذجي تجريبي يطبق أسلوب تقسيم ZeRo.
تم تصميم هذه المكتبة لقابلية التوسع حتى ما يقرب من 40 بايت من المعلمات على TPUv3s، والتي ينبغي بعد ذلك استخدام استراتيجيات التوازي المختلفة. راجع تطبيقات أخرى مثل GPT-NeoX أو DeepSpeed لذلك.
أحد الاتجاهات المستقبلية للبحث هو دمج قاعدة التعليمات البرمجية هذه مع Swarm-jax، لتحقيق المزيد من قابلية التوسع مع توازي خطوط الأنابيب.
21-07-12 : تمت إضافة دليل للضبط الدقيق
معلمة 6 مليار، نموذج إنشاء نص انحداري تم تدريبه على The Pile.
تحميل أوزان سليمة (أوزان bf16 فقط للاستدلال 9 جيجا)
تنزيل الأوزان الكاملة (بما في ذلك معلمات المحسن، 61 جيجابايت)
نقاط تفتيش مدربة جزئيا
عرض كولاب
عرض الويب
مشاركة مدونة آران
لم يكن هذا المشروع ممكنًا لولا الحوسبة التي قدمتها TPU Research Cloud بسخاء بمساعدة EleutherAI.
شكرًا لفريق Cloud TPU في Google لتوفير الوصول المبكر إلى Cloud TPU VM alpha (الآن متاح للعامة!)
شكرًا لكل من ساعدنا بطريقة أو بأخرى (مذكورة أبجديًا):
تم ترخيص أوزان GPT-J-6B بموجب الإصدار 2.0 من ترخيص Apache.
| المعلمة المفرطة | قيمة |
|---|---|
| n_parameters | 6,053,381,344 |
| n_layers | 28* |
| d_model | 4,096 |
| d_ff | 16,384 |
| n_heads | 16 |
| d_head | 256 |
| n_ctx | 2,048 |
| n_vocab | 50,257 (نفس رمز GPT-2/3) |
| ترميز الموقف | ترميزات الموضع الدوار (RoPE) |
| أبعاد حبل | 64 |
* تتكون كل طبقة من كتلة تغذية للأمام وكتلة واحدة للانتباه الذاتي
يتكون النموذج من 28 طبقة بعد النموذج 4096، وبُعد التغذية الأمامية 16384. وينقسم بُعد النموذج إلى 16 رأسًا، كل منها بعد 256. وتم تطبيق ترميز الموضع الدوار (RoPE) على 64 بُعدًا لكل رأس . تم تدريب النموذج باستخدام مفردات الترميز 50257، باستخدام نفس مجموعة BPEs مثل GPT-2/GPT-3.
يتم فرز النماذج تقريبًا حسب الأداء، أو حسب FLOPs إذا لم تكن متوفرة.
| نموذج | الأوزان | يتخبط التدريب | لامبادا PPL ↓ | لامبادا أك ↑ | فينوغراندي ↑ | هيلاسواج ↑ | بيكا ↑ | حجم مجموعة البيانات (جيجابايت) |
|---|---|---|---|---|---|---|---|---|
| فرصة | ✔ | 0 | ~ الكثير | ~0% | 50% | 25% | 25% | 0 |
| GPT-3-Ada‡ | ✘ | ----- | 9.95 | 51.6% | 52.9% | 43.4% | 70.5% | ----- |
| جي بي تي-2-1.5ب | ✔ | ----- | 10.63 | 51.21% | 59.4% | 50.9% | 70.8% | 40 |
| GTTNeo-1.3B‡ | ✔ | 3.0e21 | 7.50 | 57.2% | 55.0% | 48.9% | 71.1% | 825 |
| ميجاترون-2.5ب* | ✘ | 2.4e21 | ----- | 61.7% | ----- | ----- | ----- | 174 |
| GTTNeo-2.7B‡ | ✔ | 6.8e21 | 5.63 | 62.2% | 56.5% | 55.8% | 73.0% | 825 |
| جي بي تي-3-1.3ب*‡ | ✘ | 2.4e21 | 5.44 | 63.6% | 58.7% | 54.7% | 75.1% | ~800 |
| GPT-3-باباج‡ | ✘ | ----- | 5.58 | 62.4% | 59.0% | 54.5% | 75.5% | ----- |
| ميجاترون-8.3ب* | ✘ | 7.8e21 | ----- | 66.5% | ----- | ----- | ----- | 174 |
| جي بي تي-3-2.7ب*‡ | ✘ | 4.8e21 | 4.60 | 67.1% | 62.3% | 62.8% | 75.6% | ~800 |
| ميجاترون-11ب† | ✔ | 1.0e22 | ----- | ----- | ----- | ----- | ----- | 161 |
| جي بي تي-ي-6بي ‡ | ✔ | 1.5e22 | 3.99 | 69.7% | 65.3% | 66.1% | 76.5% | 825 |
| جي بي تي-3-6.7ب*‡ | ✘ | 1.2e22 | 4.00 | 70.3% | 64.5% | 67.4% | 78.0% | ~800 |
| جي بي تي-3-كوري‡ | ✘ | ----- | 4.00 | 69.3% | 65.6% | 68.5% | 77.9% | ----- |
| جي بي تي-3-13ب*‡ | ✘ | 2.3e22 | 3.56 | 72.5% | 67.9% | 70.9% | 78.5% | ~800 |
| جي بي تي-3-175ب*‡ | ✘ | 3.1ه23 | 3.00 | 76.2% | 70.2% | 78.9% | 81.0% | ~800 |
| جي بي تي-3-دافينشي‡ | ✘ | ----- | 3.0 | 75% | 72% | 78% | 80% | ----- |
| غوفر 230ب* | ✘ | 6.31E+23 | ----- | 74.50% | 70.10% | 79.20% | 81.80% | 1344 |
| MT-NLG 530B*‡ | ✘ | ----- | ----- | 76.6% | 73.0% | 80.2% | 82.0% | ----- |
* تمثل أرقام التقييم التي أبلغ عنها مؤلفوها، ويتم توفير جميع الأرقام الأخرى عن طريق تشغيل lm-evaluation-harness إما مع الأوزان الصادرة أو من خلال الوصول إلى واجهة برمجة التطبيقات (API). نظرًا للاختلافات الدقيقة في التنفيذ بالإضافة إلى اختلاف إطار المهام الصفرية، فقد لا تكون هذه قابلة للمقارنة بشكل مباشر. راجع منشور المدونة هذا لمزيد من التفاصيل.
† لا يوفر نموذج Megatron-11B أي مقاييس قابلة للمقارنة، والعديد من التطبيقات التي تستخدم الأوزان الصادرة لا تعيد إنتاج جودة التوليد والتقييمات. (انظر 1 2 3) وبالتالي، لم تتم محاولة التقييم.
‡ تم تدريب هذه النماذج باستخدام البيانات التي تحتوي على تلوث محتمل لمجموعة الاختبار. فشلت نماذج OpenAI GPT-3 في إلغاء تكرار بيانات التدريب لمجموعات اختبار معينة، في حين تم تدريب نماذج GPT-Neo بالإضافة إلى هذا النموذج على The Pile، والتي لم يتم إلغاء تكرارها مقابل أي مجموعات اختبار.
تم تصميم معظم البرامج النصية الموجودة في هذا المستودع ليتم تشغيلها على وحدات TPU، والتي تعد ضمن بنية TPU-VM أجهزة افتراضية يمكنها تشغيل تعليمات برمجية عشوائية. تم تصميم معظم البرامج النصية لتدوير TPU وSSH فيه لإعداد التبعيات ونسخ التعليمات البرمجية من الدليل المحلي، ثم بدء عامل Ray الذي يمكنه قبول مكالمات RPC.
تتعامل أجهزة TPUVMs مع خطوات التدريب النموذجية وتقييمها وحفظ نقاط التفتيش وتحميلها، بينما يتعامل برنامج python لبرنامج التشغيل مع تحميل البيانات والتنسيق العام (مثل وقت حفظ نقاط التفتيش وما إلى ذلك).
وهذا يعني أن معظم البرامج النصية ( train.py و eval_harness.py وما إلى ذلك) تتوقع أن يتم تشغيلها على جهاز ظاهري GCE في نفس المنطقة مثل وحدات TPU، لتقليل زمن انتقال RPC وتكلفة نقل البيانات. تتوقع البرامج النصية الأخرى (عادةً تلك التي لا تأخذ وسيطة --tpu ، مثل device_sample.py أو device_serve.py أو device_train.py ) أن يتم تشغيلها مباشرة على TPUVM. تعمل البرامج النصية Device_* فقط على الإصدار v3-8 وليس على البودات الأكبر حجمًا.
علاوة على ذلك، يوجد مثال ( resharding_example.py ) لكيفية تحويل نقاط التحقق المتوفرة (التي تحتوي على 8 أجزاء في حالة GPT-J-6B) إلى عدد أصغر، كما هو الحال عند التشغيل على وحدات معالجة الرسومات.
لضبط النموذج، قم بتشغيل device_train.py على جهاز TPU VM. باستخدام TPU v3-8، يمكنك الضبط بمعدل ~5000 رمز مميز في الثانية، وهو ما يجب أن يكون كافيًا لمجموعات البيانات الصغيرة والمتوسطة الحجم.
يرجى قراءة الدليل خطوة بخطوة للحصول على تعليمات الضبط الدقيق.
لاحظ أن هذه المكتبة لديها بعض المتطلبات المحددة لإصدار JAX. على وجه التحديد، لاستخدام نماذج v1 (بما في ذلك GPT-J 6B)، يلزم وجود jax==0.2.12 . وهذا بدوره يعتمد على jaxlib==0.1.68 . إذا لم يتم ذلك، فسوف تحصل على أخطاء xmap غامضة
ومع ذلك، لاستخدام رمز النموذج v2 (لا توجد أوزان منشورة بشكل عام)، يمكن استخدام أحدث إصدار من JAX.
للاستشهاد بهذا المستودع:
@misc{mesh-transformer-jax,
author = {Wang, Ben},
title = {{Mesh-Transformer-JAX: Model-Parallel Implementation of Transformer Language Model with JAX}},
howpublished = {url{https://github.com/kingoflolz/mesh-transformer-jax}},
year = 2021,
month = May
}
للاستشهاد بأوزان GPT-J-6B:
@misc{gpt-j,
author = {Wang, Ben and Komatsuzaki, Aran},
title = {{GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model}},
howpublished = {url{https://github.com/kingoflolz/mesh-transformer-jax}},
year = 2021,
month = May
}
إذا كنت تستخدم هذا المستودع أو أيًا من الأوزان المُدربة مسبقًا للقيام بشيء رائع، فنحن نود أن نسمع عنه. لا تتردد في فتح قضية جيثب أو التواصل عبر البريد الإلكتروني (في الملف الشخصي).
