mamba

Mamba — это открытый исследовательский код от группы state-spaces, реализующий семейство архитектур моделей пространства состояний (SSM) для обработки последовательностей. Репозиторий сопровождает статьи «Mamba: Linear-Time Sequence Modeling with Selective State Spaces», «Transformers are SSMs» (Mamba-2) и «Mamba-3: Improved Sequence Modeling using State Space Principles». Проект позиционируется как сублинейная по сложности альтернатива Transformer, особенно для данных с высокой плотностью информации, таких как языковое моделирование, и опирается на более ранние работы по структурированным пространствам состояний (S4) с аппаратно-ориентированной реализацией в духе FlashAttention. Что предоставляет репозиторий: - Selective SSM слой: основная операция selective scan, исходный код в ops/selective_scan_interface.py. - Блок Mamba: основной архитектурный модуль, оборачивающий selective SSM (modules/mamba_simple.py); используется как готовый PyTorch-модуль с настраиваемыми параметрами d_model, d_state, d_conv и expand. - Блок Mamba-2: реализован в modules/mamba2.py, упрощённый вариант — в modules/mamba2_simple.py, а также минимальный эталонный модуль SSD (structured state space duality) в modules/ssd_minimal.py. - Блок Mamba-3: реализован в modules/mamba3.py; доступны опции d_state, headdim, режим MIMO с mimo_rank, chunk_size, необязательная норма после выходной проекции и тип bfloat16. - Пример языковой модели: полный backbone плюс LM-голова в models/mixer_seq_simple.py, используется скриптами генерации. Варианты установки и сборки: Пакет устанавливается через pip (mamba-ssm) и требует Linux, Python 3.10+ и PyTorch 1.12+. По умолчанию установка не компилирует расширение selective_scan_cuda и не загружает кэшированные CUDA-колёса. Дополнительные зависимости и переменные окружения меняют это поведение: extra causal-conv1d добавляет эту зависимость; MAMBA_FORCE_BUILD=TRUE принудительно запускает локальную сборку; MAMBA_KEEP_CUDA_BUILD=TRUE включает CUDA-расширение selective scan, сначала пытаясь найти подходящее готовое CUDA/HIP-колесо, а если его нет — компилирует локально. Сочетание обоих флагов принудительно запускает локальную CUDA-компиляцию. В README отмечается, что для CUDA-сборки требуется --no-build-isolation, чтобы pip использовал уже установленный PyTorch с поддержкой CUDA. Для CUDA-сборки и выполнения на GPU также нужны NVIDIA GPU и CUDA 11.6+. Для карт AMD в README описан шаг с патчем для ROCm 6.0 (не требуется начиная с ROCm 6.1). Предобученные модели: Веса опубликованы на Hugging Face в организации state-spaces, включая 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. В README эти модели описаны как базовые, без инструкционной донастройки или других изменений под конкретные задачи; ожидается, что их качество сравнимо или лучше, чем у других архитектур, обученных на аналогичных данных, но они не сопоставимы с более крупными или донастроенными моделями. Размерности моделей следуют GPT-3-подобному масштабированию (например, 130M параметров, 24 слоя и размер модели 768, вплоть до 2.8B с 64 слоями и размерностью 2560). Инструменты оценки и инференса: Zero-shot оценка выполняется через библиотеку lm-evaluation-harness (pip install lm-eval==0.4.2); примеры команд покрывают такие задачи, как lambada_openai, hellaswag, piqa, arc_easy, arc_challenge, winogrande, openbookqa, boolq, race, truthfulqa_mc2 и mmlu. В README отмечается, что результаты по отдельным задачам могут отличаться от опубликованных на 0.1–0.3 из-за шума оценки. Скрипт бенчмарка генерации (benchmarks/benchmark_generation_mamba_simple.py) автоматически загружает модель из Hugging Face Hub, генерирует продолжения для пользовательского промпта и измеряет скорость инференса; доступны опции top-p, min-p, top-k, temperature, repetition penalty и batch size. Заметки по устранению неполадок: В README выделены две практические проблемы. Первая — точность: модели обучались с PyTorch AMP, который хранит параметры в float32 и при необходимости приводит к половинной точности; поскольку SSM чувствительны к рекуррентной динамике, фреймворки, хранящие параметры в float16 (например, DeepSpeed), могут вызывать нестабильность, и в качестве первого средства предлагается хранение параметров в fp32. Вторая — инициализация: некоторые компоненты наследуют инициализации в стиле S4 (например, целевой диапазон параметра delta через смещение линейной проекции), и фреймворки с пост-инициализационными хуками, обнуляющими смещения nn.Linear, могут потребовать дополнительной логики для их сохранения. Лицензия и цитирование: в README приведены BibTeX-записи для статей Mamba, Mamba-2 и Mamba-3, и пользователей просят цитировать их при использовании кодовой базы.