Mamba

Mamba：具有选择性状态空间的线性时间序列建模
Albert Gu*, Tri Dao*
论文：https://arxiv.org/abs/2312.00752

Mamba-2

Transformer 是 SSM：通过结构化状态空间对偶性
实现的广义模型和高效算法
Tri Dao*, Albert Gu*
论文：https://arxiv.org/abs/2405.21060

简介

Mamba 是一种新的状态空间模型架构，在信息密集型数据（如语言建模）方面表现出promising的性能，而之前的次二次模型在这方面不及 Transformer。它基于结构化状态空间模型的进展，采用了高效的硬件感知设计和实现，其精神类似于 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 版本问题，请尝试向 pip 传递 --no-build-isolation 参数。

其他要求：

Linux
NVIDIA GPU
PyTorch 1.12+
CUDA 11.6+

对于 AMD 显卡，请参阅下面的额外先决条件。

使用

我们提供了几个级别的 Mamba 模型接口。

选择性 SSM

Mamba 基于选择性 SSM 层，这是论文的重点（第 3 节；算法 2）。

源代码：ops/selective_scan_interface.py。

Mamba 模块

此存储库的主要模块是包装选择性 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(
    # 此模块大约使用 3 * expand * d_model^2 个参数
    d_model=dim, # 模型维度 d_model
    d_state=16,  # SSM 状态扩展因子
    d_conv=4,    # 局部卷积宽度
    expand=2,    # 块扩展因子
).to("cuda")
y = model(x)
assert y.shape == x.shape

Mamba-2

Mamba-2 模块在 modules/mamba2.py 中实现。

简化版本在 modules/mamba2_simple.py 中。

用法类似于 Mamba(-1)：

from mamba_ssm import Mamba2
model = Mamba2(
    # 此模块大约使用 3 * expand * d_model^2 个参数
    d_model=dim, # 模型维度 d_model
    d_state=64,  # SSM 状态扩展因子，通常为 64 或 128
    d_conv=4,    # 局部卷积宽度
    expand=2,    # 块扩展因子
).to("cuda")
y = model(x)
assert y.shape == x.shape

SSD

内部 SSD 模块的最小版本（Mamba-2 论文中的清单 1）及其在"离散"和"连续"SSM 版本之间的转换在 modules/ssd_minimal.py 中实现。

Mamba 语言模型

最后，我们提供了一个完整语言模型的示例：深度序列模型主干（带有重复的 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 数据集上训练了 3000 亿个 token，以及 mamba-2.8b-slimpj（在 SlimPajama 数据集上训练了 6000 亿个 token）。

这些模型将由下面的生成脚本自动下载。

这些模型在 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 模块。）

注意：这些是仅训练了 3000 亿个 token 的基础模型，没有进行任何形式的下游修改（如指令调优等）。预期性能应与在类似数据上训练的其他架构相当或更好，但可能不及更大或经过微调的模型。

评估

要对模型进行零样本评估（对应论文表 3），我们使用 lm-evaluation-harness 库。

通过 pip install lm-eval==0.4.2 安装 lm-evaluation-harness。
运行评估（更多文档请参阅 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

从 Hugging Face Hub 自动加载模型，
生成用户指定提示的补全内容，
对此生成过程的推理速度进行基准测试。其他可配置选项包括 top-p（核采样）概率和 softmax 温度。

示例

测试不同采样策略下的生成延迟（例如批量大小 = 1）：

python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --prompt "我的猫为一个新的语言模型编写了所有这些 CUDA 代码，然后" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name "EleutherAI/pythia-2.8b" --prompt "我的猫为一个新的语言模型编写了所有这些 CUDA 代码，然后" --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 "我的猫为一个新的语言模型编写了所有这些 CUDA 代码，然后" --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 "我的猫为一个新的语言模型编写了所有这些 CUDA 代码，然后" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

故障排除

精度

我们的模型使用 PyTorch AMP 进行混合精度训练。AMP 将模型参数保持在 float32 中，并在必要时转换为半精度。另一方面，其他框架如 DeepSpeed 将参数存储为 float16，并在必要时进行上转（例如用于优化器累积）。

我们观察到，主模型参数可能需要更高的精度，因为 SSM 对其循环动态很敏感。如果您遇到不稳定情况，首先请尝试使用存储 fp32 参数的框架（如 AMP）。

初始化

模型的某些部分使用了继承自 S4 模型先前工作的初始化方法。例如，$\Delta$ 参数通过初始化其线性投影的偏置来获得目标范围。然而，一些框架可能有后初始化钩子（例如将 nn.Linear 模块中的所有偏置项设为零）。如果是这种情况，您可能需要添加自定义逻辑（例如，这一行在我们的训练器中关闭重新初始化，但在任何其他框架中都是无操作）这是特定于训练框架的。

AMD 卡的额外先决条件

修补 ROCm

如果您使用的是 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}
}