Windows和Linux系统上的Mamba_ssm环境配置

一、Linux系统安装

如果自己的系统不是cuda11.8，那么需要先创建一个新环境来安装对应的cuda版本：

conda create -n your_env_name python=3.10.13
conda activate your_env_name
conda install cudatoolkit==11.8 -c nvidia
pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118
conda install -c "nvidia/label/cuda-11.8.0" cuda-nvcc
conda install packaging

如果已经是对应的cuda11.8版本了，就可以直接安装mamba了：

git clone https://github.com/Dao-AILab/causal-conv1d.git 
cd causal-conv1d 
git checkout v1.2.0 # current latest version tag 
CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install .
cd ..
git clone https://github.com/state-spaces/mamba.git
cd ./mamba
git checkout v1.2.0 # current latest version tag
MAMBA_FORCE_BUILD=TRUE pip install .

有一些教程写了两中安装（我们按照上面的来即可，这里只是记录下）：

pip install . # 方式一，下载whl安装，两种方式选择一个即可
MAMBA_FORCE_BUILD=TRUE pip install . # 方式二，强制在本地编译安装，Win 下无法识别此命令

注意：安装过程可能比较长，耐心等待即可。另外，以上是基于conda环境安装的。

下面是一些注意事项参考：

pip check numpy #检查numpy在当前环境中的兼容性
conda search cudatoolkit --info #查看当前源中可用的cuda版本
conda install cudatoolkit==11.8 -c nvidia #调一个cuda安装，版本号要写完整
conda install -c "nvidia/label/cuda-11.8.0" cuda-nvcc # 要安装对应的cudnn

最直接的安装，可以利用网友配置好的Docker环境 -＞ 直接使用Mamba基础环境docker镜像。
DockHub仓库地址：https://hub.docker.com/repository/docker/kom4cr0/cuda11.7-pytorch1.13-mamba1.1.1/general
代码：docker pull kom4cr0/cuda11.7-pytorch1.13-mamba1.1.1:1.1.1

参考：conda环境内安装cuda 和 Mamba的安装

二、Win系统安装

安装系统要求：cuda11.8、python3.10

同样的，如果cuda版本不对应，需要在虚拟环境中安装（创建虚拟环境并且安装cudatoolkit==11.8，这样在虚拟变量中会额外有一个cuda版本（比较便捷），后面的cuda-nvcc一定要加，不然当前虚拟环境可能找不到对应的cuda版本）：

conda create -n your_env_name python=3.10.13
conda activate your_env_name
conda install cudatoolkit==11.8 -c nvidia
pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118
conda install -c "nvidia/label/cuda-11.8.0" cuda-nvcc

conda install packaging

关于win系统的安装，首先安装Triton，这个是linux系统，但是有人编译好了：triton-2.0.0-cp310-cp310-win_amd64.whl

注意：安装前没有packaging的需要先安装（还有Triton也是）：

conda install packaging

安装完成后，安装causal_conv1d和mamba_ssm：

1）、安装causal_conv1d

安装该文件网上有几种方法，下面我将一一展示，可以尝试一下（最后一种必定成功）：

1、第一种方法

下载文件：causal-conv1d

将setup.py文件中的代码进行修改：

FORCE_BUILD = os.getenv("CAUSAL_CONV1D_FORCE_BUILD", "FALSE") == "TRUE"
SKIP_CUDA_BUILD = os.getenv("CAUSAL_CONV1D_SKIP_CUDA_BUILD", "FALSE") == "TRUE"
# For CI, we want the option to build with C++11 ABI since the nvcr images use C++11 ABI
FORCE_CXX11_ABI = os.getenv("CAUSAL_CONV1D_FORCE_CXX11_ABI", "FALSE") == "TRUE"

# 修改为：
FORCE_BUILD = os.getenv("CAUSAL_CONV1D_FORCE_BUILD", "FALSE") == "FALSE"
SKIP_CUDA_BUILD = os.getenv("CAUSAL_CONV1D_SKIP_CUDA_BUILD", "FALSE") == "FALSE"
# For CI, we want the option to build with C++11 ABI since the nvcr images use C++11 ABI
FORCE_CXX11_ABI = os.getenv("CAUSAL_CONV1D_FORCE_CXX11_ABI", "FALSE") == "FALSE"

然后cmd执行下面命令：

pip install .

2、第二种方法（感觉可靠）

在git（没有的话在git_download中进行下载）进行安装，因为在windows terminal中识别不到命令行：

CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install .

按照这个步骤安装causal-conv1d应该不是什么大问题

3、第三种方法：直接下载大神编译好的文件进行安装

下载连接：causal_conv1d-1.1.1-cp310-cp310-win_amd64

下载后，直接在cmd中pip安装即可：

pip install causal_conv1d-1.1.1-cp310-cp310-win_amd64.whl

2）、安装mamba_ssm

安装该文件网上有几种方法，下面我将一一展示，可以尝试一下（最后一种必定成功）：

1、第一种方法

下载mamba：Mamba SSM

在mamba源码的 *setup.py *修改配置：

FORCE_BUILD = os.getenv("MAMBA_FORCE_BUILD", "FALSE") == "TRUE"
SKIP_CUDA_BUILD = os.getenv("MAMBA_SKIP_CUDA_BUILD", "FALSE") == "TRUE"

# 修改为：
FORCE_BUILD = os.getenv("MAMBA_FORCE_BUILD", "FALSE") == "FALSE"
SKIP_CUDA_BUILD = os.getenv("MAMBA_SKIP_CUDA_BUILD", "FALSE") == "FALSE"

在*mamba_ssm/ops/selective_scan_interface.py *进行修改（有些教程没有这一步，可以自己尝试是否有必要）：

def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
                     return_last_state=False):
    """if return_last_state is True, returns (out, last_state)
    last_state has shape (batch, dim, dstate). Note that the gradient of the last state is
    not considered in the backward pass.
    """
    return SelectiveScanFn.apply(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
 
 
def mamba_inner_fn(
    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
    out_proj_weight, out_proj_bias,
    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
    C_proj_bias=None, delta_softplus=True
):
    return MambaInnerFn.apply(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
                              out_proj_weight, out_proj_bias,
                              A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)

# 修改为：
def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
                     return_last_state=False):
    """if return_last_state is True, returns (out, last_state)
    last_state has shape (batch, dim, dstate). Note that the gradient of the last state is
    not considered in the backward pass.
    """
    return selective_scan_ref(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
 
def mamba_inner_fn(
    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
    out_proj_weight, out_proj_bias,
    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
    C_proj_bias=None, delta_softplus=True
):
    return mamba_inner_ref(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
                              out_proj_weight, out_proj_bias,
                              A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)

然后直接mamba文件下进行安装即可：

pip install .

2、第二种方法：直接下载大神编译好的文件进行安装

下载连接：mamba_ssm-1.2.0.post1-py3-none-any

下载后，直接在cmd中pip安装即可：

pip install mamba_ssm-1.2.0.post1-py3-none-any.whl

参考：Mamba安装、Mamba 环境