mamba.py 🐍：一个简单高效的Mamba实现

这是一个基于PyTorch的 Mamba的直接实现，采用了简单的并行扫描实现方式。与顺序实现相比，该实现可显著提升速度，因为并行扫描允许在时间维度上进行并行化处理。

它兼具易读性和良好的训练性能。目前仅实现了部分其他功能，如Jamba、Vision Mamba以及muP。

更新内容

• 2024年3月8日：新增了针对Mamba和Mamba2的muP实现。这使得用户可以在小型模型上搜索最优超参数，并直接将其迁移到大型模型上。请参阅此PR。

• 2024年7月23日：mamba.py现已作为Hugging Face的transformers库的一部分。请参阅此PR。

• 2024年6月27日：已在PyPI上发布了mamba.py的软件包版本，您可以通过pip install mambapy进行安装。

• 2024年4月21日：新增了jamba.py文件，实现了 Jamba架构（Mamba与注意力层的混合）。同时，官方的CUDA实现也被添加为一种可能的后端。

• 2024年3月30日：更新了推理函数，现在支持自定义采样温度和batch_size。

• 2024年2月9日：性能更新的第一部分。对于长度小于128的小序列，与第一版相比，训练速度可提升超过20%。对于接近实际应用的设置（如NLP领域），训练速度可提升10%。请参阅此PR。

• 2024年1月22日：新增了mamba.py的MLX版本，支持推理和训练。该版本与PyTorch类似，让Mac用户也能轻松玩转Mamba模型。该版本已在迄今为止训练过的最大规模Mamba模型（28亿参数）上进行了测试。

• 2024年1月17日：新增了用于推理的step函数。它使用了Mamba的“RNN形式”来大幅加速推理过程。

---

概览

这张图展示了单个Mamba层（d_model=16, d_state=16）的训练时间（前向和后向传播），分别采用三种方法：CUDA——即官方的Mamba实现，mamba.py——即本仓库的实现，以及sequential——即选择性扫描的顺序（类似RNN）实现。

本仓库包含了一个简单易懂的代码实现，以纯PyTorch及MLX语言实现了 Mamba架构。您还可以尝试Jamba模型，该模型结合了Mamba与注意力层。本仓库的主要目标是教育用途。

仓库组织结构如下：

• 📁 mambapy：Mamba的PyTorch实现
  • pscan.py：Blelloch并行扫描的PyTorch实现
  • mamba.py：Mamba模型，如论文所述。在数值上完全等价（初始化、前向和后向传播）。
  • mamba2.py（beta）：Mamba-2模型，如论文所述。目前仅需CUDA支持，因为它仅基于原始版本进行了适配。
  • lm.py：封装Mamba(-2)模型，以便将其用作语言模型
  • jamba.py：在PyTorch中实现 Jamba模型
  • vim.py：实现 Vision Mamba。
  • 📁 onnx：将训练好的Mamba模型导出为ONNX格式，用于推理。
• 📁 mlx：基本与上述相同，但采用MLX语言实现。
• 📁 docs：包含代码注释说明的文件夹，目前重点在于并行扫描。
• 📁 examples：两个使用Mamba模型的PyTorch示例，以及一个训练文件。

muP已实现，并兼容两种Mamba模型（详情见下文）。

使用方法

您可以下载本仓库，也可以通过pip install mambapy进行安装。

最基本的用法是使用Mamba对象（mamba.py），它根据配置实现了一个简单的Mamba模型。无需嵌入层或头部：输入为(B, L, D)，输出也为(B, L, D)。

import torch
from mambapy.mamba import Mamba, MambaConfig

config = MambaConfig(d_model=16, n_layers=2)
model = Mamba(config)

B, L, D = 2, 64, 16
x = torch.randn(B, L, D)
y = model(x)

assert y.shape == x.shape

您也可以通过导入mamba2.py中的Mamba2Config和Mamba2对象来使用Mamba-2。

类LM（lm.py）基于Mamba或Mamba-2对象，提供经典的语言模型API。使用方法如下：

from mambapy.lm import LM, MambaConfig

config = MambaConfig(d_model=16, n_layers=4) # 核心模型
model = MambaLM(config, vocab_size=32000) # 封装为LM

x = torch.randint(high=32000, size=(16, 64))
logits = model(x) # (B, L, vocab_size)

它简单地封装了一个Mamba(-2)对象，加上嵌入层、最终归一化和语言建模头部。

您可以直接使用预训练的Mamba模型：

from mambapy.lm import from_pretrained
from transformers import AutoTokenizer

model = from_pretrained('state-spaces/mamba-130m').to("cuda")
tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')

output = model.generate(tokenizer, "Mamba is a type of")

这是mamba.py模块的结构：

Jamba

您也可以训练和运行Jamba模型的推理。请查看jamba.py文件，它构建了一个Jamba对象，将来自mamba.py的Mamba层与注意力层交错排列。

这是jamba.py中模块的结构：

API与Mamba和MambaLM模型相同。 您可以这样加载预训练的Jamba模型：

from mambapy.jamba_lm import from_pretrained
from transformers import AutoTokenizer

model = from_pretrained('TechxGenus/Mini-Jamba').to("cuda")
tokenizer = AutoTokenizer.from_pretrained('TechxGenus/Mini-Jamba')

output = model.generate(tokenizer, "def min(arr):")

📁 示例

目前有两个基本示例可供使用（部分可能已过时）：

• example_llm.ipynb：加载一个带有预训练权重的Mamba模型（从HuggingFace获取，参数量从1.3亿到28亿不等）
• example_e2e_training.ipynb：一个端到端训练示例，其中Mamba模型被用作一个简单3x3网格游戏的世界模型（训练尚未完成，模型规模应更大）。

如果您想要一个完整的训练示例（如llama2.c中所示），可以查看我创建的othello_mamba仓库。借助这个仓库，您可以从零开始训练Mamba或Jamba模型，使用bfloat16，轻松地与Transformer结合，自定义数据集等等。

muP

muP 是一种技术，能够将超参数（例如学习率）从小规模模型迁移到大规模模型。例如，有研究表明，可以将200万参数模型的学习率迁移到100亿参数模型上，且二者保持一致。这在实际超参数搜索中极为有用：您只需在小规模模型上进行扫描，找到最佳超参数组合，既快速又经济，然后自动获得适用于大规模模型的最佳超参数。

muP之所以能做到这一点，是因为它以特定方式初始化并缩放了模型权重的学习率。这些修改的效果如下：

如果不使用muP，结果是：

这里展示的是不同宽度（d_model）下激活值的变化情况，从t=1（初始化）到t=5（训练5个步骤）。在标准参数化（SP）下，网络的激活值随宽度变化很大；而在muP下，激活值则几乎不受宽度影响。

直观地说，如果网络的激活值（“信号”）无论宽度如何都保持一致，那么我们很容易想象，最优超参数也就不受宽度影响了。

这正是我们在实践中观察到的现象：当我们扫描寻找最佳学习率时：

在标准参数化下，随着模型变大，最佳学习率会随之调整；而使用muP时，最佳学习率则大致保持不变。小模型只有17.2万个参数，而大模型却超过1亿！

如需了解更多关于muP的一般信息，可以阅读论文；若想了解我对Mamba的muP实现的推导过程以及代码中具体的变化，请参阅相关PR。

---

性能

本节提供了mamba.py与官方Mamba实现之间更全面的性能对比。 总体而言，正如本文第一张图所示，两者在序列长度上的渐近性能大致相同。您可以将mamba.py视为普通的Transformer实现，而官方Mamba实现更接近FlashAttention v1。两者各有优势。

不过，这两种实现是否在其他参数上的渐近性能也相同呢？

d_model的渐近性能

我们可以看到，随着d_model增大，两种实现的表现完全一致。两者的差距大致保持不变。（总体来说，mamba.py慢约2倍）

d_state的渐近性能

这张图非常重要。我们发现，随着d_state增大，两者的渐近性能并不相同。提醒一下，d_state，即论文中的$N$，是状态扩展因子：输入的每个通道会被扩展为隐藏状态的$N$个通道。

注意：CUDA版本似乎并未受到d_state增大的影响。这是因为基准测试是在批大小为1的情况下进行的：GPU并未达到满负荷，因此d_state增大的影响并不明显。同样的情况也会出现在小型模型或短输入长度的情况下。请参阅此问题。

这在实际中有意义吗？截至目前，所有预训练的Mamba模型（最高达28亿参数）都使用了d_state=16，因此这种d_state带来的性能变化在此情况下并不重要。由于d_state本不该增长（与序列长度或d_model不同），这并不是灾难性的结果，但值得考虑。

然而，将这一观察与Albert Gu和Tri Dao在Mamba论文中的说法联系起来就很有意思：核心思想是利用现代加速器（GPU）的特性，仅在内存层次结构中更高效的层级上显式化状态ℎ。 他们还在附录D中描述了选择性扫描的主要数据移动过程：主要在SRAM中工作，可将内存读写次数减少$O(N)$倍。这解释了我们在这里看到的不同渐近行为。

当d_state=16（如state-spaces/mamba-2.8b-slimpj中所用）时，两者的差距相对较小；但当d_state=64（目前未用于任何模型）时，差距会显著拉大。（请注意第二张图中的OOM问题）

以上所有图表均是在A100 80GB上，以批大小为1计算得出的。 这是单个Mamba块前向和后向传播的综合度量。

之前的分析表明，内核融合的重要性：它可将内存访问次数减少$O(N)$，从而加快整个过程。

但同时也要考虑内存需求：官方Mamba实现的后向传播采用了重新计算策略——与其在内存中保存前向传播过程中计算出的激活值，不如在需要时直接在后向传播中重新计算。这大大降低了Mamba模型在训练时的内存需求。而本仓库并未实现这一功能。

因此，本仓库实现了Mamba论文中提到的三种技术之一，即所谓的“硬件感知选择性扫描”：并行扫描。 我们将探讨内核融合如何影响速度，同时讨论重新计算对内存需求的影响。

---

资源与更多信息

• Mamba论文：介绍了本仓库中实现的Mamba架构，该架构能够以线性时间对序列进行建模。
• Mamba实现，采用PyTorch编写，但使用了CUDA编写的并行扫描。这是速度最快的版本。
• Mamba的最小PyTorch实现，将扫描操作实现为一个顺序循环（其性能略逊于第一张图中的“顺序”曲线）。这段代码基本遵循官方Mamba实现中的此文件，但用torch.nn.Conv1d替换了CUDA卷积，并用顺序循环替换了CUDA编写的选择性扫描。本仓库的代码结构也沿用了这两个文件的组织方式。
• 前缀和及其应用，作者：Guy E. Blelloch（1993）。
• 对序列长度并行化线性递归神经网络：通过在序列上应用并行扫描，以消除顺序for循环。
• x.com/francoisfleuret：原始pscan实现。

待办事项

• 完成文档
• Mamba 2
• 清理vim.py
• 根据性能更新，更新性能图表
• 绘制三种不同Mamba实现（官方、朴素、mamba.py）的训练内存消耗图
• 使用ThunderKittens实现pscan？
• Jamba？推理和/或微调？
• 增加d_model后的更多测试（添加性能章节）
• 一个阶跃函数，用于（自回归）推理。
• 一个训练函数，类似于llama2.c

性能相关：

• 展开pscan.py中的for循环，以获得更好的性能（参见François Fleuret的pscan）（尽管这会牺牲一些可读性）
• 专门为反向传播编写一个逆向并行扫描。（目前，我们需要在扫描前后翻转数组）
• 某种程度上降低内存占用（如有必要，可以牺牲速度）
• 使用torch.compile()。据我测试，目前它还无法正常工作。似乎它不认可自定义的PScan自动求导函数。需要进一步调查。（参见PR#1)

引用

如果您在研究中发现本项目有用并希望引用，请使用以下BibTex条目：

@software{mambapy,
  author = {Alexandre Torres--Leguet},
  title = {mamba.py: 一个简单、可 hack 且高效的纯PyTorch和MLX版Mamba实现},
  url = {https://github.com/alxndrTL/mamba.py},
  version = {1.0},
  year = {2024},
}

mamba.py 快速上手指南

环境准备

• 系统要求：Linux / macOS / Windows（推荐 Linux）
• Python 版本：≥ 3.8
• PyTorch：≥ 2.0（推荐使用 CUDA 11.8 或 12.1）
• 推荐硬件：NVIDIA GPU（支持 CUDA），或 Apple Silicon（支持 MLX 版本）

国内用户可使用清华源加速安装：
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple mambapy

安装步骤

# 通过 PyPI 安装（推荐）
pip install mambapy

# 或从源码安装（如需修改代码）
git clone https://github.com/alxndrTL/mamba.py
cd mamba.py
pip install -e .

基本使用

1. 使用基础 Mamba 模型

import torch
from mambapy.mamba import Mamba, MambaConfig

# 配置模型
config = MambaConfig(d_model=16, n_layers=2)
model = Mamba(config)

# 输入输出：(B, L, D)
B, L, D = 2, 64, 16
x = torch.randn(B, L, D)
y = model(x)

assert y.shape == x.shape  # 输出维度与输入一致

2. 使用语言模型封装（推荐用于实际任务）

from mambapy.lm import LM, MambaConfig
import torch

config = MambaConfig(d_model=16, n_layers=4)
model = LM(config, vocab_size=32000)  # 封装为语言模型

x = torch.randint(high=32000, size=(16, 64))  # 输入 token ID
logits = model(x)  # 输出: (B, L, vocab_size)

3. 加载预训练模型（推荐）

from mambapy.lm import from_pretrained
from transformers import AutoTokenizer

# 加载预训练 Mamba-130M
model = from_pretrained('state-spaces/mamba-130m').to("cuda")
tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')

# 生成文本
output = model.generate(tokenizer, "Mamba 是一种")
print(output)

✅ 支持 Mamba2、Jamba 模型，使用方式类似，只需导入对应类（如 Mamba2、JambaLM）即可。