当前环境配置 #

pytorch-ignite                           0.5.3
pytorch-lightning                        2.6.1
torch                                    2.10.0+cu128
torchao                                  0.10.0
torchaudio                               2.10.0+cu128
torchcodec                               0.10.0+cu128
torchdata                                0.11.0
torchinfo                                1.8.0
torchmetrics                             1.9.0
torchsummary                             1.5.1
torchtune                                0.6.1
torchvision                              0.25.0+cu128

遇到的编译错误 #

尝试Kaggle预编译Wheel、GitHub源码安装均失败，核心报错如下：

causal-conv1d编译失败 #

error: subprocess-exited-with-error
  
  × Building wheel for causal-conv1d (pyproject.toml) did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  ERROR: Failed building wheel for causal-conv1d
ERROR: ERROR: Failed to build installable wheels for some pyproject.toml based projects (causal-conv1d)

mamba-ssm编译失败 #

Building wheels for collected packages: mamba-ssm
  error: subprocess-exited-with-error
  
  × Building wheel for mamba-ssm (pyproject.toml) did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  Building wheel for mamba-ssm (pyproject.toml) ... error
  ERROR: Failed building wheel for mamba-ssm
Failed to build mamba-ssm
ERROR: ERROR: Failed to build installable wheels for some pyproject.toml based projects (mamba-ssm)

解决步骤 #

1. 对齐CUDA与PyTorch版本兼容性 #

你的PyTorch绑定了CUDA 12.8，但mamba-ssm的底层依赖对高版本CUDA支持滞后。建议降级到CUDA 12.1兼容的PyTorch版本：

pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu121

2. 直接安装预编译轮子 #

优先使用官方预编译的CUDA兼容轮子，跳过源码编译环节：

# 先安装causal-conv1d预编译轮
pip install causal-conv1d==1.2.0 --find-links https://github.com/state-spaces/mamba/releases/download/v1.2.0/
# 再安装mamba-ssm
pip install mamba-ssm==1.2.0 --find-links https://github.com/state-spaces/mamba/releases/download/v1.2.0/

如果上述链接失效，可指定PyTorch CUDA索引源：

pip install causal-conv1d mamba-ssm --extra-index-url https://download.pytorch.org/whl/cu121

3. 补全编译环境依赖（源码编译备用） #

如果必须源码编译，先安装Kaggle缺失的编译工具：

apt-get update && apt-get install -y gcc g++ build-essential cmake

同时设置CUDA环境变量：

export CUDA_HOME=/usr/local/cuda-12.1
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

4. Conda安装替代方案 #

用conda管理依赖，避免pip的编译冲突：

conda install -c conda-forge mamba-ssm

若conda源无对应版本，先安装conda版PyTorch，再用pip安装预编译轮。

5. 安装验证 #

完成后运行以下代码确认安装成功：

import torch
from mamba_ssm import Mamba

model = Mamba(
    d_model=512,
    n_layer=6,
    vocab_size=10000
)
x = torch.randn(2, 100, 512)
output = model(x)
print(output.shape)  # 预期输出 (2, 100, 512)

内容的提问来源于stack exchange，提问作者pandiri veeresh kumar