mamba
v2.2.2
Мамба: моделирование последовательностей линейного времени с выборочными пространствами состояний
Альберт Гу*, Три Дао*
Статья: https://arxiv.org/abs/2312.00752.
Трансформаторы — это ССМ: обобщенные модели и эффективные алгоритмы
Через структурированную двойственность пространства состояний
Три Дао*, Альберт Гу*
Статья: https://arxiv.org/abs/2405.21060.
Mamba — это новая архитектура модели пространства состояний, демонстрирующая многообещающую производительность при работе с насыщенными информацией данными, такими как языковое моделирование, где предыдущие субквадратичные модели не дотягивают до Transformers. Он основан на развитии структурированных моделей пространства состояний с эффективным аппаратно-ориентированным дизайном и реализацией в духе 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Блок Мамба-2 реализован в файле groups/mamba2.py.
Более простая версия находится в модулях/mamba2_simple.py.
Использование аналогично Мамбе(-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 находится по адресу elements/ssd_minimal.py.
Наконец, мы приводим пример полной языковой модели: костяк модели глубокой последовательности (с повторяющимися блоками Mamba) + заголовок языковой модели.
Источник: модели/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 вдвое больше, чем у Transformer аналогичного размера, поскольку для каждого «слоя» (блок MHA + блок MLP) Transformer необходимы два блока Mamba.)
Примечание. Это базовые модели, обученные только для токенов 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Чтобы провести оценку моделей Мамба-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С Мамбой-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}
}
