CentOS 7服务器上快速安装mamba函数库

本次预配置虚拟环境为cuda 11.8+torch 2.2.2+python 3.10

1. 创建conda虚拟环境：conda create -n mamba python=3.10

激活环境：conda activate mamba

2. 安装Pytorch环境：

   conda install pytorch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 pytorch-cuda=11.8 -c pytorch -c nvidia

3. 安装conv1d库(可自选版本)：

   pip install https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.0.post8/causal_conv1d-1.5.0.post8+cu11torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

4. 安装mamba_ssm库【可自选版本】：

   pip install https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu11torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

注：在选择causal_conv1d和mamba_ssm的时候，如果不知道自己的torch和cuda是什么版本可以打印输出：

import torch
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())

5. 测试：

   from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
   from mamba_ssm.ops.triton.selective_state_update import selective_state_update

6. 通过程序测试

mamba_ssm测试：

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(y)
print(y.shape)

输出：

tensor([[[-0.0035, 0.0065, 0.0156, ..., 0.0144, 0.0181, -0.0319],

0.0144, -0.0009, -0.0034, ..., 0.0666, 0.0301, 0.0096\], \[-0.0102, -0.0006, 0.0264, ..., -0.0478, 0.0402, -0.0818\], ..., \[-0.0177, -0.0073, 0.0390, ..., -0.0340, -0.0025, -0.0271\], \[ 0.0138, 0.0031, -0.0223, ..., -0.0278, 0.0072, -0.0143\], \[-0.0076, -0.0186, 0.0079, ..., -0.0062, 0.0021, -0.0283\]\], \[\[ 0.0079, 0.0089, -0.0111, ..., -0.0010, 0.0097, 0.0033\], \[ 0.0176, 0.0140, 0.0087, ..., -0.0223, 0.0294, -0.0158\], \[-0.0149, 0.0228, 0.0126, ..., 0.0382, -0.0015, -0.0182\], ..., \[ 0.0322, -0.0258, 0.0312, ..., -0.0045, -0.0036, -0.0509\], \[ 0.0164, 0.0001, -0.0060, ..., 0.0722, 0.0058, 0.0061\], \[ 0.0365, 0.0176, 0.0470, ..., 0.0103, 0.0257, 0.0039\]\]\], device='cuda:0', grad_fn=\) torch.Size(\[2, 64, 16\]) mamba_ssm 2测试： from mamba_ssm import Mamba2 import torch 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) assert y.shape == x.shape print(y) print(y.shape) tensor(\[\[\[-0.7895, -0.0058, -0.6374, ..., -0.1971, -0.1220, 0.9073\], \[-0.2546, 0.3763, -0.6773, ..., -0.2671, 0.4855, -0.7709\], \[ 0.4188, 0.3607, 0.7131, ..., 0.9439, -0.0798, -1.2252\], ..., \[-1.3136, -1.1002, -0.5330, ..., 1.5189, 0.2091, 1.0726\], \[-0.4520, -0.6626, -0.3810, ..., 0.3964, 0.0947, 0.7275\], \[-0.3024, -0.2375, 0.2435, ..., 0.4073, 0.4688, 0.6197\]\], \[\[ 0.0257, 0.7625, 0.7594, ..., 0.3531, 0.3276, 0.4292\], \[-1.0428, 0.8166, 0.1294, ..., 0.8236, 0.0515, -0.3141\], \[ 0.1267, -0.6214, 0.1667, ..., -0.4576, 0.7774, 0.7242\], ..., \[ 0.0609, -0.0283, 0.4718, ..., 0.7035, 0.2011, 0.2541\], \[ 0.3708, -0.0039, 0.2280, ..., 0.9191, -0.6267, -0.2572\], \[-0.2993, 0.0933, 0.2601, ..., 0.9123, -1.0403, 0.7488\]\]\], device='cuda:0', grad_fn=\) torch.Size(\[2, 64, 512\]) 参考： 1.[https://zhuanlan.zhihu.com/p/27156724975](https://zhuanlan.zhihu.com/p/27156724975 "https://zhuanlan.zhihu.com/p/27156724975") 2.[https://zhuanlan.zhihu.com/p/25916604332](https://zhuanlan.zhihu.com/p/25916604332 "https://zhuanlan.zhihu.com/p/25916604332")