记ubuntu20.04 不同conda环境多版本cuda安装与切换
需求描述
先看一下我的配置：[前提–>我已经安装了显卡驱动和11.3的cuda toolkit]
nvidia-smi
NVIDIA-SMI: 535.183.01 Driver Version: 535.183.01 CUDA Version: 12.2
早年base环境中装的cuda toolkit是11.3的版本：（对应的是1.12的torch）
nvcc -V
nvcc: NVIDIA ® Cuda compiler driver
Copyright © 2005-2021 NVIDIA Corporation
Built on Sun_Mar_21_19:15:46_PDT_2021
Cuda compilation tools, release 11.3, V11.3.58
Build cuda_11.3.r11.3/compiler.29745058_0
现在想要运行mamba，11.3下安装会报错：
RuntimeError: causal_conv1d is only supported on CUDA 11.6 and above. Note: make sure nvcc has a supported version by running nvcc -V.
而且torch也最好是2.x版本，所以我需要创建一个conda环境，再安装版本在12.2以下11.6以上的cuda，我这里选择的是cuda11.8+torch2.6.0。最终实现base环境下链接向cuda11.3的torch1.12版本正常运行，conda环境（umamba）下链接向cuda11.8的torch2.6.0正常运行。
具体步骤
- 新建conda环境，并激活，后续操作都在新环境中进行
conda create -n umamba python=3.11 # 'umamba' is my_env_name
conda activate umamba
- 安装torch [–>官网]，选择你的环境对应的配置，我这里采用pip安装：
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
为了加快速度且尽量不中断，这里建议使用清华源，在上面的命令后添加：
-i https://pypi.tuna.tsinghua.edu.cn/simple
我这里安装完之后是可以正常使用torch的（import torch; torch.__version__; torch.version.cuda; torch.cuda.is_available()这几个命令都能正常输出，也能正常跑非mamba的网络训练），GPT给出的解释是尽管nvcc -V链接的是11.3的cuda，但是官方发布的PyTorch自带了运行所需的 CUDA 11.8 runtime libraries。
- 下载11.8的cuda[–>官网]：
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
- 安装cuda11.8
sudo sh cuda_11.8.0_520.61.05_linux.run
这里安装时需要注意以下几点：
第一步continue后同意一堆协议说明，Do you accept the above EULA? (accept/decline/quit): accept
#输入accept接受协议；
第二步(注意!!!)用方向下键↓，移动到 Driver 上，回车取消勾选 Driver，由于安装的多版本 CUDA ，安装第一个 CUDA 的时候已经安装过驱动，所以不要重复安装！然后下键移动到 Install 上，回车开始安装；
第三步有一个类似于Do you want to install a symbolic link at /usr/local/cuda?的问题，建议选择no。因为之前安装过另一个版本的cuda，除非你确定想要用这个新版本的cuda，否则这里就建议选no，不然指定该链接后会将cuda指向这个新的版本；（如果还有Install the CUDA 10.0 Samples?也选择no，不用安装样例）
其他就回车，然后等待，安装成功后会print出类似于下面的信息：
- 利用activate和deactivate的命令实现不同conda环境下的cuda指向切换。首先确保激活umamba环境，在这个环境里写一个专属的 env-specific 环境变量配置方法，gpt推荐写到环境的 etc/conda/activate.d 和etc/conda/deactivate.d中。
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
nano $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
添加内容如下：
#!/bin/bash
export CUDA_HOME=/usr/local/cuda-11.8
export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
再创建停用脚本：
mkdir -p $CONDA_PREFIX/etc/conda/deactivate.d
nano $CONDA_PREFIX/etc/conda/deactivate.d/env_vars.sh
添加内容如下：
#!/bin/bash
export PATH=${PATH///usr/local/cuda-11.8/bin:/}
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH///usr/local/cuda-11.8/lib64:/}
unset CUDA_HOME
- 最后测试一下
conda deactivate
conda activate umamba
nvcc -V
- 安装mamba核心组件 causal-conv1d 和 mamba-ssm。进入github[–>causal-conv1d官方]还有[–>mamba-ssm官方]，选中并下载安装对应版本的 [causal-conv1d 1.5.0]和[mamba-ssm]：
pip install causal_conv1d-1.5.0.post8+cu11torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl
pip install mamba_ssm-2.2.4+cu11torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl
最后用代码测试一下：
import torch
from mamba_ssm import Mamba
batch, length, dim = 2, 4, 4
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)
输出如下： 
完。
参考链接
- Chatgpt
- Ubuntu CUDA相关问题全解
- ubuntu22.04多版本安装cuda及快速切换
- Mamba在Ubuntu系统上稳定安装运行
- Ubuntu20.04下成功安装Mamba的运行环境
