Ubuntu下VMamba环境配置：详细步骤与错误解决方案

因对比实验需要，近日配置VMamba环境，记录如下：

0. 我的环境

• Ubuntu 22.04
• A800 80 GB * 1
• 实际环境安装了 CUDA11.8

1. 前置要求

• 强制要求: CUDA >= 11.6
• 建议要求：torch>=2.0 CUDA>=11.8 (本文采用CUDA==11.8)

2. CUDA 环境

配置CUDA环境，包括下载 CUDA Toolkit 11.8 和相应的cudnn，也可以在conda虚拟环境中配置不高于实际环境版本的CUDA (见后文)

3. 详细环境配置

3.1 基础包配置

1. 创建虚拟环境并激活，非常推荐python=3.10

conda create -n vm python=3.10.13 -y
conda activate vm
conda install packaging

2. （可选）安装cudatoolkit(如果实际环境已经安装了CUDA Toolkit 11.8则可跳过)

conda install cudatoolkit==11.8 -c nvidia

3. 安装相应版本的torch包，可以采用我的版本：

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

4. （可选）安装nvcc

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

5. (可选) 安装gcc/g++

如果你是在虚拟环境中安装的CUDA 11.8, 而实际主机安装的是其他更高版本，则有可能出现g++/gcc不适配的情况，在3.2.3编译selective_scan时会报错：

3.2 VMamba 相关依赖配置

本部分是比较容易出错和比较麻烦的部分，以下步骤供参考，推荐参考这篇文章

1. 项目

git clone https://github.com/MzeroMiko/VMamba.git
cd VMamba 
pip install -r requirements.txt
pip install "numpy<2"

个人建议将timm升级至最新版本

2. 手动安装两个重要的包

前往causal-conv1d发布页下载v1.1.3版本匹配的wheel, 例如我下载了https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.1.3/causal_conv1d-1.1.3+cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

前往mamba-ssm发布页下载v1.1.2版本匹配的wheel, 例如我下载了https://github.com/state-spaces/mamba/releases/download/v1.1.2/mamba_ssm-1.1.2+cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

安装wheel:

pip install causal_conv1d-1.1.3{....}.whl mamba_ssm-1.1.2{....}.whl

（替换{...}为你选择的具体版本）

3. 编译 selective_scan

cd kernels/selective_scan && pip install . --no-build-isolation

这一步可能会有点久, 亲测如果不加--no-build-isolation会报错ModuleNotFoundError: No module named 'torch'

4. 错误记录

我在跑图像分类训练时，遇到了与Issue#264相同的错误：

assert selective_scan_backend in [None, "oflex", "mamba", "torch"]
AssertionError

参考该页的解决方法：
将 VMamba/classification/models/vmamba.py 中的

v2=partial(self.forward_corev2, force_fp32=(not self.disable_force32), selective_scan_backend="core")

替换为

v2=partial(self.forward_corev2, force_fp32=(not self.disable_force32), selective_scan_backend="oflex"),

参考

非常感谢以下文章/帖子/网页，帮助我完成了VMamba环境的正确配置：

1. https://avoid.overfit.cn/post/d5bf0b0481de441f89b43a4be54a361a
2. https://bcdaka.github.io/posts/79428dcdc803cd1cdeb2061f3b54fe5e/#CUDA_118_32
3. https://zhuanlan.zhihu.com/p/687359086
4. https://github.com/MzeroMiko/VMamba/issues/264