Parameter count

[state-spaces]

Parameter count

[tridao]

That would be strange, I'm getting 129M params:

In [1]: import torch

In [2]: from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

In [3]: model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-130m")

In [4]: sum(p.numel() for p in model.parameters())
Out[4]: 129135360

[jmercat]

I was going to respond that:

from mamba_ssm import MambaLMHeadModel

params = {
    "d_model": 2560,
    "n_layer": 32,
    "vocab_size": 50432,
}

model = MambaLMHeadModel(**params)

num_param = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Number of parameters (local): {num_param}")

model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b")

num_param = sum(p.numel() for p in model.parameters())
print(f"Number of parameters (downloaded): {num_param}")

gives:

Number of parameters (local): 1449006080
Number of parameters (downloaded): 2768345600

But I see you've updated the number of layers to the different configs. So problem solved.