mamba-architecture

Mamba - 选择性状态空间模型

快速开始

Mamba 是一种状态空间模型架构，在序列建模中实现了 O(n) 线性复杂度。

安装：

# Install causal-conv1d (optional, for efficiency)
pip install causal-conv1d>=1.4.0

# Install Mamba
pip install mamba-ssm
# Or both together
pip install mamba-ssm[causal-conv1d]

先决条件：Linux、NVIDIA GPU、PyTorch 1.12+、CUDA 11.6+

基本用法（Mamba 块）：

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,      # Model dimension
    d_state=16,       # SSM state dimension
    d_conv=4,         # Conv1d kernel size
    expand=2          # Expansion factor
).to("cuda")

y = model(x)  # O(n) complexity!
assert y.shape == x.shape

常见工作流

工作流 1：使用 Mamba-2 构建语言模型

完整语言模型及生成：

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from mamba_ssm.models.config_mamba import MambaConfig
import torch

# Configure Mamba-2 LM
config = MambaConfig(
    d_model=1024,           # Hidden dimension
    n_layer=24,             # Number of layers
    vocab_size=50277,       # Vocabulary size
    ssm_cfg=dict(
        layer="Mamba2",     # Use Mamba-2
        d_state=128,        # Larger state for Mamba-2
        headdim=64,         # Head dimension
        ngroups=1           # Number of groups
    )
)

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)

# Generate text
input_ids = torch.randint(0, 1000, (1, 20), device="cuda", dtype=torch.long)
output = model.generate(
    input_ids=input_ids,
    max_length=100,
    temperature=0.7,
    top_p=0.9
)

工作流 2：使用预训练 Mamba 模型

从 HuggingFace 加载：

from transformers import AutoTokenizer
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

# Load pretrained model
model_name = "state-spaces/mamba-2.8b"
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")  # Use compatible tokenizer
model = MambaLMHeadModel.from_pretrained(model_name, device="cuda", dtype=torch.float16)

# Generate
prompt = "The future of AI is"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
output_ids = model.generate(
    input_ids=input_ids,
    max_length=200,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.2
)
generated_text = tokenizer.decode(output_ids[0])
print(generated_text)

可用模型：

• state-spaces/mamba-130m
• state-spaces/mamba-370m
• state-spaces/mamba-790m
• state-spaces/mamba-1.4b
• state-spaces/mamba-2.8b

工作流 3：Mamba-1 与 Mamba-2

Mamba-1（较小状态）：

from mamba_ssm import Mamba

model = Mamba(
    d_model=256,
    d_state=16,      # Smaller state dimension
    d_conv=4,
    expand=2
).to("cuda")

Mamba-2（多头、较大状态）：

from mamba_ssm import Mamba2

model = Mamba2(
    d_model=256,
    d_state=128,     # Larger state dimension
    d_conv=4,
    expand=2,
    headdim=64,      # Head dimension for multi-head
    ngroups=1        # Parallel groups
).to("cuda")

主要区别：

• 状态大小：Mamba-1（d_state=16）与 Mamba-2（d_state=128）
• 架构：Mamba-2 采用多头结构
• 归一化：Mamba-2 使用 RMSNorm
• 分布式：Mamba-2 支持张量并行

工作流 4：与 Transformers 基准对比

生成速度对比：

# Benchmark Mamba
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "state-spaces/mamba-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

# Benchmark Transformer
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "EleutherAI/pythia-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

预期结果：

• Mamba：推理速度快 5 倍
• 内存：无需 KV 缓存
• 扩展性：随序列长度线性增长

何时选择 Mamba 而非替代方案

以下情况适合使用 Mamba：

• 需要处理长序列（100K+ 个 token）
• 希望比 Transformers 推理更快
• 内存受限（无需 KV 缓存）
• 构建流式应用
• 需要线性扩展

优势：

• O(n) 复杂度：线性而非二次
• 推理速度快 5 倍：没有注意力开销
• 无需 KV 缓存：降低内存用量
• 百万 token 序列：硬件效率高
• 流式处理：每个 token 占用固定内存

以下情况可考虑其他方案：

• Transformers：需要一流性能，且算力充足
• RWKV：希望使用 RNN 与 Transformer 的混合架构
• RetNet：需要基于保留机制的架构
• Hyena：希望采用基于卷积的方法

常见问题

问题：CUDA 显存不足

减小批次大小或使用梯度检查点：

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)
model.gradient_checkpointing_enable()  # Enable checkpointing

问题：安装缓慢

安装二进制 wheel（而非源码）：

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

问题：缺少 causal-conv1d

单独安装：

pip install causal-conv1d>=1.4.0

问题：无法从 HuggingFace 加载模型

使用 MambaLMHeadModel.from_pretrained（而非 AutoModel）：

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b")

进阶主题

选择性 SSM：参阅 references/selective-ssm.md，了解数学表述、状态空间方程，以及选择性如何实现 O(n) 复杂度。

Mamba-2 架构：参阅 references/mamba2-details.md，了解多头结构、张量并行和分布式训练设置。

性能优化：参阅 references/performance.md，了解硬件感知设计、CUDA 内核和内存效率优化技术。

硬件要求

• GPU：支持 CUDA 11.6+ 的 NVIDIA GPU
• 显存：
  • 130M 模型：2GB
  • 370M 模型：4GB
  • 790M 模型：8GB
  • 1.4B 模型：14GB
  • 2.8B 模型：28GB（FP16）
• 推理：比 Transformers 快 5 倍
• 内存：无需 KV 缓存（用量低于 Transformers）

性能（与 Transformers 相比）：

• 速度：推理速度快 5 倍
• 内存：减少 50%（无需 KV 缓存）
• 扩展性：线性而非二次

资源

• 论文（Mamba-1）：https://arxiv.org/abs/2312.00752（2023 年 12 月）
• 论文（Mamba-2）：https://arxiv.org/abs/2405.21060（2024 年 5 月）
• GitHub：https://github.com/state-spaces/mamba ⭐ 13,000+
• 模型：https://huggingface.co/state-spaces
• 文档：仓库 README 和 wiki

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 明确输入。 根据数据或用户确认格式、标识符和参数。任何会影响结果的值都要询问，不要猜测。
3. 先运行小规模版本。 执行冒烟测试：使用小型模型或子集，运行少量步骤，并确认损失下降且检查点成功保存。
4. 按照以上说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 汇报。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果发生这种情况	请这样做
CUDA 内存不足	减小批次大小，启用梯度累积/检查点或混合精度，或对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否有无效值，并启用梯度裁剪。
这些说明中的函数、标志或端点在已安装版本中缺失	查阅已安装版本自己的文档（help()、--help、官方文档），进行调整，并告知用户。切勿编造 API。
所需输入、标识符或参数含糊不清	询问用户，或在运行前明确说明所作假设。

诚信原则

• 切勿捏造结果、参数、标识符、引用或统计数据。如果无法运行或验证某项内容，请明确说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前先估算计算成本。
• 此处与版本相关的细节可能已过时：请对照已安装版本的官方文档进行确认。
• 对于需要付费、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征得同意。

相关技能

• rwkv-architecture：介绍 RWKV（Receptance Weighted Key Value）架构，这是一种 RNN/Transformer 混合架构，具有 O(n) 推理能力且无需 KV 缓存，包括 RWKV…
• optimizing-attention-flash：通过 PyTorch 原生的 scaled_dot_product_attention（PyTorch 2.2+）或 flash-attn 库，为 transformer 模型启用 Flash Attention，…