mamba
v2.2.2
مامبا: نمذجة التسلسل الزمني الخطي مع مساحات الحالة الانتقائية
ألبرت جو*، تري داو*
الورقة: https://arxiv.org/abs/2312.00752
المحولات هي SSMs: النماذج المعممة والخوارزميات الفعالة
من خلال ازدواجية مساحة الدولة المنظمة
تري داو*، ألبرت جو*
الورقة: https://arxiv.org/abs/2405.21060
مامبا عبارة عن بنية جديدة لنموذج الفضاء الحكومي تُظهر أداءً واعدًا في البيانات كثيفة المعلومات مثل نمذجة اللغة، حيث لا ترقى النماذج التربيعية السابقة إلى مستوى المحولات. وهو يعتمد على خط التقدم في نماذج مساحة الحالة المنظمة، مع تصميم وتنفيذ فعال يراعي الأجهزة بروح FlashAttention.
pip install causal-conv1d>=1.4.0 : تنفيذ فعال لطبقة Conv1d سببية بسيطة مستخدمة داخل كتلة Mamba.pip install mamba-ssm : حزمة Mamba الأساسية.pip install mamba-ssm[causal-conv1d] : لتثبيت حزمة Mamba الأساسية وcausal-conv1d.pip install mamba-ssm[dev] : لتثبيت حزمة Mamba الأساسية وتبعيات التطوير. ويمكن أيضًا بناؤه من المصدر باستخدام pip install . من هذا المستودع.
إذا اشتكت pip من إصدارات PyTorch، فحاول تمرير --no-build-isolation إلى pip .
متطلبات أخرى:
بالنسبة لبطاقات AMD، راجع المتطلبات الأساسية الإضافية أدناه.
نعرض عدة مستويات من الواجهة مع نموذج Mamba.
يعتمد مامبا على طبقة SSM انتقائية، وهي محور الورقة (القسم 3؛ الخوارزمية 2).
المصدر: ops/selective_scan_interface.py.
الوحدة الرئيسية لهذا المستودع هي كتلة بنية Mamba التي تغطي SSM الانتقائي.
المصدر: الوحدات/mamba_simple.py.
الاستخدام:
import torch
from mamba_ssm import Mamba
batch , length , dim = 2 , 64 , 16
x = torch . randn ( batch , length , dim ). to ( "cuda" )
model = Mamba (
# This module uses roughly 3 * expand * d_model^2 parameters
d_model = dim , # Model dimension d_model
d_state = 16 , # SSM state expansion factor
d_conv = 4 , # Local convolution width
expand = 2 , # Block expansion factor
). to ( "cuda" )
y = model ( x )
assert y . shape == x . shapeيتم تنفيذ كتلة Mamba-2 في الوحدات النمطية/mamba2.py.
يوجد إصدار أبسط في الوحدات/mamba2_simple.py
الاستخدام مشابه لـ Mamba(-1):
from mamba_ssm import Mamba2
model = Mamba2 (
# This module uses roughly 3 * expand * d_model^2 parameters
d_model = dim , # Model dimension d_model
d_state = 64 , # SSM state expansion factor, typically 64 or 128
d_conv = 4 , # Local convolution width
expand = 2 , # Block expansion factor
). to ( "cuda" )
y = model ( x )
assert y . shape == x . shape يوجد إصدار مصغر من وحدة SSD الداخلية (القائمة 1 من ورقة Mamba-2) مع التحويل بين إصدارات SSM "المنفصلة" و"المستمرة" على Modules/ssd_minimal.py.
أخيرًا، نقدم مثالاً لنموذج لغة كامل: العمود الفقري لنموذج التسلسل العميق (مع كتل مامبا المتكررة) + رأس نموذج اللغة.
المصدر:models/mixer_seq_simple.py.
هذا مثال على كيفية دمج Mamba في شبكة عصبية شاملة. يتم استخدام هذا المثال في البرامج النصية للإنشاء أدناه.
يتم تحميل النماذج المدربة مسبقًا إلى Hugging Face: mamba-130m , mamba-370m , mamba-790m , mamba-1.4b , mamba-2.8b , mamba2-130m , mamba2-370m , mamba2-780m , mamba2-1.3b , mamba2-2.7b ، transformerpp-2.7b و mamba2attn-2.7b ، تم تدريبهما على 300B من الرموز المميزة على Pile، بالإضافة إلى mamba-2.8b-slimpj (تم تدريبهما على 600B من الرموز المميزة في مجموعة بيانات SlimPajama).
سيتم تنزيل النماذج تلقائيًا بواسطة البرنامج النصي للإنشاء أدناه.
تم تدريب هذه النماذج على Pile، وتتبع أبعاد النموذج القياسية الموصوفة في GPT-3 والتي يتبعها العديد من النماذج مفتوحة المصدر:
| حدود | طبقات | النموذج خافت. |
|---|---|---|
| 130 م | 24 | 768 |
| 370 م | 48 | 1024 |
| 790 م | 48 | 1536 |
| 1.4 ب | 48 | 2048 |
| 2.8 ب | 64 | 2560 |
(يضاعف عدد طبقات Mamba عدد طبقات المحول ذي الحجم المماثل، حيث يلزم وجود كتلتي Mamba لكل "طبقة" (كتلة MHA + كتلة MLP) للمحول.)
ملاحظة: هذه نماذج أساسية تم تدريبها فقط لـ 300B من الرموز المميزة، دون أي شكل من أشكال التعديل النهائي (ضبط التعليمات، وما إلى ذلك). من المتوقع أن يكون الأداء قابلاً للمقارنة أو أفضل من البنى الأخرى التي تم تدريبها على بيانات مماثلة، ولكن ليس لمطابقة النماذج الأكبر أو المضبوطة بدقة.
لإجراء تقييمات صفرية للنماذج (المقابلة للجدول 3 من الورقة)، نستخدم مكتبة lm-evaluation-harness.
lm-evaluation-harness عن طريق pip install lm-eval==0.4.2 .lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-130m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
python evals/lm_harness_eval.py --model hf --model_args pretrained=EleutherAI/pythia-160m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande --device cuda --batch_size 64 لإعادة إنتاج النتائج على نموذج mamba-2.8b-slimpj المذكورة في منشورات المدونات:
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks boolq,piqa,hellaswag,winogrande,arc_easy,arc_challenge,openbookqa,race,truthfulqa_mc2 --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks mmlu --num_fewshot 5 --device cuda --batch_size 256لإجراء التقييمات على نماذج Mamba-2، ما عليك سوى استبدال أسماء النماذج:
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/transformerpp-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2attn-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256لاحظ أن نتيجة كل مهمة قد تختلف عن القيم المبلغ عنها بنسبة 0.1-0.3 بسبب الضوضاء في عملية التقييم.
معايير البرنامج النصي/benchmark_generation_mamba_simple.py
تتضمن الخيارات الأخرى القابلة للتكوين احتمالية top-p (أخذ عينات النواة)، ودرجة حرارة softmax.
لاختبار زمن وصول التوليد (على سبيل المثال، حجم الدفعة = 1) باستخدام استراتيجيات مختلفة لأخذ العينات:
python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name " EleutherAI/pythia-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --minp 0.05 --topk 0 --temperature 0.7 --repetition-penalty 1.2لاختبار إنتاجية التوليد بمطالبات عشوائية (على سبيل المثال، حجم دفعة كبيرة):
python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --batch 64
python benchmarks/benchmark_generation_mamba_simple.py --model-name " EleutherAI/pythia-2.8b " --batch 64مع Mamba-2، ما عليك سوى تغيير اسم الطراز:
python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba2-2.7b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2تم تدريب نماذجنا باستخدام PyTorch AMP للحصول على دقة مختلطة. يحتفظ AMP بمعلمات النموذج في float32 ويلقي بنصف الدقة عند الضرورة. من ناحية أخرى، تقوم أطر العمل الأخرى مثل DeepSpeed بتخزين المعلمات في float16 والبث عند الضرورة (على سبيل المثال لتراكم المحسن).
لقد لاحظنا أن الدقة الأعلى لمعلمات النموذج الرئيسية قد تكون ضرورية، لأن أجهزة SSM حساسة لديناميكياتها المتكررة. إذا كنت تعاني من عدم الاستقرار، كخطوة أولى، يرجى تجربة إطار عمل يخزن المعلمات في fp32 (مثل AMP).
تحتوي بعض أجزاء النموذج على عمليات تهيئة موروثة من العمل السابق على نماذج S4. على سبيل المثال، nn.Linear على الصفر). إذا كان الأمر كذلك، فقد يتعين عليك إضافة منطق مخصص (على سبيل المثال، يقوم هذا السطر بإيقاف إعادة التهيئة في مدربنا، ولكنه سيكون محظورًا في أي إطار عمل آخر) خاص بإطار العمل التدريبي.
إذا كنت تستخدم ROCm 6.0، فقم بتنفيذ الخطوات التالية لتجنب الأخطاء أثناء الترجمة. هذا غير مطلوب لـ ROCm 6.1 وما بعده.
حدد موقع دليل تثبيت ROCm الخاص بك. يوجد هذا عادة في /opt/rocm/ ، ولكن قد يختلف حسب التثبيت لديك.
تطبيق التصحيح. قم بالتشغيل باستخدام sudo في حالة مواجهة مشكلات في الأذونات.
patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch إذا كنت تستخدم قاعدة التعليمات البرمجية هذه، أو وجدت عملنا ذا قيمة، فيرجى الاستشهاد بـ Mamba:
@article{mamba,
title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
author={Gu, Albert and Dao, Tri},
journal={arXiv preprint arXiv:2312.00752},
year={2023}
}
@inproceedings{mamba2,
title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
author={Dao, Tri and Gu, Albert},
booktitle={International Conference on Machine Learning (ICML)},
year={2024}
}
