Linux服务器Mamba2安装及example运行问题

Mamba 从发表到现在已经有段时间了，先前的Mamba代码有些地方不完善或者需要对源码做些修改后编译，最近可能需要用到Mamba，发现之前实现的Vision Mamba块，注释掉了一些源码可能会导致训练速度下降，然后官方实现了Mamba2，这里尝试安装，做一些记录，防止服务器重置🫤

CUDA 安装

在安装其它东西之前，需要安装CUDA，这里参照了 Vision Mamba 的环境安装 cuda-11.8 版本，

直接到 CUDA 官网找对应版本就行 CUDA官网

在服务器上运行这两行代码，然后安装CUDA，网上有很多教程我这里服务器已经装好了，不在赘述

wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
sudo sh cuda_11.8.0_520.61.05_linux.run

Conda 创建环境

Anaconda 创建一个虚拟环境配置mamba，创建并激活，同样，这里创建python 3.10 的环境

conda create -n your_env_name python=3.10.13
conda activate your_env_name

Pytorch 安装

pytorch 这里安装2.1.1 版本，直接进入官网去找相应命令就行，PyTorch官网

找到对应的命令就行，我这里使用 pip 安装，安装会有点慢，可以替换清华源安装。

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

pip 清华源临时使用

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple some-package

永久替换

python -m pip install --upgrade pip
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

causal-conv1d安装

Mamba 依赖causal-con1d 库，这里可以直接使用 pip 安装

pip install causal-conv1d

这里可以根据需要自己指定版本，但部分人可能在这里安装报错，

可以通过GitHub 下载到本地安装，causal-conv1d GitHub

找到对应的版本然后下载到服务器上或者复制链接使用wget 命令下载，再使用 pip 命令安装，这里以 1.4.0版本为例，对应 CUDA 11.8， torch-2.1.1，python-3.10

wget https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.4.0/causal_conv1d-1.4.0+cu118torch2.1cxx11abiTRUE-cp310-cp310-linux_x86_64.wh

pip install causal_conv1d-1.4.0+cu118torch2.1cxx11abiTRUE-cp310-cp310-linux_x86_64.whl

当然也可以将GitHub 项目下载到本地编译

git clone https://github.com/Dao-AILab/causal-conv1d.git
cd causal-conv1d
CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install .

网上操作较多，这里不在赘述

Mamba 安装

Mamba 安装也和 causal-conv1d 一样的几种方法，这里懒得写了，我直接 pip 安装成功了，由于需要编译安装过程需要编译些东西，会需要等很久，但我这里运行的时候出了点问题。Mamba Github

问题解决

Mamba 的 Github 给了两端代码，用于测试是否安装成功

Mamba

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

Mamba2

from mamba_ssm import Mamba2
model = Mamba2(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=64,  # SSM state expansion factor, typically 64 or 128
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

这里第一段代码直接运行成功，但测试 Mamba2 时报错

AttributeError: 'NoneType' object has no attribute 'causal_conv1d_fwd'

这里我尝试了 Mamba 给的其它测试代码（需要下载预训练模型，需要一段时间，也可以根据需要修改 --model-name 参数）

Mamba

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 "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

Mamba2

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

这里Mamba2 测试也是出现了一些问题，会出现一个 dconv 和 d_conv 的问题，这里我尝试修改了源码但仍有问题，于是重新安装了 causal-conv1d，Mamba Github issue 中有人指出需要 causal-conv1d>=1.2.0，但我安装为 1.4.0 仍有问题，我猜可能我从GitHub 中下载的版本是 1.4.0 的问题，于是重新安装。

然后 下面这个测试就调通了

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

但上边的测试还没有成功，出现以下错误

RuntimeError: causal_conv1d with channel last layout requires strides (x.stride(0) and x.stride(2)) to be multiples of 8

就是 causal_conv1d要求步幅(x.s arstride(0)和x.s arstride(2))为8的倍数

网上去找了解决方法，主要是说要求 d_model * expand / headdim 是 8 的 倍数

修改成如下代码 重新运行就成功了

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
print("Mamba result",y.shape)
assert y.shape == x.shape

import torch
from mamba_ssm import Mamba2
batch, length, dim = 2, 64, 512
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba2(
    # This module uses roughly 3 * expand * d_model^2 parameters
    # make sure d_model * expand / headdim = multiple of 8
    d_model=dim, # Model dimension d_model
    d_state=64,  # SSM state expansion factor, typically 64 or 128
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
    headdim=64,  # default 64
).to("cuda")
y = model(x)
print("Mamba2 result",y.shape)
assert y.shape == x.shape