mamba
v2.2.2
Mamba：具有选择性状态空间的线性时间序列建模
Albert Gu*、Tri Dao*
论文：https://arxiv.org/abs/2312.00752
Transformer 是 SSM：广义模型和高效算法
通过结构化状态空间对偶性
Tri Dao*、Albert Gu*
论文：https://arxiv.org/abs/2405.21060
Mamba 是一种新的状态空间模型架构，在语言建模等信息密集数据上显示出良好的性能，而之前的二次模型在 Transformers 方面存在不足。它基于结构化状态空间模型的进展，并本着 FlashAttention 的精神进行高效的硬件感知设计和实现。
pip install causal-conv1d>=1.4.0 ：Mamba 块内使用的简单因果 Conv1d 层的有效实现。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 模型的几个级别的接口。
Mamba 基于选择性 SSM 层，这是本文的重点（第 3 节；算法 2）。
来源：ops/selective_scan_interface.py。
该存储库的主要模块是包装选择性 SSM 的 Mamba 架构块。
来源：modules/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 . shapeMamba-2 块在 module/mamba2.py 中实现。
更简单的版本位于modules/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 模块的最小版本（来自 Mamba-2 论文的清单 1）以及“离散”和“连续”SSM 版本之间的转换位于 module/ssd_minimal.py。
最后，我们提供了一个完整的语言模型的示例：深度序列模型主干（具有重复的 Mamba 块）+ 语言模型头。
来源：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 ，在 Pile 上使用 300B 令牌进行训练，以及mamba-2.8b-slimpj （在 SlimPajama 数据集上使用 600B 令牌进行训练）。
模型将由下面的生成脚本自动下载。
这些模型在 Pile 上进行训练，并遵循 GPT-3 和许多开源模型所描述的标准模型维度：
| 参数 | 层数 | 型号暗淡。 |
|---|---|---|
| 130M | 24 | 第768章 |
| 370M | 48 | 1024 |
| 790M | 48 | 1536 |
| 1.4B | 48 | 2048 |
| 2.8B | 64 | 2560 |
（Mamba 的层数是具有相似大小的 Transformer 的两倍，因为 Transformer 的每个“层”（MHA 块 + MLP 块）需要两个 Mamba 块。）
注意：这些是仅针对 300B 令牌训练的基础模型，没有任何形式的下游修改（指令调整等）。性能预计与在类似数据上训练的其他架构相当或更好，但无法匹配更大或微调的模型。
为了对模型进行零样本评估（对应于本文的表 3），我们使用 lm-evaluation-harness 库。
pip install lm-eval==0.4.2安装lm-evaluation-harness 。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 的差异。
脚本 benchmarks/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}
}
