mamba-minimal

使用 PyTorch 单文件实现的简单、极简版 Mamba。

特点：

• 前向传播和反向传播的输出与官方实现数值等效
• 代码经过简化、可读性强并带有注释

不包含：

• 速度优化。官方实现经过了深度优化，这些优化是 Mamba 论文的核心贡献。为保持可读性，本实现大部分代码保持简单。
• 完整的参数初始化（虽然这可以在不牺牲可读性的前提下添加）

演示

查看 demo.ipynb 获取提示补全的示例。

from model import Mamba
from transformers import AutoTokenizer

model = Mamba.from_pretrained('state-spaces/mamba-370m')
tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')

generate(model, tokenizer, 'Mamba is the')

Mamba is the world's longest venomous snake with an estimated length of over 150 m. With such a large size and a venomous bite, Mamba kills by stabbing the victim (which is more painful and less effective than a single stab of the bite)

150 米... 🫢 太吓人了！

参考文献

Mamba 架构由 Albert Gu 和 Tri Dao 在论文 Mamba: Linear-Time Sequence Modeling with Selective State Spaces 中提出。

官方实现位于：https://github.com/state-spaces/mamba/tree/main

mamba-minimal 快速上手指南

环境准备

• 系统要求：支持 Python 的主流操作系统（Linux / macOS / Windows）
• 前置依赖：
  • Python 3.8 或更高版本
  • PyTorch 1.12 或更高版本（需根据你的 CUDA 版本或 CPU 环境安装）
  • Git（用于克隆仓库）

安装步骤

1. 克隆仓库

   git clone https://github.com/<原作者仓库地址>/mamba-minimal.git
   cd mamba-minimal
   

2. 创建并激活 Python 虚拟环境（推荐）

   python -m venv venv
   # Linux/macOS
   source venv/bin/activate
   # Windows
   venv\Scripts\activate
   

3. 安装 Python 依赖包

   pip install torch transformers
   

   （可选）如需加速下载，可使用国内 PyPI 镜像源：

   pip install torch transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

以下是一个最简单的文本生成示例，展示了如何加载预训练模型并进行推理。

1. 确保你已进入项目根目录，并已激活虚拟环境。

2. 创建 Python 脚本（例如 demo.py），内容如下：

   from model import Mamba
   from transformers import AutoTokenizer
   
   # 加载预训练模型和分词器
   model = Mamba.from_pretrained('state-spaces/mamba-370m')
   tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')
   
   # 定义一个简单的生成函数（需自行实现或参考项目中的 demo.ipynb）
   def generate(model, tokenizer, prompt, max_length=50):
       inputs = tokenizer(prompt, return_tensors="pt")
       outputs = model.generate(**inputs, max_length=max_length)
       return tokenizer.decode(outputs[0], skip_special_tokens=True)
   
   # 执行生成
   result = generate(model, tokenizer, 'Mamba is the')
   print(result)
   

3. 运行脚本

   python demo.py
   

   预期会输出一段由模型生成的关于“Mamba”的文本。

注意：此 mamba-minimal 项目侧重于代码的简洁性与可读性，其运行速度远未经过优化。如需高性能版本，请参考官方实现 state-spaces/mamba。