Awesome

BlackMamba

BlackMamba: Mixture of Experts for State-space models
Quentin Anthony*, Yury Tokpanov*, Paolo Glorioso*, Beren Millidge*
Paper: https://arxiv.org/abs/2402.01771

About

In this repository we provide inference code for our BlackMamba model.

BlackMamba is an novel architecture which combines state-space models (SSMs) with mixture of experts (MoE). It uses Mamba as its SSM block and switch transformer as its MoE block base. BlackMamba is extremely low latency for generation and inference, providing significant speedups over all of classical transformers, MoEs, and Mamba SSM models. Additionally, due to its SSM sequence mixer, BlackMamba retains linear computational complexity in the sequence length.

Requirements

pip install causal-conv1d>=1.1.0: required for Mamba. The rest of the kernels should be built locally.

Other requirements:

Linux NVIDIA GPU PyTorch 1.12+ CUDA 11.6+

Quick installation in a fresh Python environment

• pip install torch packaging
• pip install . to install from source from this repository

Pretrained Models

Our pretrained models are uploaded to our HuggingFace:

• BlackMamba 340M/1.5B*
• BlackMamba 630M/2.8B*

*Since models are MoE, they're named according to (Forward Pass Parameters) / (Total Parameters) for clarity.

Usage

from mamba_model import MambaModel
import torch

model = MambaModel.from_pretrained(pretrained_model_name="Zyphra/BlackMamba-2.8B")
model = model.cuda().half()
inputs = torch.tensor([1, 2]).cuda().long().unsqueeze(0)
out = model(inputs)