Bamba 9B V2

🚀 Bamba 9B v2

Bamba-9B-v2 is a decoder-only language model based on the Mamba-2 architecture, designed for a wide range of text generation tasks. It outperforms Llama 3.1 8B on leaderboard scores.

🚀 Quick Start

Bamba-9B-v2 is a powerful decoder-only language model. You can quickly start using it for inference, training, fine-tuning, etc. The following sections will provide detailed guides.

✨ Features

• High Performance: Trained with an additional 1T tokens, significantly improving on Bamba v1.
• Wide Range of Tasks: Capable of handling various text generation tasks.
• Efficient Storage and Inference: Support for quantization to make storage and inference more efficient.

📦 Installation

Besides PyTorch, you would need a few extra dependencies for Mamba models.

We found some of these dependencies picky on PyTorch versions when doing pip install, so the best way is to build from source for all Mamba dependencies if you hit dependency issue with your env:

git clone https://github.com/Dao-AILab/causal-conv1d.git
cd causal-conv1d && pip install . && cd ..
git clone https://github.com/state-spaces/mamba.git
cd mamba && pip install . && cd ..
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention && pip install . && cd ..

For users using our HF versions of the model, you would need to install the latest transformers which includes our newly merged implementation for our Bamba models:

pip install git+https://github.com/huggingface/transformers.git

💻 Usage Examples

Basic Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ibm-ai-platform/Bamba-9B-v2")
tokenizer = AutoTokenizer.from_pretrained("ibm-ai-platform/Bamba-9B-v2")

message = ["Mamba is a snake with following properties  "]
inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
response = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])

📚 Documentation

Model Details

We introduce Bamba-9B-v2, a decoder-only language model based on the Mamba-2 architecture and is designed to handle a wide range of text generation tasks. Bamba v2 is 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.

The current release includes the following models:

Original checkpoints (in dcp format) were also uploaded to public bucket:

bucket: bamba-public
endpoint-url: https://s3.us-east.cloud-object-storage.appdomain.cloud

example command to list original Bamba distributed checkpoints:

aws --endpoint-url https://s3.us-east.cloud-object-storage.appdomain.cloud s3 ls s3://bamba-public/checkpoints/pretraining/phase_two/3_1t/step_500/

Training

We trained our Bamba model with FSDP using our training repo here. Note that this training effort was started before FSDP2 and also long before we contributed Mamba2-Hybrid to HF, so we were doing FSDP1 training with official Mamba implementation. For users trying to reproduce the training you now have much more options with our newly contributed HF-version of Mamba2-Hybrid.

Benchmark scores

Base pretrained models

*For the v2 leaderboard results, we perform normalization and report the normalized results. Further details on our evaluation and normalization detailes along with run and analysis scripts can be found here.

Fine-tuning

This example shows how to fine tune the bamba model for a specific task using SFT Trainer.

Quantization

We can create a (FP8) quantized model using fms-model-optimizer, which will make the storage and inference even more efficient.

python -m fms_mo.run_quant \
    --model_name_or_path <"path_to_original_model"> \
    --quant_method fp8 \
    --torch_dtype bfloat16 \
    --output_dir <"path_to_save_new_model">

Model size comparison before and after FP8:

More details about fms-model-optimizer can be found here.

Llama.cpp

There is preliminary work to enable running Bamba architecture models using llama.cpp. This is work-in-progress, so should only be used as a guide for the adventurous!

Known Limitations

• Currently, inference is only supported on CPUs
• Models quantized with llama-quantize exhibit bad performance

Setup

To enable Bamba support, you'll need to build from source using Gabe's fork.

git clone --branch BambaArchitecture git@github.com:gabe-l-hart/llama.cpp.git
cd llama.cpp
mkdir build
cd build
# NOTE: To build with debug symbols and extra logging, use CMAKE_BUILD_TYPE=Debug
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j

Conversion to GGUF

You can use a pre-converted GGUF file from Huggingface (e.g. bamba-9b.gguf). If one doesn't exist, you can use the convert_hf_to_gguf.py script from Gabe's fork to perform the conversion manually.

# Install the python dependencies
cd /path/to/llama.cpp
pip install -r requirements/requirements-convert_hf_to_gguf.txt

# Perform the conversion
./convert_hf_to_gguf.py /path/to/bamba-model --outfile /path/to/bamba-model/bamba-model.gguf

Run with llama-cli

# Run the model with no layers on the GPU (CPU-only)
cd /path/to/llama.cpp
./bin/llama-cli  -ngl 0 -m /path/to/bamba-model/bamba-model.gguf -p "Tell me a story about a developer and their dog"

Quantization with llama-quantize

You can (optionally) quantize the GGUF model using llama.cpp's built in quantizaiton tool llama-quantize.

# Run the quantization (see llama-quantize --help for all quant types)
cd /path/to/llama.cpp
./build/bin/llama-quantize /path/to/bamba-model/bamba-model.gguf Q4_K_M

🔧 Technical Details

The Bamba-9B-v2 model is based on the Mamba-2 architecture. It is trained with an additional 1T tokens, which significantly improves its performance compared to Bamba v1. The model is trained using FSDP and has achieved good results on various benchmarks.

📄 License

This project is licensed under the Apache-2.0 license.

Contributors

• Data collection and curation: We acknowledge and thank AllenAI team for making a high quality open source dataset Dolma as well as Hugging Face data team for making FineWeb-edu and Cosmopedia available. These are tremendous contributions and enable us to create the model today.
• 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, 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 Gomez, 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, and Mukesh Khare.

We would also like to thank the community, in particular Pablo Montalvo-Leroux and Vaibhav Srivastav from Hugging Face who provided valuable feedback to this blog and the PRs into transformers. Further, we would like to thank Tyler Michael Smith from Neural Magic, who is shepherding the integration with vLLM.

A huge shoutout to Meta PyTorch, AllenAI, and Hugging Face teams for their contributions to the open initative, FSDP allowed us to smoothly train this model and the data from Dolma and Fineweb/Cosmopedia made this model today!