Bamba-9B: Inference-Efficient Hybrid Mamba2 Model

TL;DR

Bamba-9B is an inference‑efficient Hybrid Mamba2 model released by IBM, Princeton, CMU, and UIUC. It shows 2.5× throughput and 2× latency gains versus standard transformers in vLLM and is available for immediate use in transformers, vLLM, TRL, and llama.cpp.

Motivation

Transformer inference is limited by the KV‑cache bottleneck that grows with context length. Hybrid Mamba2 architectures keep the KV‑cache size constant, addressing this bottleneck. Bamba‑9B validates the hybrid Mamba2 approach at the 7B‑10B scale using fully open data and provides reproducible checkpoints to encourage community experimentation.

Use in transformers

To run Bamba‑9B with the 🤗 Transformers library, load the model and tokenizer with AutoModelForCausalLM and AutoTokenizer, then call generate. Example code:

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("ibm-fms/Bamba-9B
tokenizer = AutoTokenizer.from_pretrained("ibm-fms/Bamba-9B
txt = ["Mamba is a snake with following properties  
inputs = tokenizer(txt, return_tensors='pt', return_token_type_ids=False)
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])

Evaluations

Comparison with SoTA transformer models

Bamba‑9B scores 62.31 average on the HF OpenLLM v1 leaderboard, slightly below Meta Llama 3.1 8B (63.51) but above Olmo2 7B (66.17) and IBM Granite v3 8B (67.47) on some metrics. On the v2 leaderboard Bamba‑9B averages 10.91, lower than Llama 3.1 8B (14.27) but comparable when math and MMLU tasks are excluded; the average then nears Llama 3.1 8B’s 44.68 versus Bamba‑9B’s 45.53.

Comparison with transformer models trained on similar token budget

Bamba‑9B (2.2T tokens) averages 62.31, outperforming Olmo1.5 7B trained on 2T tokens (55.8). Even when comparing the 2T‑token Bamba‑9B checkpoint (59.11) to Llama2 7B (53.78) and IBM Granite 7B (52.07), Bamba‑9B remains higher, indicating competitiveness given equal data.

Comparison with other Mamba/Mamba2 models

Bamba‑9B averages 62.31, higher than NVIDIA Mamba2 Hybrid 8B* (58.78) and Zamba 7B (64.36) on some metrics, and lower than Falcon Mamba 7B (65.31). The table shows that hybrid Mamba2 models can deliver competitive results while offering up to 5× theoretical inference efficiency.

Inference efficiency

Bamba‑9B achieves up to 2.5× better throughput and 2× lower latency than Meta Llama 3.1 8B on an NVIDIA H100 80GB GPU using vLLM, measured across batch sizes and sequence lengths from 1K to 64K tokens. The arithmetic‑intensity analysis predicts a potential 5× speedup when the decode stage becomes memory‑bound; current vLLM results are limited by missing chunked pre‑fill support, transformer‑style memory allocation, and unoptimized Mamba2 kernels for H100.

Model architecture

Bamba‑9B uses 32 total layers: 3 full‑attention layers and 29 Mamba2 layers, with an MLP expansion factor of 3.5, a 128k vocab size, RoPE embeddings, and GQA (8 KV‑heads, 32 heads). Compared to the NVIDIA hybrid Mamba2 8B model, Bamba‑9B reduces attention layers from 4 to 3 and adds RoPE.

Data

Training used Dolma v1.7 for the first phase, then FineWeb‑edu and Cosmopedia for an additional 200B tokens. All data were tokenized on an internal Red Hat OpenShift cluster using the Ray framework. The data mix for phase one is visualized in the blog post.

Pre‑training

Pre‑training proceeded in phases: ablation at 1.8B/100B tokens, then 3B/2T tokens with Dolma, followed by a 9B/2T token run, and finally a 200B token fine‑tuning phase with FineWeb‑edu and Cosmopedia. Training hyper‑parameters: cosine LR schedule, peak 3e‑4, quadratic warmup over 2000 steps, decay 0.033, end LR 1e‑5, AdamW (β1=0.9, β2=0.95), weight decay 0.1, sequence length 4096, global batch size 1.5M tokens, 192 A100 GPUs on IBM Cloud Vela for ~2 months. Three job interruptions occurred due to deployment errors and hardware failures, detected by the Autopilot system.

Data loader

The released stateful data loader provides checkpoint‑resumable, auto‑rescaling, zero‑overhead shuffled streaming, asynchronous distributed operation without peer‑to‑peer traffic, dynamic data mixing and on‑the‑fly tokenization, and is PyTorch native, modular, and extensible. It has been battle‑tested on hundreds of training jobs and integrated with Torch Titan.

Quantization

Using the FMS Model Optimizer framework with llm‑compressor, Bamba‑9B checkpoints were quantized to fp8, resulting in negligible accuracy loss: OpenLLM v1 average dropped from 62.31 to 61.5 (−0.1) and v2 average from 10.91 to 10.04 (−0.9). fp8 inference enablement in vLLM is pending kernel updates for Mamba2 layers.

Context length extension

Applying LongRoPE to the full‑attention layers extends Bamba‑9B’s context length. Preliminary PhoneBook retrieval tests show the extended model outperforms the base Bamba‑9B, Llama2‑7B, and Llama3‑8B up to 16K tokens without tuning, and matches Llama3.1‑8B performance. At 32K tokens Llama3.1‑8B leads.

Summary

Bamba‑9B, a hybrid Mamba2 model from IBM, Princeton, CMU, and UIUC trained on 2.2T open tokens, delivers 2.5× throughput and 2× latency gains over Llama 3.1 8B in vLLM, is immediately usable in transformers, vLLM, TRL, and llama.cpp, and comes with training, tuning, and extended‑pretraining recipes plus a stateful data loader.

Future work

Plans include continued pretraining on additional data, SFT with community‑suggested mixes, supervised fine‑tuning using Tulu‑3, Orca‑AgentInstruct, and Daring‑Anteater datasets, enabling chunked pre‑fill and proper memory allocation in vLLM, adding fp8 kernels for faster inference, applying torch.compile and fp8 training, and extending context length to 1M+ tokens.

Contributors

Data collection and curation: AllenAI (Dolma) and Hugging Face (FineWeb‑edu, Cosmopedia). Data preprocessing: IBM team members Tuan Hoang Trong, Syed Zawad, Jay Gala, Ryan Gordon using the IBM Data Prep Kit. Model architecture: Tri Dao (Princeton), Albert Gu (CMU), Linsong Chu (IBM), Davis Wertheimer (IBM), Minjia Zhang (UIUC), Mudhakar Srivatsa (IBM), Raghu Ganti (IBM). Model training: IBM team members Linsong Chu, Divya Kumari, Davis Wertheimer, Raghu Ganti, Dakshi Agrawal. Model tuning: IBM team members Sukriti Sharma, Anh Uong via TRL. Model inference: IBM and community contributors Fabian Lim, Antoni Viros i Martin, Adnan Hoque, Jamie Yang, Nelson Nimura Gonzalez, Joshua Rosenkranz, Nick Hill, Gabe Goodhart. Quantization: IBM team members Naigang Wang, Charlie Liu. Evaluations: IBM evaluation team led by Yotam Perlitz, Ofir Arviv, Michal Shmueli‑Scheuer, Haoechen Shen, Minjia Zhang (UIUC). Leadership acknowledgments: Priya Nagpurkar, David Cox, Sriram Raghavan, Aya Soffer, Ruchir Puri, Mukesh Khare. Community thanks: Pablo Montalvo‑Leroux, Aritra Roy Gosthipaty, Vaibhav Srivastav (Hugging Face), Stas Bekman (Contextual AI), Tyler Michael Smith (Neural Magic). Thanks also to Meta PyTorch, AllenAI, and Hugging Face for open‑source contributions.

Appendix: Arithmetic Intensity

The appendix derives compute and memory equations for attention and Bamba models, showing that Bamba‑9B’s decode‑stage memory advantage can yield up to a 5× speedover Llama for long sequences (>16K tokens). Current vLLM measurements of 2.5× throughput and 2× latency are limited by missing chunked pre‑fill, transformer‑style memory allocation, and unoptimized Mamba2 kernels on H100.