在网上试了很多解决办法,终于解决。
1.安装
想要运行mamba-2,需要安装causal-conv1d>=1.4.0这是官方github中说的。
GitHub - state-spaces/mamba: Mamba SSM architecture
关于mamba_ssm版本我安装了2.2.2。
下图是我安装的两个版本,
注意:我选择的是两个FALSE版本,我第一次安装的时候causal-conv1d选择了TRUE版本所以出现了上述问题。后面,我卸载了之后重装了FALSE版本就没有这个问题了。
还有就是,如果你是第一次安装可以参考我之前写的博客。
https://blog.csdn.net/lihaiyuan_0324/article/details/138076262https://blog.csdn.net/lihaiyuan_0324/article/details/138076262
如果你是已经安装过1.0版本的,最好是先卸载在装1.0版本的,因为我第一次没卸载1.0,直接安装1.4还是出现了标题的问题,但是后面卸载后重装就没有这个问题了。
causal-conv1d:
mamba_ssm:Releases · state-spaces/mamba · GitHub
2.AttributeError: module 'triton.language' has no attribute 'cumsum'
我解决标题的问题之后,又出现了,上述问题。
按照作者的回复,Please use triton >= 2.1.0
我安装了triton = 2.2.0
就成功了。
3.测试代码
import torchfrom mamba_ssm import Mambabatch, length, dim = 2, 64, 16x = 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)print("Mamba result", y.shape)assert y.shape == x.shapeimport torchfrom mamba_ssm import Mamba2batch, length, dim = 2, 64, 512x = torch.randn(batch, length, dim).to("cuda")model = Mamba2( # This module uses roughly 3 * expand * d_model^2 parameters # make sure d_model * expand / headdim = multiple of 8 d_model=dim, # Model dimension d_model d_state=64, # SSM state expansion factor, typically 64 or 128 d_conv=4, # Local convolution width expand=2, # Block expansion factor headdim=64, # default 64).to("cuda")y = model(x)print("Mamba2 result", y.shape)assert y.shape == x.shape
以下是运行结果
Mamba result torch.Size([2, 64, 16])
Mamba2 result torch.Size([2, 64, 512])
希望可以帮助到你!