DOCSEARCHbyAlgolia
Submit Docs

Mamba State Space Model

repository·main·Indexed 12 days ago

https://github.com/state-spaces/mamba

A high-performance state space model (SSM) architecture for linear-time sequence modeling. It provides hardware-aware implementations for language modeling and information-dense sequence tasks as an alternative to Transformers, featuring Mamba, Mamba2, and Mamba3 blocks, as well as the MambaLMHeadModel.

Tokens
2.7K
Snippets
10
Records
12
Agent score
48%

What's inside Mamba

  1. Install Mamba from source

    main

    To use Mamba-3 or install from the latest source tree, use the following commands. To include the CUDA selective_scan_cuda extension, use the MAMBA_KEEP_CUDA_BUILD flag. To force local compilation of the CUDA extension, also include MAMBA_FORCE_BUILD=TRUE.

    ModeCommand
    Source default (no CUDA scan)pip install . --no-build-isolation
    Source from GitHub (no CUDA scan)pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation
    Source CUDA selective scan opt-inMAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation
    Source forced local CUDA buildMAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE pip install . --no-build-isolation
    pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation

  2. Install mamba-ssm

    main

    Mamba requires Linux, Python 3.10+, and PyTorch 1.12+. For CUDA-enabled features (selective_scan_cuda), an NVIDIA GPU and CUDA 11.6+ are required.

    Important: Always use --no-build-isolation when installing to ensure pip uses your existing CUDA-enabled PyTorch instead of installing a CPU-only version in an isolated environment.

    Installation Modes

    ModeCommandDescription
    Core packagepip install mamba-ssm --no-build-isolationInstalls without selective_scan_cuda and without compiling CUDA extensions.
    Core + causal-conv1dpip install "mamba-ssm[causal-conv1d]" --no-build-isolationInstalls core package and causal-conv1d extra, without selective_scan_cuda.
    Force local core buildMAMBA_FORCE_BUILD=TRUE pip install mamba-ssm --no-build-isolationBuilds the pure Python wheel locally.
    CUDA selective scan opt-inMAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolationInstalls selective_scan_cuda (tries prebuilt wheel, then compiles locally).
    Force local CUDA buildMAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolationSkips cached wheels and compiles selective_scan_cuda locally.
    MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolation

  3. Patch ROCm 6.0 for AMD cards

    main

    If you are using an AMD card with ROCm 6.0, you must apply a patch to avoid compilation errors. This step is not required for ROCm 6.1 or later.

    1. Locate your ROCm installation directory (typically /opt/rocm/).
    2. Apply the patch using the following command (use sudo if necessary):
    patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch

  4. Run zero-shot evaluations with lm-evaluation-harness

    main

    To evaluate Mamba models using the lm-evaluation-harness library:

    1. Install the harness: pip install lm-eval==0.4.2.
    2. Run evaluation using the mamba_ssm model type.

    Example for Mamba-1 (130M):

    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

    Example for Mamba-2 (2.7B):

    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
    # Install harness
pip install lm-eval==0.4.2

# Run evaluation
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

  5. Troubleshoot model instability with precision settings

    main

    Mamba models are sensitive to recurrent dynamics and may require higher precision for main model parameters. If you experience instabilities, ensure you are using a framework that keeps model parameters in float32 and only casts to half precision when necessary (e.g., PyTorch AMP).

    Avoid frameworks that store parameters in float16 and upcast only when necessary (such as DeepSpeed), as this can lead to instability in SSMs.

  6. Troubleshoot initialization issues with post-initialization hooks

    main

    Mamba uses specific initializations for certain parameters, such as the $\Delta$ parameter, which relies on a targeted range for the bias of its linear projection.

    If your training framework uses post-initialization hooks (for example, a hook that sets all nn.Linear bias terms to zero), it may overwrite these critical initializations. You may need to implement custom logic to prevent your framework from re-initializing these specific terms.

  7. Benchmark Mamba generation latency and throughput

    main

    Use the benchmarks/benchmark_generation_mamba_simple.py script to test inference performance. This script can autoload models from Hugging Face and benchmark generation speed.

    Test Latency (Sampling Strategies)

    To test latency with specific sampling parameters like top-p (--topp), temperature (--temperature), or repetition-penalty (--repetition-penalty):

    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

    Test Throughput (Large Batch)

    To test throughput with random prompts using a large batch size:

    python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --batch 64

  8. Use the Mamba2 block

    main

    The Mamba2 block implements the Mamba-2 architecture. Usage is similar to the original Mamba block.

    Parameters:

    • d_model: Model dimension.
    • d_state: SSM state expansion factor (typically 64 or 128).
    • d_conv: Local convolution width.
    • expand: Block expansion factor.

    Note: This module uses roughly 3 * expand * d_model^2 parameters.

    from mamba_ssm import Mamba2

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba2(
    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

  9. Use the Mamba block

    main

    The Mamba module is the main architecture block wrapping the selective SSM. It is suitable for standard sequence modeling tasks.

    Parameters:

    • d_model: Model dimension.
    • d_state: SSM state expansion factor.
    • d_conv: Local convolution width.
    • expand: Block expansion factor.

    Note: This module uses roughly 3 * expand * d_model^2 parameters.

    import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    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

  10. Use the Mamba3 block

    main

    The Mamba3 block implements the Mamba-3 architecture, an inference-first state space model.

    Parameters:

    • d_model: Model dimension.
    • d_state: SSM state size.
    • headdim: SSM headdim.
    • is_mimo: Use MIMO mode.
    • mimo_rank: MIMO rank when is_mimo=True.
    • chunk_size: 64/mimo_rank if x is in bf16, else 32/mimo_rank.
    • is_outproj_norm: Additional post SSM norm.
    • dtype: Data type (e.g., torch.bfloat16).

    Note: This module uses roughly 6 * d_model^2 parameters.

    from mamba_ssm import Mamba3
import torch

batch, length, dim = 2, 2048, 768
x = torch.randn(batch, length, dim).to(torch.bfloat16).to("cuda")
model = Mamba3(
    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

  11. Use selective scan operations

    main

    The package exports low-level selective scan operations used for the SSM (State Space Model) computations:

    • selective_scan_fn: The core selective scan function.
    • mamba_inner_fn: The inner function for Mamba computations.

    These are typically used when implementing custom Mamba-like layers or performing low-level kernel operations.

    from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn

  12. Import Mamba model architectures

    main

    The mamba_ssm package provides several core model architectures for different use cases:

    • Mamba: The standard Mamba implementation.
    • Mamba2: The Mamba-2 architecture.
    • Mamba3: The Mamba-3 architecture.
    • MambaLMHeadModel: A Mamba-based Language Model head implementation (found in mamba_ssm.models.mixer_seq_simple).
    from mamba_ssm import Mamba, Mamba2, Mamba3, MambaLMHeadModel