Mamba

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

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

Mamba-3：利用状态空间原理改进序列建模
通过结构化状态空间对偶性
Aakash Lahoti*，Kevin Y. Li*，Berlin Chen*，Caitlin Wang*，Aviv Bick，J. Zico Kolter，Tri Dao†，Albert Gu†
论文：https://arxiv.org/abs/2603.15569

关于

Mamba 是一种新型状态空间模型架构，在语言建模等信息密集型数据上表现出令人鼓舞的性能，而此前的次二次复杂度模型则难以匹敌 Transformer。它基于 结构化状态空间模型 的研究进展，采用高效的硬件感知设计与实现，其精神与 FlashAttention 一脉相承。

安装

请先安装 PyTorch，然后执行以下步骤：

• 【可选】pip install causal-conv1d>=1.4.0 --no-build-isolation：这是 Mamba 模块内部使用的一种高效因果一维卷积层实现。
• pip install mamba-ssm --no-build-isolation：核心 Mamba 软件包。
• pip install mamba-ssm[causal-conv1d] --no-build-isolation：同时安装核心 Mamba 软件包和 causal-conv1d。

必须使用 --no-build-isolation 参数，以便 pip 使用您已有的支持 CUDA 的 PyTorch 版本，而不是在隔离的构建环境中安装 CPU 版本的 PyTorch。

注意：若要使用 Mamba-3，请从源码安装：MAMBA_FORCE_BUILD=TRUE pip install --no-cache-dir --force-reinstall git+https://github.com/state-spaces/mamba.git --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-3

Mamba-3 模块的实现位于 modules/mamba3.py。

使用方法如下：

from mamba_ssm import Mamba3
batch, length, dim = 2, 2048, 768
x = torch.randn(batch, length, dim).to(torch.bfloat16).to("cuda")
model = Mamba3(
    # 该模块大约占用 6 * d_model^2 个参数
    d_model=dim, # 模型维度 d_model
    d_state=128,  # SSM 状态大小
    headdim=64, # SSM 头维度
    is_mimo=True, # 使用 MIMO 模式
    mimo_rank=4, # 当 is_mimo=True 时的 MIMO 秩
    chunk_size=16, # 如果输入为 bf16，则为 64/mimo_rank；否则为 32/mimo_rank
    is_outproj_norm=False, # 是否添加 SSM 后的归一化层
    dtype=torch.bfloat16,
).to("cuda")
y = model(x)
assert y.shape == x.shape

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 库。

1. 使用 pip install lm-eval==0.4.2 安装 lm-evaluation-harness。
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

1. 自动加载来自 Hugging Face Hub 的模型，
2. 根据用户指定的提示生成文本补全，
3. 对该生成过程的推理速度进行基准测试。

其他可配置选项包括 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 及以上版本不需要此操作。

1. 找到您的 ROCm 安装目录。通常位于 /opt/rocm/，但具体位置可能因安装而异。

2. 应用补丁。如果遇到权限问题，请使用 sudo 运行。

    patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch 
   

引用

如果您使用此代码库，或以其他方式认为我们的工作有价值，请引用 Mamba：

@article{mamba,
  title={Mamba: 基于选择性状态空间的线性时间序列建模},
  author={Gu, Albert 和 Dao, Tri},
  journal={arXiv 预印本 arXiv:2312.00752},
  year={2023}
}

@inproceedings{mamba2,
  title={Transformer 就是 SSM：通过结构化状态空间对偶性实现通用模型和高效算法},
  author={Dao, Tri 和 Gu, Albert},
  booktitle={国际机器学习大会 (ICML)},
  year={2024}
}

@misc{lahoti2026mamba3improvedsequencemodeling,
      title={Mamba-3：利用状态空间原理改进序列建模}, 
      author={Aakash Lahoti、Kevin Y. Li、Berlin Chen、Caitlin Wang、Aviv Bick、J. Zico Kolter、Tri Dao 和 Albert Gu},
      year={2026},
      eprint={2603.15569},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2603.15569}, 
}

Mamba 快速上手指南

Mamba 是一种基于选择性状态空间（Selective State Space）的新型架构，专为处理高密度信息序列（如语言建模）设计。它在保持线性时间复杂度的同时，性能可媲美 Transformer。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux
• GPU: NVIDIA GPU (推荐) 或 AMD GPU (需额外配置)
• CUDA: 11.6 或更高版本
• PyTorch: 1.12 或更高版本 (必须预先安装且支持 CUDA)

注意: 安装过程中请务必使用 --no-build-isolation 参数，以确保 pip 使用您现有的 CUDA-enabled PyTorch 环境，而不是在隔离环境中重新安装 CPU 版本的 torch。

安装步骤

首先安装 PyTorch（此处以官方命令为例，国内用户可使用清华源或阿里源加速）：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

接着安装 Mamba 核心依赖及主包：

1. 安装因果卷积层 (可选但推荐):

   pip install causal-conv1d>=1.4.0 --no-build-isolation
   

2. 安装 Mamba 核心包:

   pip install mamba-ssm --no-build-isolation
   

   或者一次性安装核心包与因果卷积层:

   pip install mamba-ssm[causal-conv1d] --no-build-isolation
   

3. 安装 Mamba-3 (如需使用最新版本): Mamba-3 需要从源码强制构建：

   MAMBA_FORCE_BUILD=TRUE pip install --no-cache-dir --force-reinstall git+https://github.com/state-spaces/mamba.git --no-build-isolation
   

基本使用

Mamba 提供了不同版本的模块（Mamba, Mamba-2, Mamba-3）。以下是各版本最基础的使用示例。

1. 使用 Mamba (v1)

import torch
from mamba_ssm import Mamba

# 定义输入维度
batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")

# 初始化模型
model = Mamba(
    d_model=dim,   # 模型维度
    d_state=16,    # SSM 状态扩展因子
    d_conv=4,      # 局部卷积宽度
    expand=2,      # 块扩展因子
).to("cuda")

# 前向传播
y = model(x)
assert y.shape == x.shape

2. 使用 Mamba-2

Mamba-2 是改进版本，通常具有更大的状态维度。

from mamba_ssm import Mamba2

# 假设 x 已定义 (同上)
model = Mamba2(
    d_model=dim,   # 模型维度
    d_state=64,    # SSM 状态扩展因子 (通常为 64 或 128)
    d_conv=4,      # 局部卷积宽度
    expand=2,      # 块扩展因子
).to("cuda")

y = model(x)
assert y.shape == x.shape

3. 使用 Mamba-3

Mamba-3 引入了更多高级特性（如 MIMO 模式），建议使用 bfloat16 精度。

from mamba_ssm import Mamba3

batch, length, dim = 2, 2048, 768
x = torch.randn(batch, length, dim).to(torch.bfloat16).to("cuda")

model = Mamba3(
    d_model=dim,        # 模型维度
    d_state=128,        # SSM 状态大小
    headdim=64,         # SSM head 维度
    is_mimo=True,       # 启用 MIMO 模式
    mimo_rank=4,        # MIMO 秩
    chunk_size=16,      # 分块大小 (bf16 下通常为 64/mimo_rank)
    is_outproj_norm=False,
    dtype=torch.bfloat16,
).to("cuda")

y = model(x)
assert y.shape == x.shape

加载预训练模型进行推理

您可以直接使用 Hugging Face 上的预训练模型（如 state-spaces/mamba-2.8b）进行文本生成或评估。运行以下脚本可自动下载模型并测试生成速度：

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