Bamba-9B-v2 - Fast and powerful!

This is a joint work between IBM, Princeton, CMU, and UIUC. Full list of authors and contributors is at the bottom of the page.

During Christmas of 2024, IBM, Princeton, CMU, and UIUC released, Bamba v1, a performant Mamba2 based pretrained model with full data lineage trained to 2T tokens. Since then, we have been busy cooking an update with new datasets. Today, we are excited to release Bamba v2, trained for an additional 1T tokens that significantly improves on Bamba v1. The L1 and L2 leaderboard scores outperform Llama 3.1 8B, which was trained with nearly 5x the amount of data. All of this with the inference speedup that we get from Mamba2 based architecture, which with the latest vLLM is 2-2.5x faster than similar sized transformer models.

Artifacts 📦

1. Hugging Face Bamba collection
2. GitHub repo with inference, training, and tuning scripts
3. vLLM RFC

Fast and Powerful ⚡🏎️

Evaluation setup ⚙️ 🖥️: We rerun all the benchmarks following the setup and scripts here for all models. For the v2 leaderboard results, we perform normalization and report the normalized results. In all the evaluations, higher is better except where indicated.

While the HF leaderboards themselves are not active anymore since models evolve, comparing the key benchmarks is an important measure of model capabilities. We provide these comparisons below for various benchmarks. We observe that compared to other SoTA models that are trained to at least 10T+ tokens (and in many cases 15T+), Bamba 9B v2 outperforms the popular Llama 3.1 8B base model on both L1 and L2 averages.

HF OpenLLM v1 benchmarks + OpenbookQA, Boolq, and PIQA

HF OpenLLM v2 benchmarks

Training recipe 📖🍳

Given the limited GPU budget (192 A100s), we did not have the option of training the model for 10T+ tokens. Instead, we decided to explore infusing the existing model with newer data as well as experimenting with techniques like model merging, which was highlighted in the Olmo2 model training. Our training recipe is outlined in the below diagram.

Training receipe for Bamba 9B v2

We took the 2T base checkpoint (aka Bamba 9b v1) and extended it by adding Olmo Mix, released by AllenAI as part of Olmo2 training recipe. We use a constant learning rate schedule to go from 2T to 2.5T tokens, specifically 2e-5. The precise data mix is in the below table.

Table 1: 2-2.5T Olmo2 mix data

We then used a mix of synthetic data from Nemotron-CC and Hugging Face datasets to continue training on 500B additional tokens, putting us at 3T tokens. During this phase, we launch two jobs (serially) to perform an ablation on learning rate schedules, one with constant learning rate at 2e-5 and another with cosine learning rate going from 2e-5 ending at 2e-6. Our general observation is that cosine and constant learning rate schedulers are complementary. The data mixes for 2.5T to 3T are as follows.

Table 2: 2.5-3T Custom mix data

Finally, we anneal both these models using very high quality data for 100B additional tokens and merge the final annealed models using MergeKit. Specifically, we observe that simple weighted averaging works best for us. The resulting model is Bamba 9B v2! The data mix for annealing is below.

Table 3: 3T-3.1T Custom mix data for annealing

Note on instruction tuning 🧠🔧📘 We have experimented with Tuluv3 data for creating an instruction following model leveraging Open Instruct. We observe that the model's performance improves significantly across various benchmarks with a L1 average of 64.69 and an L2 average of 24.68. We are working on an instruction following model that has non-restricted data.

vLLM integration 📥🧠📤

We are deeply engaged with the vLLM community on adding support for Mamba2 attention. In our first release, we worked with various Mamba2 model developers to add support for tensor parallel support and reduce memory copies.

We are currently working on three major updates:

(a) vLLM KV-cache management: The allocation of pages and managing KV-cache is fundamentally different from that of transformer based models. We are working with the community toward achieving a generic solution that can handle any Hybrid model including Mamba2.

(b) Chunked prefill: We are well aware that chunked prefill can give tremendous improvements in real workloads. While the current kernel works well for transformer based models, we need to implement a new kernel for Mamba2 based models.

(c) Faster decode kernels: We are working on Triton implementation of the Conv1D kernel to ensure that the model can run in a performant manner on AMD GPUs as well. Further, we are exploring the fusion of the 4-5 kernels at decode time to reduce the latency of token generation even further.

We anticipate that with the above changes, we can achieve a first class citizen support for Mamba2 based models and significantly better performance than what is today (2-2.5x for longer sequence lengths/larger batches).

Call for Action 📢👉🚀

We are committed to keeping open datasets with complete reproduction of our results. We call on the community to help improve the model on multiple fronts:

1. Test time scaling and GRPO on the model
2. Improve inference performance in vLLM (we expect 4-5x better than corresponding transformer models)
3. Help us improve MATH skills!

Contributors

• Data collection and curation: We acknowledge and thank AllenAI team for making a high quality open source datasets Dolma, OlmoMix, and DolminoMix, Hugging Face data team for making FineWeb* and SmolLM corpus, and NVIDIA team for creating the Nemotron-CC. These are tremendous contributions which enabled us to create and improve the model.
• Data preprocessing: We thank IBM's internal data preprocessing team, specifically Tuan Hoang Trong, Syed Zawad, Jay Gala, and Ryan Gordon for helping tokenize the data at scale. The code for tokenization is available here.
• Model architecture: The model architecture design was jointly done by Princeton, CMU, IBM, and UIUC and involved the following folks: Tri Dao (Princeton), Albert Gu (CMU), Linsong Chu (IBM), Davis Wertheimer (IBM), Minjia Zhang (UIUC), Mudhakar Srivatsa (IBM), and Raghu Ganti (IBM).
• Model training: Model training was performed primarily by the IBM team using the Mamba2 kernels and layer implementation from Tri Dao and Albert Gu. The following folks from IBM were primarily involved: Linsong Chu, Divya Kumari, Davis Wertheimer, Garrett Goon, Raghu Ganti, and Dakshi Agrawal.
• Model tuning: Tuning of the model was enabled and verified in TRL by the IBM team, involving Sukriti Sharma and Anh Uong.
• Model inference: Model inference in transformers, vLLM, and llama.cpp builds on the kernels written by Princeton and CMU. The IBM team is working with the community to enable it in various ecosystems. The team includes Fabian Lim, Antoni viros i Martin, Adnan Hoque, Jamie Yang, Nelson Nimura Gonzalez, Joshua Rosenkranz, Nick Hill, and Gabe Goodhart.
• Quantization: Quantization is led by the IBM team - Naigang Wang and Charlie Liu.
• Evaluations: Evaluations are led by a team in IBM with long context evaluations being performed by UIUC, involving the following folks: Yotam Perlitz, Ofir Arviv, Michal Shmueli-Scheuer (IBM), Haoechen Shen, and Minjia Zhang (UIUC).

Finally, we would like to thank our leadership for their support in this effort - Priya Nagpurkar, David Cox, Sriram Raghavan, Aya Soffer, Ruchir Puri, and Mukesh Khare.

We would also like to thank the community - Hugging Face for transformers enablement and vLLM for inference optimizations.

A huge shoutout to Meta PyTorch, AllenAI, and Hugging Face teams for their contributions to the open initative, PyTorch FSDP allowed us to smoothly train this model and the various high quality datasets that allowed us to update the model further!