Vim 官方代码链接为:https://github.com/hustvl/Vim,在原来博客 “Mamba 环境安装踩坑问题汇总及解决方法” 基础上,不绕过selective_scan_cuda进行 Vim 环境安装,这样可以获得和 Linux 一样的速度。注意,Vim (Vision Mamba)和 Vmamba (VMamba: Visual State Space Model)虽然都是基于mamba,但是它们不是同一篇!
(安装问题 / 资源自取 / 论文合作想法请+vx:931744281
)
conda create -n mamba python=3.10
conda activate mamba
conda install cudatoolkit==11.8
pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118
pip install setuptools==68.2.2
conda install nvidia/label/cuda-11.8.0::cuda-nvcc_win-64
conda install packaging
pip install triton-2.0.0-cp310-cp310-win_amd64.whl
其中 triton-2.0.0-cp310-cp310-win_amd64.whl
获取参看原来博客 “Mamba 环境安装踩坑问题汇总及解决方法” 。
然后下载 Vim 的官方代码:
git clone https://github.com/hustvl/Vim.git
causal-conv1d
的安装,Vim 官方代码仓里给了其源码,因此只需要利用其源码安装即可,类似 “Window 下Mamba 环境安装踩坑问题汇总及解决方法 (无需绕过selective_scan_cuda)”cd causal-conv1d
set CAUSAL_CONV1D_FORCE_BUILD=TRUE
pip install .
官方没有编译好的适用于Windows版本的 whl,因此需要用上述步骤来手动编译。
笔者编译好了 Windows 下的 causal_conv1d-1.0.0-cp310-cp310-win_amd64.whl,亦可直接下载安装(只适用于torch 2.1)。
pip install causal_conv1d-1.0.0-cp310-cp310-win_amd64.whl
完成前期工作后进入下一步正式编译。注意安装成功后会在相应环境(xxx\conda\envs\xxx\Lib\site-packages\
)中生成 causal_conv1d_cuda.cp310-win_amd64.pyd
文件,此文件对应 causal_conv1d_cuda
包。
mamba-ssm
的编译Vim 官方对 mamba-ssm
的源码进行了修改,所以其与原版有不同。
既可以在安装完原版的基础上再修改相应环境(xxx\conda\envs\xxx\Lib\site-packages\
)中的源码文件,参考“Window 下Mamba 环境安装踩坑问题汇总及解决方法 (无需绕过selective_scan_cuda)”。
也可以直接强行利用Vim的源码进行编译。方法如下:
cd mamba-1p1p1 # 先切换到Vim 下面的这个目录
setup.py
修改第41行配置:FORCE_BUILD = os.getenv("MAMBA_FORCE_BUILD", "TRUE") == "TRUE"
修改第261行配置为(可选,相当于在 BuildExtension
后多加了一个.with_options(use_ninja=False)
):
cmdclass={"bdist_wheel": CachedWheelsCommand, "build_ext": BuildExtension.with_options(use_ninja=False)}
csrc/selective_scan/selective_scan_fwd_kernel.cuh
的 void selective_scan_fwd_launch
函数改为void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
// Only kNRows == 1 is tested for now, which ofc doesn't differ from previously when we had each block
// processing 1 row.
static constexpr int kNRows = 1;
BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
BOOL_SWITCH(params.is_variable_B, kIsVariableB, [&] {
BOOL_SWITCH(params.is_variable_C, kIsVariableC, [&] {
BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] {
using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ, input_t, weight_t>;
// constexpr int kSmemSize = Ktraits::kSmemSize;
static constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t);
// printf("smem_size = %d\n", kSmemSize);
dim3 grid(params.batch, params.dim / kNRows);
auto kernel = &selective_scan_fwd_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
});
});
});
}
csrc/selective_scan/static_switch.h
的 BOOL_SWITCH
函数改为#define BOOL_SWITCH(COND, CONST_NAME, ...) \
[&] { \
if (COND) { \
static constexpr bool CONST_NAME = true; \
return __VA_ARGS__(); \
} else { \
static constexpr bool CONST_NAME = false; \
return __VA_ARGS__(); \
} \
}()
(这两步是将 constexpr
改为 static constexpr
)
csrc/selective_scan/selective_scan_bwd_kernel.cuh
和 csrc/selective_scan/selective_scan_fwd_kernel.cuh
文件开头加入:#ifndef M_LOG2E
#define M_LOG2E 1.4426950408889634074
#endif
pip install .
一般即可顺利编译,成功安装。pip install mamba_ssm-1.1.1-cp310-cp310-win_amd64.whl --no-dependencies causal_conv1d
由于此时没有绕过selective_scan_cuda,在虚拟环境中(xxx\conda\envs\xxx\Lib\site-packages\
)产生了 selective-scan-cuda.cp310-win-amd64.pyd 文件,所以运行速度较快。