mamba

Mamba is an open-source research codebase from the state-spaces group implementing a family of state space model (SSM) architectures for sequence modeling. It accompanies the papers "Mamba: Linear-Time Sequence Modeling with Selective State Spaces," "Transformers are SSMs" (Mamba-2), and "Mamba-3: Improved Sequence Modeling using State Space Principles." The project positions itself as a subquadratic alternative to Transformers, particularly for information-dense data such as language modeling, and builds on prior structured state space work (S4) with a hardware-aware implementation in the spirit of FlashAttention. What the repository provides: - Selective SSM layer: the core selective scan operation, with source in ops/selective_scan_interface.py. - Mamba block: the main architecture module wrapping the selective SSM (modules/mamba_simple.py), usable as a drop-in PyTorch module with configurable d_model, d_state, d_conv and expand parameters. - Mamba-2 block: implemented in modules/mamba2.py with a simpler variant in modules/mamba2_simple.py, plus a minimal SSD (structured state space duality) reference module in modules/ssd_minimal.py. - Mamba-3 block: implemented in modules/mamba3.py, exposing options such as d_state, headdim, MIMO mode with mimo_rank, chunk_size, an optional output-projection norm, and bfloat16 dtype. - Language model example: a full backbone plus LM head in models/mixer_seq_simple.py, used by the generation scripts. Installation and build options: The package is installed via pip (mamba-ssm) and requires Linux, Python 3.10+ and PyTorch 1.12+. By default the install does not compile the selective_scan_cuda extension and does not fetch cached CUDA wheels. Optional extras and environment flags change this behavior: the causal-conv1d extra adds that dependency; MAMBA_FORCE_BUILD=TRUE forces a local build; MAMBA_KEEP_CUDA_BUILD=TRUE opts into the CUDA selective scan extension, first trying a matching prebuilt CUDA/HIP wheel and compiling locally if none is available. Combining both flags forces local CUDA compilation. The README notes that --no-build-isolation is required for CUDA builds so pip uses the existing CUDA-enabled PyTorch. CUDA builds and GPU execution additionally require an NVIDIA GPU and CUDA 11.6+. For AMD cards, the README documents a ROCm 6.0 patch step (not needed from ROCm 6.1 onward). Pretrained models: Weights are published on Hugging Face under the state-spaces organization, including 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 and mamba2attn-2.7b, trained on 300B tokens on the Pile, plus mamba-2.8b-slimpj trained on 600B tokens on SlimPajama. The README describes these as base models without instruction tuning or other downstream modification, and states that performance is expected to be comparable to or better than other architectures trained on similar data, but not to match larger or fine-tuned models. Model dimensions follow GPT-3-style scaling (e.g. 130M parameters with 24 layers and 768 model dimension, up to 2.8B with 64 layers and 2560 dimension). Evaluation and inference tooling: Zero-shot evaluation is done through the lm-evaluation-harness library (pip install lm-eval==0.4.2), with example commands covering tasks such as lambada_openai, hellaswag, piqa, arc_easy, arc_challenge, winogrande, openbookqa, boolq, race, truthfulqa_mc2 and mmlu. The README notes that per-task results may differ from reported values by 0.1-0.3 due to evaluation noise. A generation benchmark script (benchmarks/benchmark_generation_mamba_simple.py) autoloads a model from the Hugging Face Hub, generates completions for a user prompt, and measures inference speed, with options for top-p, min-p, top-k, temperature, repetition penalty and batch size. Troubleshooting notes: The README highlights two practical issues. First, precision: models were trained with PyTorch AMP, which keeps parameters in float32 and casts to half precision as needed; because SSMs are sensitive to recurrent dynamics, frameworks that store parameters in float16 (such as DeepSpeed) may cause instability, and fp32 parameter storage is suggested as a first remedy. Second, initialization: some components inherit S4-style initializations (for example, the delta parameter's targeted range via linear-projection bias), and frameworks with post-initialization hooks that zero out nn.Linear biases may need custom logic to preserve them. Licensing and citation: the README provides BibTeX entries for the Mamba, Mamba-2 and Mamba-3 papers and asks users to cite them when using the codebase.