Mamba环境配置---（Ubuntu）

这里配置的是ubuntu的Mamba(-1 and -2)环境。

之前在windows下按照网上的步骤配置了mamba环境，但是发现不是cuda算子下运行（应该是的），后来还发现即使没用cuda加速，但只能运行mamba-1版本，mamba2不能运行，可测试下面的代码，是否能运行mamba2：

注：这代码从mamba官方代码中选取

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

运行完后发现报错：

 因此开始在ubuntu系统进行环境配置！！！

---

Ubuntu环境配置：

1、创建自己的conda环境

conda create -n mamba_demo python=3.10.13

注：这里的 -n 参数设置为mamba_demo（可以取自己的名字）。

2、进入环境

conda activate mamba_demo

3、创建torch环境和对应的cuda11.8版本

conda install pytorch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 pytorch-cuda=11.8 -c pytorch -c nvidia

4、继续安装packaging和cuda-nvcc

conda install -c "nvidia/label/cuda-11.8.0" cuda-nvcc
conda install packaging

5、 直接下载causal-conv1d和mamba-ssm的.whl文件，并安装：

这里直接在windows下通过迅雷下载的，ubuntu从github下载太慢。

这里的causal-conv1d和mamba-ssm的.whl版本分别是1.4.0和2.2.2。

（注：这里直接pip install 或者下载.zip再pip install . 会出现报错要么卡住，所以最好下载.whl进行下载）

causal-conv1d安装：

pip install causal_conv1d-1.4.0+cu118torch2.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

mamba-ssm安装：

pip install mamba_ssm-2.2.2+cu118torch2.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

6、代码测试：

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
print('!!!!!')

import torch
from mamba_ssm import Mamba2

dim = 512
x = torch.randn(batch, length, dim).to("cuda")
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
print('!!!!!')

测试结果：

---

附.whl还有测试的代码压缩包

链接: https://pan.baidu.com/s/1DCV5LM-VFOMtLh4cyhqUMg?pwd=1234 提取码: 1234

---

也可以自行从github上下载对应的.whl