Mamba：选择性状态空间模型训练与推理指南

本文基于官方仓库 state-spaces/mamba 的 README 完整翻译整理，覆盖 Mamba、Mamba-2、Mamba-3 三代模型的安装、接口、预训练权重、评测、推理与常见问题。所有命令、参数、版本要求与模型名均与官方 README 一致，未做编造。

一、项目官方信息

二、项目背景

Mamba 是一种全新的状态空间模型（State Space Model, SSM）架构，在语言建模这类信息密集的数据上展现出可与 Transformer 媲美的性能，而此前的各种亚二次复杂度模型都无法达到 Transformer 的水平。它基于 结构化状态空间模型（S4） 这条技术脉络，并以 FlashAttention 为精神做了面向硬件感知（hardware-aware）的高效设计与实现。所谓"选择性"，是指模型可以根据输入内容动态地决定保留哪些信息、遗忘哪些信息——这与传统 SSM 中参数固定的做法不同，也是 Mamba 能在语言建模上追平 Transformer 的关键；而"硬件感知"则借鉴 FlashAttention 的思路，把算法 kernel 写得贴合 GPU 内存层级，避免反复读写显存，从而在长序列上既省显存又快。相比 Transformer 的注意力机制随序列长度平方增长，Mamba 的选择性扫描在序列长度上是线性的，因此特别适合长文档、长代码、基因组序列等超长上下文场景。

仓库随三代论文演进：第一代 Mamba 提出"选择性状态空间"，让模型根据输入动态选择保留或遗忘信息，从而在语言建模上达到线性时间复杂度——对比 Transformer 随序列长度平方增长的注意力，长序列推理成本大幅降低；第二代 Mamba-2 提出"结构化状态空间对偶"（SSD），证明 Transformer 本质上是 SSM 的一种，并给出统一的高效算法；第三代 Mamba-3 以"推理优先"（inference-first）的状态空间原则改进序列建模，重点优化推理阶段的吞吐与显存。三篇论文的作者均为 Albert Gu 与 Tri Dao（Mamba-3 另含多位共同作者）。

三、环境要求与安装

安装前先装 PyTorch（建议较新的稳定版，以兼容 Mamba-3）。需要特别注意：默认情况下 mamba-ssm 只安装核心包，不会编译 selective_scan_cuda 扩展，也不会下载缓存的 CUDA 预编译 wheel。这种设计让只想先看代码、做 CPU 实验或在没有 NVIDIA GPU 的机器上也能快速装上；但要真正发挥 Mamba 的速度，就必须开启 CUDA selective scan 扩展。官方提供五种安装模式：

--no-build-isolation 是 CUDA 构建所必需的，这样 pip 才会使用你环境里已有的、带 CUDA 的 PyTorch，而不是在隔离构建环境里另装一个 CPU 版 torch——否则编译出的扩展找不到正确的 CUDA 头文件与库。

从源码安装使用同样的默认值与开关：

# 源码默认，不含 selective_scan_cuda
pip install . --no-build-isolation
# 从 GitHub 源码默认安装
pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation
# 源码强制本地核心构建
MAMBA_FORCE_BUILD=TRUE pip install . --no-build-isolation
# 源码按需开启 CUDA selective scan
MAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation
# 源码强制本地 CUDA 编译
MAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation

官方说明：要用最新源码树里的 Mamba-3，请从源码安装；默认源码安装不含 selective_scan_cuda，如需 CUDA selective scan 扩展加 MAMBA_KEEP_CUDA_BUILD=TRUE，再加 MAMBA_FORCE_BUILD=TRUE 则强制本地 CUDA 编译（跳过预编译 wheel 缓存）。这意味着想第一时间用上 Mamba-3 的新特性，不能只 pip install mamba-ssm，而要 git clone 后从源码装，并显式打开 CUDA 构建开关。

核心要求：Linux、Python 3.10+、PyTorch 1.12+。CUDA selective_scan_cuda 构建与 GPU 执行的附加要求：NVIDIA GPU、CUDA 11.6+。AMD 显卡见后文补丁说明。Windows 官方未列入支持范围，建议在 Linux 容器或 WSL2 中使用。

四、使用方法

仓库暴露多个层级的接口，从最底层的 SSM 算子到完整语言模型。

4.1 选择性 SSM 层

Mamba 基于一个选择性 SSM 层，这是论文的核心（第 3 节；算法 2）。源码见 ops/selective_scan_interface.py。这一层是整个仓库的底层算子，包含 CUDA 加速的选择性扫描实现；如果只想研究或定制 SSM 数学，可以从这里入手，而不必直接用上层 Mamba 块。

4.2 Mamba Block

本仓库的主模块是包装了选择性 SSM 的 Mamba 架构块，源码见 modules/mamba_simple.py。它把选择性 SSM、局部卷积与门控组合成一个可即插即用的模块，输入输出形状一致，方便像搭积木一样堆叠成网络。各参数含义：d_model 是模型维度，d_state 是 SSM 状态扩展因子（Mamba-1 默认 16），d_conv 是局部卷积宽度（默认 4），expand 是块扩展因子（默认 2）：

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

4.3 Mamba-2

Mamba-2 块在 modules/mamba2.py 实现，简化版在 modules/mamba2_simple.py。与 Mamba(-1) 相比，Mamba-2 把 SSM 状态扩展因子 d_state 从 16 提升到通常 64 或 128，以利用 SSD 结构带来的并行效率；其余参数含义与 Mamba(-1) 一致。用法与 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 模块的极简版本（Mamba-2 论文 Listing 1），带"离散"与"连续"SSM 版本之间的转换，在 modules/ssd_minimal.py，适合逐行研读 SSD 算法而不必面对工程化的复杂代码。

4.4 Mamba-3

Mamba-3 块在 modules/mamba3.py 实现，按"推理优先"原则重新设计，参数也与前两代不同：约使用 6 * d_model^2 个参数（Mamba-1/2 为 3 * expand * d_model^2），并新增 headdim、is_mimo、mimo_rank、chunk_size 等控制项。其中 is_mimo=True 启用 MIMO 模式，mimo_rank 控制 MIMO 秩，chunk_size 在 bf16 输入时取 64/mimo_rank、否则取 32/mimo_rank。用法如下：

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, # 若 x 为 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

4.5 Mamba 语言模型

仓库还提供一个完整语言模型示例：由重复 Mamba 块构成的深度序列模型主干 + 语言模型头，源码见 models/mixer_seq_simple.py。它演示了如何把 Mamba 集成进端到端神经网络，包括 embedding、堆叠 Mamba 块、最终归一化与分类头；下一节的生成脚本就基于它，可以直接拿来加载预训练权重做生成，省去自己拼装模型结构的工作，也是把 Mamba 当 backbone 做下游任务时最好的起点。

五、预训练模型

预训练模型已上传到 Hugging Face · state-spaces，包括 Mamba-1、Mamba-2 两个世代的一系列尺寸：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，它们都在 The Pile 上训练了 300B token；另有 mamba-2.8b-slimpj 在 SlimPajama 数据集上训练了 600B token。生成脚本会自动从 Hugging Face Hub 下载这些模型，无需手工管理权重文件，首次运行时联网拉取即可。

这些模型在 The Pile 上训练，采用 GPT-3 定义、也被许多开源模型沿用的标准模型尺寸：

注意：Mamba 的层数是相近 Transformer 的两倍，因为 Transformer 的每个"层"（MHA 块 + MLP 块）需要两个 Mamba 块。官方特别提醒：这些都是只训练了 300B token 的基座模型，没有任何形式的下游改造（指令微调等）；性能预期与在类似数据上训练的其他架构相当或更好，但不要指望它们能匹配更大的或已微调的模型。其中 mamba-2.8b-slimpj 是在 SlimPajama 上训练了 600B token 的版本，是开源中效果最好的 Mamba 基座；transformerpp-2.7b 与 mamba2attn-2.7b 则是 Mamba-2 论文中用于对照的变体模型，用来比较不同架构设计的优劣。

六、评测与推理

6.1 零样本评测

要跑论文表 3 对应的零样本评测，使用 lm-evaluation-harness。先安装指定版本的评测工具（lm-eval==0.4.2），再用 lm_eval 命令指定模型后端、预训练权重与任务列表：

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

复现博客中 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

官方提示：由于评测过程中的噪声，每个任务的结果可能与报告值相差 0.1–0.3，做横向对比时不必纠结小数点后一位的差异。任务列表中的 lambada_openai、hellaswag、piqa、arc_easy、arc_challenge、winogrande、openbookqa 是常见的零样本语言理解基准；mmlu 用 --num_fewshot 5 跑 5-shot 多任务知识评测，覆盖人文、社科、理工等多个学科方向。

6.2 生成与推理基准

脚本 benchmarks/benchmark_generation_mamba_simple.py 会：自动从 Hugging Face Hub 加载模型、按用户给定提示生成补全、并对生成速度做基准测试。其他可配置项包括 top-p（核采样）概率与 softmax 温度。--topp 控制核采样的保留概率，--temperature 控制输出随机性，--repetition-penalty 惩罚重复 token，--minp 与 --topk 是另一种最小概率/ top-k 采样控制。

测试生成延迟（如 batch size=1）并比较不同采样策略。下面三条命令分别演示默认 top-p 采样、以及 min-p 采样；其中也附了一条用同尺寸 Pythia（Transformer）模型做对照的命令，方便直接对比 Mamba 与 Transformer 的生成速度：

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

测试随机提示下的生成吞吐（如大 batch size）：

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

吞吐测试用随机提示、不指定 --prompt，目的是压测模型在大 batch 下的每秒生成 token 数；同样给出同尺寸 Pythia 做对照，方便量化 Mamba 相对 Transformer 的吞吐优势。用 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

七、常见问题 FAQ（Troubleshooting）

1. 训练不稳定怎么办？ 我们的模型用 PyTorch AMP 混合精度训练：AMP 把模型参数保持在 float32，必要时再转半精度；而 DeepSpeed 等框架把参数存在 float16、必要时再上转。我们观察到，由于 SSM 对其循环动力学敏感，主模型参数可能需要更高精度。如果遇到不稳定，第一步请尝试使用把参数存在 fp32 的框架（如 AMP），而不是盲目调小学习率或减少 batch。
2. 初始化不匹配？ 模型部分模块的初始化继承自 S4 的前期工作，例如通过初始化线性投影的偏置，让 Δ 参数落在目标范围内（见 mamba_simple.py 第 102 行附近）。但有些框架会有"初始化后钩子"（例如把 nn.Linear 里所有偏置清零），这会把精心设计的 Δ 初始化冲掉。如果遇到这种情况，可能需要添加针对自己训练框架的自定义逻辑——例如仓库在自己的 trainer 里有一行代码关闭了这种重新初始化，但在其他框架里这是空操作，需要按框架补写。
3. AMD 显卡编译报错？ 如果你用 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

4. 为什么默认安装不含 CUDA 加速？ 这是官方刻意设计的默认行为：pip install mamba-ssm --no-build-isolation 只装纯 Python 核心，方便无 GPU 环境或只读源码；需要 CUDA selective scan 时显式设置 MAMBA_KEEP_CUDA_BUILD=TRUE，它会先找预编译 wheel，找不到再本地编译。生产环境建议直接走 CUDA 构建，否则选择性扫描会退化为慢得多的纯 Python 实现。
5. Mamba、Mamba-2、Mamba-3 怎么选？ 新代码建议直接用 Mamba-3（需从源码安装，默认 wheel 不含 Mamba-3）；Mamba-2 提供 SSD 统一视角且有成熟权重（130M 到 2.7B 全规格开箱即用）；Mamba(-1) 适合复现早期论文结果或兼容已有历史代码。三者接口风格一致，迁移成本主要在参数名与维度上，迁移时建议对照官方示例代码逐步调整。
6. 预训练模型能直接对话吗？ 不能。这些都是基座模型，只做了 300B（slimpj 为 600B）token 预训练，未经指令微调；要对话需要自己后训练（SFT）或选社区微调版，直接拿来提问会输出续写而非回答。
7. 如何引用？ 如果使用本代码库或觉得我们的工作有价值，请分别引用三篇论文：Mamba (arXiv:2312.00752)、Mamba-2 (ICML 2024)、Mamba-3 (arXiv:2603.15569)，BibTeX 见仓库引用章节，按所使用的模型代际对应引用即可。

八、跳转链接

如果你想复现论文结果、自己训练或把 Mamba 集成进下游任务，建议从三篇论文与官方实现源码入手；遇到编译或训练问题，先对照本文 Troubleshooting 章节检查精度与初始化设置，AMD 用户记得打 ROCm 6.0 补丁。以下是本文整理涉及的全部入口：

• 仓库主页：github.com/state-spaces/mamba
• 权重：Hugging Face · state-spaces
• 论文：Mamba、Mamba-2、Mamba-3
• 上游：S4、FlashAttention
• 评测：lm-evaluation-harness
• 示例数据：The Pile