Mamba
Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Albert Gu*, Tri Dao*
Paper: https://arxiv.org/abs/2312.00752
Transformers are SSMs: Generalized Models and Efficient Algorithms
Through Structured State Space Duality
Tri Dao*, Albert Gu*
Paper: https://arxiv.org/abs/2405.21060
Mamba-3: Improved Sequence Modeling using State Space Principles
Through Structured State Space Duality
Aakash Lahoti*, Kevin Y. Li*, Berlin Chen*, Caitlin Wang*, Aviv Bick, J. Zico Kolter, Tri Dao†, Albert Gu†
Paper: https://arxiv.org/abs/2603.15569
About
Mamba is a new state space model architecture showing promising performance on information-dense data such as language modeling, where previous subquadratic models fall short of Transformers.
It is based on the line of progress on structured state space models,
with an efficient hardware-aware design and implementation in the spirit of FlashAttention.
Installation
Install PyTorch first. By default, mamba-ssm installs the core package without compiling the
selective_scan_cuda extension and without downloading cached CUDA-enabled wheels.
| Install mode |
Command |
Behavior |
| Core package |
pip install mamba-ssm --no-build-isolation |
Installs mamba-ssm without selective_scan_cuda and does not compile CUDA extensions. |
| Core package plus causal-conv1d |
pip install "mamba-ssm[causal-conv1d]" --no-build-isolation |
Installs the core package and the causal-conv1d extra, still without selective_scan_cuda. |
| Force local core package build |
MAMBA_FORCE_BUILD=TRUE pip install mamba-ssm --no-build-isolation |
Builds the default pure Python wheel locally, still without selective_scan_cuda. |
| CUDA selective scan opt-in |
MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolation |
Installs selective_scan_cuda; pip first tries a matching prebuilt CUDA/HIP wheel, then compiles locally if no wheel is available. |
| Force local CUDA selective scan build |
MAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolation |
Skips cached wheels and compiles selective_scan_cuda locally. |
--no-build-isolation is required for CUDA builds so that pip uses your existing CUDA-enabled
PyTorch instead of installing torch-cpu in an isolated build environment.
Source installs use the same defaults and opt-in flags:
| Install mode |
Command |
Source default, no selective_scan_cuda |
pip install . --no-build-isolation |
Source default from GitHub, no selective_scan_cuda |
pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation |
Source forced local core build, no selective_scan_cuda |
MAMBA_FORCE_BUILD=TRUE pip install . --no-build-isolation |
| Source CUDA selective scan opt-in |
MAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation |
| Source forced local CUDA selective scan build |
MAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation |
NOTE: To use Mamba-3 from the latest source tree, install from source. For the default source
install without selective_scan_cuda, run pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation.
For the CUDA selective scan extension, add MAMBA_KEEP_CUDA_BUILD=TRUE; add
MAMBA_FORCE_BUILD=TRUE as well to force local CUDA compilation.
Core requirements:
- Linux
- Python 3.10+
- PyTorch 1.12+
Additional requirements for CUDA selective_scan_cuda builds and GPU execution:
For AMD cards, see additional prerequisites below.
Usage
We expose several levels of interface with the Mamba model.
Selective SSM
Mamba is based on a selective SSM layer, which is the focus of the paper (Section 3; Algorithm 2).
Source: ops/selective_scan_interface.py.
Mamba Block
The main module of this repository is the Mamba architecture block wrapping the selective SSM.
Source: modules/mamba_simple.py.
Usage:
import torch
from mamba_ssm import Mamba
batch, length, dim = 2, 64, 16
x = 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)
assert y.shape == x.shape
Mamba-2
The Mamba-2 block is implemented at modules/mamba2.py.
A simpler version is at modules/mamba2_simple.py
The usage is similar to Mamba(-1):
from mamba_ssm import Mamba2
model = Mamba2(
# This module uses roughly 3 * expand * d_model^2 parameters
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
).to("cuda")
y = model(x)
assert y.shape == x.shape
SSD
A minimal version of the inner SSD module (Listing 1 from the Mamba-2 paper) with conversion between "discrete" and "continuous" SSM versions
is at modules/ssd_minimal.py.
Mamba-3
The Mamba-3 block is implemented at modules/mamba3.py.
The usage is as follows:
from mamba_ssm import Mamba3
batch, length, dim = 2, 2048, 768
x = torch.randn(batch, length, dim).to(torch.bfloat16).to("cuda")
model = Mamba3(
# This module uses roughly 6 * d_model^2 parameters
d_model=dim, # Model dimension d_model
d_state=128, # SSM state size
headdim=64, # SSM headdim
is_mimo=True, # Use MIMO mode
mimo_rank=4, # MIMO rank when is_mimo=True
chunk_size=16, # 64/mimo_rank if x is in bf16, else 32/mimo_rank
is_outproj_norm=False, # Additional post SSM norm
dtype=torch.bfloat16,
).to("cuda")
y = model(x)
assert y.shape == x.shape
Mamba Language Model
Finally, we provide an example of a complete language model: a deep sequence model backbone (with repeating Mamba blocks) + language model head.
Source: models/mixer_seq_simple.py.
This is an example of how to integrate Mamba into an end-to-end neural network.
This example is used in the generation scripts below.
Pretrained Models
Pretrained models are uploaded to
Hugging Face: mamba-130m, mamba-370m,
mamba-790m, mamba-1.4b, mamba-2.8b, mamba2-130m, mamba2-370m,
mamba2-780m, mamba2-1.3b, mamba2-2.7b, transformerpp-2.7b, mamba2attn-2.7b, trained on 300B tokens on the Pile, as well as mamba-2.8b-slimpj
(trained on 600B tokens on the SlimPajama dataset).
The models will be autodownloaded by the generation script below.
These models were trained on the Pile, and follow the standard model dimensions described by GPT-3 and followed by many open source models:
| Parameters |
Layers |
Model dim. |
| 130M |
24 |
768 |
| 370M |
48 |
1024 |
| 790M |
48 |
1536 |
| 1.4B |
48 |
2048 |
| 2.8B |
64 |
2560 |
(The layer count of Mamba doubles that of a Transformer with similar size, as two Mamba blocks are needed for each "layer" (MHA block + MLP block) of a Transformer.)
Note: these are base models trained only for 300B tokens, without any form of downstream modification (instruction tuning, etc.).
Performance is expected to be comparable or better than other architectures trained on similar data, but not to match larger or fine-tuned models.
Evaluations
To run zero-shot evaluations of models (corresponding to Table 3 of the paper),
we use the
lm-evaluation-harness
library.
- Install
lm-evaluation-harness by pip install lm-eval==0.4.2.
- Run evaluation with (more documentation at the lm-evaluation-harness repo):
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-130m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
python evals/lm_harness_eval.py --model hf --model_args pretrained=EleutherAI/pythia-160m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande --device cuda --batch_size 64
To reproduce the results on the mamba-2.8b-slimpj model reported in the blogposts:
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks boolq,piqa,hellaswag,winogrande,arc_easy,arc_challenge,openbookqa,race,truthfulqa_mc2 --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks mmlu --num_fewshot 5 --device cuda --batch_size 256
To run evaluations on Mamba-2 models, simply replace the model names:
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/transformerpp-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2attn-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
Note that the result of each task might differ from reported values by 0.1-0.3 due to noise in the evaluation process.
Inference
The script benchmarks/benchmark_generation_mamba_simple.py
- autoloads a model from the Hugging Face Hub,
- generates completions of a user-specified prompt,
- benchmarks the inference speed of this generation.
Other configurable options include the top-p (nucleus sampling) probability, and the softmax temperature.
Examples
To test generation latency (e.g. batch size = 1) with different sampling strategies:
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --prompt "My cat wrote all this CUDA code for a new language model and" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name "EleutherAI/pythia-2.8b" --prompt "My cat wrote all this CUDA code for a new language model and" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --prompt "My cat wrote all this CUDA code for a new language model and" --minp 0.05 --topk 0 --temperature 0.7 --repetition-penalty 1.2
To test generation throughput with random prompts (e.g. large batch size):
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --batch 64
python benchmarks/benchmark_generation_mamba_simple.py --model-name "EleutherAI/pythia-2.8b" --batch 64
With Mamba-2, you just need to change the model name:
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba2-2.7b" --prompt "My cat wrote all this CUDA code for a new language model and" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
Troubleshooting
Precision
Our models were trained using PyTorch AMP for mixed precision. AMP keeps model parameters in float32 and casts to half precision when necessary.
On the other hand, other frameworks like DeepSpeed store parameters in float16 and upcasts when necessary (e.g. for optimizer accumulation).
We've observed that higher precision for the main model parameters may be necessary, because SSMs are sensitive to their recurrent dynamics. If you are experiencing instabilities,
as a first step please try a framework storing parameters in fp32 (such as AMP).
Initialization
Some parts of the model have initializations inherited from prior work on S4 models.
For example, the $\Delta$ parameter has a targeted range by initializing the bias of its linear projection.
However, some frameworks may have post-initialization hooks (e.g. setting all bias terms in nn.Linear modules to zero).
If this is the case, you may have to add custom logic (e.g. this line turns off re-initializing in our trainer, but would be a no-op in any other framework)
that is specific to the training framework.
Additional Prerequisites for AMD cards
Patching ROCm
If you are on ROCm 6.0, run the following steps to avoid errors during compilation. This is not required for ROCm 6.1 onwards.
Locate your ROCm installation directory. This is typically found at /opt/rocm/, but may vary depending on your installation.
Apply the Patch. Run with sudo in case you encounter permission issues.
patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch
Citation
If you use this codebase, or otherwise find our work valuable, please cite Mamba:
@article{mamba,
title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
author={Gu, Albert and Dao, Tri},
journal={arXiv preprint arXiv:2312.00752},
year={2023}
}
@inproceedings{mamba2,
title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
author={Dao, Tri and Gu, Albert},
booktitle={International Conference on Machine Learning (ICML)},
year={2024}
}
@misc{lahoti2026mamba3improvedsequencemodeling,
title={Mamba-3: Improved Sequence Modeling using State Space Principles},
author={Aakash Lahoti and Kevin Y. Li and Berlin Chen and Caitlin Wang and Aviv Bick and J. Zico Kolter and Tri Dao and Albert Gu},
year={2026},
eprint={2603.15569},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2603.15569},
}
Mamba 是一种用于序列建模的状态空间模型架构,计算量随序列长度线性增长。在语言建模等信息密集型任务上表现不俗——以往那些低于二次复杂度的模型在这些任务上不如 Transformer;这个仓库是它的 PyTorch 实现。