这里配置的是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