Mamba-2: IndexError: map::at

[Prophet-Kathleen]

Mamba-2: IndexError: map::at

[maksymdolgikh]

I found a workaround, at least for T4. Set the model to half-precision to avoid this error:

model = MambaLMHeadModel.from_pretrained('state-spaces/mamba2-130m')
model = model.half()

Let me know if you understand why this works :)

[tridao]

Looks like a Triton error, which GPU do you use?

[Prophet-Kathleen]

rtx2080ti with driver version 535.171.04 right now, i am trying to use different version of python or newest trition 2.3.1 or pytorch
to fix this problem

[tridao]

I'm not sure triton supports GPUs before Ampere (e.g. 2080) very well

[Prophet-Kathleen]

I just borrowed a rtx3060 (driver version 535.171.04) to test the code and it works. and triton version is 2.3.1. thanks

[ghaddarAbs]

same issue here but using V100 :(

[jsie7]

With newest triton version (2.3.1), this seems mainly related to the used GPU. I also ran into this error on an RTX 2080 Ti, thus I tried to reproduce the error on different GPUs I have available.

Setup: Nvidia driver 535.161.07, Cuda 11.8, Triton 2.3.1, mamba-ssm v2.0.3
Working GPUs: V100, RTX 3090, RTX 4090, A100 (40GB & 80GB)
Index Error (map::at): RTX 2080 Ti, Titan RTX, Quadro RTX 6000
--> it seems the error only occurs for the Turing microarchitecture

@ghaddarAbs V100 works for me; maybe update triton to 2.3.1?

[ghaddarAbs]

@jsie7 thanks for suggesting I will try it out ... which torch version you used ?

[jsie7]

@jsie7 thanks for suggesting I will try it out ... which torch version you used ?

I'm using v2.0.1

[SolomidHero]

I experienced same problem :(

Driver Version: 550.54.14      
CUDA Version: 12.4 # by nvcc -V
Tesla T4  
triton==2.3.1

I only installed torch with cuda12.1 support.

[jsie7]

I experienced same problem :(

Driver Version: 550.54.14      
CUDA Version: 12.4 # by nvcc -V
Tesla T4  
triton==2.3.1

I only installed torch with cuda12.1 support.

The Tesla T4 is also based on the Turing microarchitecture. This just further confirms that it's an issue with that architecture.

[JulienSiems]

Thanks @maksymdolgikh this worked for me, also no idea why this works though.

[samjoc]

I found a workaround, at least for T4. Set the model to half-precision to avoid this error:

model = MambaLMHeadModel.from_pretrained('state-spaces/mamba2-130m')
model = model.half()

Let me know if you understand why this works :)

worked on 2080Ti for me too.