Mamba环境配置教程【自用】

1. 新建一个Conda虚拟环境

conda create -n mamba python=3.10

2. 进入该环境

conda activate mamba

3. 安装torch（建议2.3.1版本）以及相应的 torchvison、torchaudio
直接进入pytorch离线包下载网址，在里面寻找对应的pytorch以及torchvison、torchaudio
CSDN资源

下载完成后，进入这些文件的目录下，直接使用下面三个指令进行安装即可

pip install torch-2.3.1+cu118-cp310-cp310-linux_x86_64.whl 
pip install torchvision-0.18.1+cu118-cp310-cp310-linux_x86_64.whl 
pip install torchaudio-2.3.1+cu118-cp310-cp310-linux_x86_64.whl

4. 安装triton和transformers库

pip install triton==2.3.1
pip install transformers==4.43.3

5. 安装完这些我们最基本Pytorch环境以及配置完成，接下来就是Mamba所需的一些依赖了，由于Mamba需要底层的C++进行编译，所以还需要手动安装一下cuda-nvcc这个库，直接使用conda命令即可

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

6. 最后就是下载最重要的 causal-conv1d 和mamba-ssm库。在这里我们同样选择离线安装的方式，来避免大量奇葩的编译bug。首先进入下面各自的github网址种进行下载对应版本
causal-conv1d —— 1.4.0

mamba-ssm —— 2.2.2

和安装pytorch一样，进入下载的.whl文件所在文件夹，直接使用以下指令进行安装

pip install causal_conv1d-1.4.0+cu118torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
pip install mamba_ssm-2.2.2+cu118torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

7. 安装好环境后，验证一下Mamba块能否成功运行，直接复制下面代码保存问mamba2_test.py，并运行

# Copyright (c) 2024, Tri Dao, Albert Gu.import math
import torch
import torch.nn as nn
import torch.nn.functional as Ffrom einops import rearrange, repeattry:from causal_conv1d import causal_conv1d_fn
except ImportError:causal_conv1d_fn = Nonetry:from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated, LayerNorm
except ImportError:RMSNormGated, LayerNorm = None, Nonefrom mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combinedclass Mamba2Simple(nn.Module):def __init__(self,d_model,d_state=128,d_conv=4,conv_init=None,expand=2,headdim=64,ngroups=1,A_init_range=(1, 16),dt_min=0.001,dt_max=0.1,dt_init_floor=1e-4,dt_limit=(0.0, float("inf")),learnable_init_states=False,activation="swish",bias=False,conv_bias=True,# Fused kernel and sharding optionschunk_size=256,use_mem_eff_path=True,layer_idx=None,  # Absorb kwarg for general moduledevice=None,dtype=None,):factory_kwargs = {"device": device, "dtype": dtype}super().__init__()self.d_model = d_modelself.d_state = d_stateself.d_conv = d_convself.conv_init = conv_initself.expand = expandself.d_inner = self.expand * self.d_modelself.headdim = headdimself.ngroups = ngroupsassert self.d_inner % self.headdim == 0self.nheads = self.d_inner // self.headdimself.dt_limit = dt_limitself.learnable_init_states = learnable_init_statesself.activation = activationself.chunk_size = chunk_sizeself.use_mem_eff_path = use_mem_eff_pathself.layer_idx = layer_idx# Order: [z, x, B, C, dt]d_in_proj = 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheadsself.in_proj = nn.Linear(self.d_model, d_in_proj, bias=bias, **factory_kwargs)conv_dim = self.d_inner + 2 * self.ngroups * self.d_stateself.conv1d = nn.Conv1d(in_channels=conv_dim,out_channels=conv_dim,bias=conv_bias,kernel_size=d_conv,groups=conv_dim,padding=d_conv - 1,**factory_kwargs,)if self.conv_init is not None:nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)# self.conv1d.weight._no_weight_decay = Trueif self.learnable_init_states:self.init_states = nn.Parameter(torch.zeros(self.nheads, self.headdim, self.d_state, **factory_kwargs))self.init_states._no_weight_decay = Trueself.act = nn.SiLU()# Initialize log dt biasdt = torch.exp(torch.rand(self.nheads, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))+ math.log(dt_min))dt = torch.clamp(dt, min=dt_init_floor)# Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759inv_dt = dt + torch.log(-torch.expm1(-dt))self.dt_bias = nn.Parameter(inv_dt)# Just to be explicit. Without this we already don't put wd on dt_bias because of the check# name.endswith("bias") in param_grouping.pyself.dt_bias._no_weight_decay = True# A parameterassert A_init_range[0] > 0 and A_init_range[1] >= A_init_range[0]A = torch.empty(self.nheads, dtype=torch.float32, device=device).uniform_(*A_init_range)A_log = torch.log(A).to(dtype=dtype)self.A_log = nn.Parameter(A_log)# self.register_buffer("A_log", torch.zeros(self.nheads, dtype=torch.float32, device=device), persistent=True)self.A_log._no_weight_decay = True# D "skip" parameterself.D = nn.Parameter(torch.ones(self.nheads, device=device))self.D._no_weight_decay = True# Extra normalization layer right before output projectionassert RMSNormGated is not Noneself.norm = RMSNormGated(self.d_inner, eps=1e-5, norm_before_gate=False, **factory_kwargs)self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)def forward(self, u, seq_idx=None):"""u: (B, L, D)Returns: same shape as u"""batch, seqlen, dim = u.shapezxbcdt = self.in_proj(u)  # (B, L, d_in_proj)A = -torch.exp(self.A_log)  # (nheads) or (d_inner, d_state)initial_states=repeat(self.init_states, "... -> b ...", b=batch) if self.learnable_init_states else Nonedt_limit_kwargs = {} if self.dt_limit == (0.0, float("inf")) else dict(dt_limit=self.dt_limit)if self.use_mem_eff_path:# Fully fused pathout = mamba_split_conv1d_scan_combined(zxbcdt,rearrange(self.conv1d.weight, "d 1 w -> d w"),self.conv1d.bias,self.dt_bias,A,D=self.D,chunk_size=self.chunk_size,seq_idx=seq_idx,activation=self.activation,rmsnorm_weight=self.norm.weight,rmsnorm_eps=self.norm.eps,outproj_weight=self.out_proj.weight,outproj_bias=self.out_proj.bias,headdim=self.headdim,ngroups=self.ngroups,norm_before_gate=False,initial_states=initial_states,**dt_limit_kwargs,)else:z, xBC, dt = torch.split(zxbcdt, [self.d_inner, self.d_inner + 2 * self.ngroups * self.d_state, self.nheads], dim=-1)dt = F.softplus(dt + self.dt_bias)  # (B, L, nheads)assert self.activation in ["silu", "swish"]# 1D Convolutionif causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:xBC = self.act(self.conv1d(xBC.transpose(1, 2)).transpose(1, 2))  # (B, L, self.d_inner + 2 * ngroups * d_state)xBC = xBC[:, :seqlen, :]else:xBC = causal_conv1d_fn(x=xBC.transpose(1, 2),weight=rearrange(self.conv1d.weight, "d 1 w -> d w"),bias=self.conv1d.bias,activation=self.activation,).transpose(1, 2)# Split into 3 main branches: X, B, C# These correspond to V, K, Q respectively in the SSM/attention dualityx, B, C = torch.split(xBC, [self.d_inner, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)y = mamba_chunk_scan_combined(rearrange(x, "b l (h p) -> b l h p", p=self.headdim),dt,A,rearrange(B, "b l (g n) -> b l g n", g=self.ngroups),rearrange(C, "b l (g n) -> b l g n", g=self.ngroups),chunk_size=self.chunk_size,D=self.D,z=None,seq_idx=seq_idx,initial_states=initial_states,**dt_limit_kwargs,)y = rearrange(y, "b l h p -> b l (h p)")# Multiply "gate" branch and apply extra normalization layery = self.norm(y, z)out = self.out_proj(y)return outif __name__ == '__main__':model = Mamba2Simple(256).cuda()inputs = torch.randn(2, 128, 256).cuda()pred = model(inputs)print(pred.size())      

参考文献