Window 下Mamba 环境安装踩坑问题汇总及解决方法 (无需绕过selective_scan_cuda)

导航

背景

Mamba 官方代码链接为:https://github.com/state-spaces/mamba,在原来博客 “Mamba 环境安装踩坑问题汇总及解决方法” 基础上,不绕过selective_scan_cuda,进行 Mamba 环境安装,这样可以获得和 Linux 一样的速度1

安装问题 / 资源自取 / 论文合作想法请+vx931744281

Windows 下环境准备

  1. 前期环境准备,同原来博客 “Mamba 环境安装踩坑问题汇总及解决方法” ,具体为:
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 环境安装踩坑问题汇总及解决方法” 。

  1. causal-conv1d 的安装,同原来博客 “Mamba 环境安装踩坑问题汇总及解决方法” ,具体细化为:
git clone https://github.com/Dao-AILab/causal-conv1d.git
cd causal-conv1d
git checkout v1.1.1  # 安装最新版的话,此步可省略
set CAUSAL_CONV1D_FORCE_BUILD=TRUE
pip install .

官方没有编译好的适用于Windows版本的 whl,因此需要用上述步骤来手动编译。笔者编译好了 Windows 下的 causal_conv1d-1.1.1-cp310-cp310-win_amd64.whl,亦可直接下载安装。

pip install causal_conv1d-1.1.1-cp310-cp310-win_amd64.whl
  1. mamba-ssm 环境准备,下载工程文件,即
git clone https://github.com/state-spaces/mamba.git
cd mamba
git checkout v1.1.1   # 安装最新版的话,此步可省略

注意,上述过程中,新版 mamba-ssm 需要搭配新版 causal-conv1d,要不然函数不兼容。完成前期工作后进入下一步正式编译。

Windows 下 mamba-ssm 的编译

  • 在mamba源码 setup.py 修改第41行配置:
FORCE_BUILD = os.getenv("MAMBA_FORCE_BUILD", "TRUE") == "TRUE"
  • csrc/selective_scan/selective_scan_fwd_kernel.cuhvoid selective_scan_fwd_launch 函数改为
void selective_scan_fwd_launch(SSMParamsBase &params, 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.hBOOL_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/cus/selective_scan_bwd_kernel.cuhcsrc/selective_scan/cus/selective_scan_fwd_kernel.cuh 文件开头加入:
#ifndef M_LOG2E
#define M_LOG2E 1.4426950408889634074
#endif
  • 完成上述修改后,执行 pip install . 一般即可顺利编译,成功安装。
  • 本人编译好的Windows 下的whl 也有:mamba-ssm-1.1.3 ,可直接下载安装或联系本人vx自取。利用 whl 安装命令为:
pip install mamba-ssm-1.1.3-cp310-cp310-win-amd64.whl

由于此时没有绕过selective_scan_cuda,在虚拟环境中产生了 selective-scan-cuda.cp310-win-amd64.pyd 文件,所以运行速度较快。

Windows 下 mamba-ssm 的编译出现的问题及解决(20240714)

1. 基本报错信息

如果不进行修改利用 `pip install .`` 直接编译源码时会出现如下报错:

subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.
      
      The above exception was the direct cause of the following exception:
      
      Traceback (most recent call last):

        File "/data/xxx/anaconda3/envs/xxx/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 2116, in _run_ninja_build
          raise RuntimeError(message) from e
      RuntimeError: Error compiling objects for extension
      [end of output]

这是最基本的报错信息,只要编译出错就会输出这些,如果在其上面没有看到具体报错,可在 setup.py 里面,将

cmdclass={"bdist_wheel": CachedWheelsCommand, "build_ext": BuildExtension}

改为

cmdclass={"bdist_wheel": CachedWheelsCommand, "build_ext": BuildExtension.with_options(use_ninja=False)}

pytorch默认使用ninjia作为backend2,禁用掉可以看到具体的报错,但是编译速度实测会变慢,所以解决bug后可以改回来。

注:有的博客将 anaconda环境下的 lib/python3.6/site-packages/torch/utils/cpp_extension.py文件里的[‘ninja’,‘-v’]改成[‘ninja’,‘–v’] 或者[‘ninja’,‘–version’] 是错误的做法,治标不治本。

2. “M_LOG2E” is undefined

在Windows下会出现如下大量报错:

 xxx\mamba-1.1.3\csrc\selective_scan\selective_scan_bwd_kernel.cuh(221): error: identifier "M_LOG2E" is undefined

出现这种情况的原因,可参考 issue

Note for the owners: The reason for needing #define is stated here: https://stackoverflow.com/a/56974843:
“On windows it is using the Microsoft compiler for that. So the Microsoft compiler is correct to disallow VLA, and there is no way to avoid this AFAIK. Your code works on linux, because on linux nvcc uses the g++ host compiler, and it allows (in a non-standard-compliant way) the use of a VLA in C++ host code.”

因此,只需在csrc/selective_scan/cus/selective_scan_bwd_kernel.cuhcsrc/selective_scan/cus/selective_scan_fwd_kernel.cuh 文件加入以下代码即可

#ifndef M_LOG2E
#define M_LOG2E 1.4426950408889634074
#endif

3. error C2975: “kIsVariableC_”

在Windows下会出现如下大量报错:

error C2975: “kIsVariableC_”:“Selective_Scan_bwd_kernel_traits”的模板参数无效,应为编译时常量表达式

csrc/selective_scan/static_switch.h 函数里的 constexpr 改为 static constexpr,参考 issue。具体步骤参看前一节。

4. error C2975: “kNRows_”

在Windows下会出现如下大量报错:

xxx\mamba-1.1.3\csrc\selective_scan\selective_scan_fwd_kernel.cuh(314): error C2975: “kNRows_”:“Selective_Scan_fwd_kernel_traits”的模板参数无效,应为编译时常量表达式

csrc/selective_scan/selective_scan_fwd_kernel.cuh 函数 void selective_scan_fwd_launch 里的 constexpr 改为 static constexpr,参考 issue。具体步骤参看前一节。


  1. Windows Support #12 ↩︎

  2. 出现错误“subprocess.CalledProcessError: Command ‘[‘ninja‘, ‘-v‘]‘ returned non-zero exit status 1”解决方法 ↩︎

相关推荐

  1. Windows 11的安装门槛

    2024-07-15 19:48:04       29 阅读
  2. Microsoft登录:安装Windows 11 23H2的两种方法

    2024-07-15 19:48:04       50 阅读

最近更新

  1. docker php8.1+nginx base 镜像 dockerfile 配置

    2024-07-15 19:48:04       52 阅读
  2. Could not load dynamic library ‘cudart64_100.dll‘

    2024-07-15 19:48:04       54 阅读
  3. 在Django里面运行非项目文件

    2024-07-15 19:48:04       45 阅读
  4. Python语言-面向对象

    2024-07-15 19:48:04       55 阅读

热门阅读

  1. golang基础用法

    2024-07-15 19:48:04       14 阅读
  2. shell脚本传参调用http接口

    2024-07-15 19:48:04       15 阅读
  3. JsonCPP源码分析——分配器和配置器

    2024-07-15 19:48:04       14 阅读
  4. Web开发:<div>标签作用

    2024-07-15 19:48:04       17 阅读
  5. 查找子串方法总结

    2024-07-15 19:48:04       22 阅读
  6. 两个事务update同一条数据会发生什么?

    2024-07-15 19:48:04       18 阅读