Bamba 9B V2

🚀 Bamba 9B v2 模型

Bamba-9B-v2 是一個僅解碼器的語言模型，基於 Mamba-2 架構構建，旨在處理廣泛的文本生成任務。該模型在 Bamba v1 的基礎上，額外訓練了 1T 個標記，顯著提升了性能。其在 L1 和 L2 排行榜上的得分超過了 Llama 3.1 8B，而 Llama 3.1 8B 的訓練數據量幾乎是 Bamba-9B-v2 的 5 倍。

🚀 快速開始

你可以利用我們新貢獻的 HF 集成在 Bamba 模型上運行推理：

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])

✨ 主要特性

• 基於 Mamba-2 架構，能夠處理廣泛的文本生成任務。
• 在 Bamba v1 的基礎上額外訓練 1T 個標記，性能顯著提升。
• 在 L1 和 L2 排行榜上的得分超過 Llama 3.1 8B。

📦 安裝指南

除了 PyTorch，你還需要為 Mamba 模型安裝一些 額外依賴。

我們發現，在使用 pip 安裝時，其中一些依賴對 PyTorch 版本比較挑剔。因此，如果你在環境中遇到依賴問題，最好的方法是從源代碼構建所有 Mamba 依賴：

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 ..

對於使用我們 HF 版本模型的用戶，你需要安裝最新的 transformers，其中包含了我們為 Bamba 模型新合併的實現：

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

💻 使用示例

基礎用法

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])

📚 詳細文檔

模型詳情

模型參數

當前發佈的模型

原始檢查點

原始檢查點（dcp 格式）也已上傳到公共存儲桶：

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

列出原始 Bamba 分佈式檢查點的示例命令：

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/

訓練

我們使用 FSDP 在 此訓練倉庫 中訓練了我們的 Bamba 模型。

請注意，這項訓練工作是在 FSDP2 之前開始的，並且在我們將 Mamba2-Hybrid 貢獻給 HF 之前很久就開始了。因此，我們使用 官方 Mamba 實現 進行了 FSDP1 訓練。

對於試圖復現訓練的用戶，現在你可以使用我們新貢獻的 HF 版本的 Mamba2-Hybrid 有更多選擇。

基準測試分數

基礎預訓練模型

*對於 v2 排行榜結果，我們進行了 歸一化 並報告了歸一化結果。有關我們的評估和歸一化細節以及運行和分析腳本的更多詳細信息，請參閱 此處。

微調

此示例 展示瞭如何使用 SFT Trainer 針對特定任務微調 Bamba 模型。

量化

我們可以使用 fms-model-optimizer 創建一個 (FP8) 量化模型，這將使存儲和推理更加高效。

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">

FP8 量化前後的模型大小比較：

有關 fms-model-optimizer 的更多詳細信息，請參閱 此處。

Llama.cpp

目前有初步工作支持使用 llama.cpp 運行 Bamba 架構模型。這仍在進行中，因此僅供勇於嘗試的用戶參考！

已知限制

• 目前，僅支持在 CPU 上進行推理。
• 使用 llama-quantize 量化的模型性能較差。

安裝設置

要啟用 Bamba 支持，你需要使用 Gabe 的分支 從源代碼構建。

git clone --branch BambaArchitecture git@github.com:gabe-l-hart/llama.cpp.git
cd llama.cpp
mkdir build
cd build
# 注意：要使用調試符號和額外日誌進行構建，請使用 CMAKE_BUILD_TYPE=Debug
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j

轉換為 GGUF

你可以使用 Huggingface 上預先轉換的 GGUF 文件（例如 bamba-9b.gguf）。如果沒有可用的文件，你可以使用 Gabe 分支中的 convert_hf_to_gguf.py 腳本手動進行轉換。

# 安裝 Python 依賴
cd /path/to/llama.cpp
pip install -r requirements/requirements-convert_hf_to_gguf.txt

# 進行轉換
./convert_hf_to_gguf.py /path/to/bamba-model --outfile /path/to/bamba-model/bamba-model.gguf

使用 llama-cli 運行

# 在沒有 GPU 層的情況下運行模型（僅使用 CPU）
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"

使用 llama-quantize 進行量化

你可以（可選）使用 llama.cpp 內置的量化工具 llama-quantize 對 GGUF 模型進行量化。

# 運行量化（請參閱 llama-quantize --help 瞭解所有量化類型）
cd /path/to/llama.cpp
./build/bin/llama-quantize /path/to/bamba-model/bamba-model.gguf Q4_K_M

貢獻者

• 數據收集和整理：我們感謝 AllenAI 團隊提供了高質量的開源數據集 Dolma，以及 Hugging Face 數據團隊提供了 FineWeb-edu 和 Cosmopedia。這些巨大的貢獻使我們能夠創建今天的模型。
• 數據預處理：我們感謝 IBM 內部的數據預處理團隊，特別是 Tuan Hoang Trong、Syed Zawad、Jay Gala 和 Ryan Gordon，他們幫助大規模地對數據進行了分詞。分詞代碼可在 此處 找到。
• 模型架構：模型架構設計由普林斯頓大學、卡內基梅隆大學、IBM 和伊利諾伊大學厄巴納 - 香檳分校共同完成，參與人員包括 Tri Dao（普林斯頓大學）、Albert Gu（卡內基梅隆大學）、Linsong Chu（IBM）、Davis Wertheimer（IBM）、Minjia Zhang（伊利諾伊大學厄巴納 - 香檳分校）、Mudhakar Srivatsa（IBM）和 Raghu Ganti（IBM）。
• 模型訓練：模型訓練主要由 IBM 團隊使用 Tri Dao 和 Albert Gu 的 Mamba2 內核和層實現完成。IBM 的以下人員主要參與其中：Linsong Chu、Divya Kumari、Davis Wertheimer、Raghu Ganti 和 Dakshi Agrawal。
• 模型微調：IBM 團隊在 TRL 中啟用並驗證了模型的微調，參與人員包括 Sukriti Sharma 和 Anh Uong。
• 模型推理：transformers、vLLM 和 llama.cpp 中的模型推理基於普林斯頓大學和卡內基梅隆大學編寫的內核。IBM 團隊正在與社區合作，在各種生態系統中啟用它，團隊成員包括 Fabian Lim、Antoni viros i Martin、Adnan Hoque、Jamie Yang、Nelson Nimura Gomez、Joshua Rosenkranz、Nick Hill 和 Gabe Goodhart。
• 量化：量化由 IBM 團隊領導 - Naigang Wang 和 Charlie Liu。
• 評估：評估由 IBM 的一個團隊領導，長上下文評估由伊利諾伊大學厄巴納 - 香檳分校進行，參與人員包括 Yotam Perlitz、Ofir Arviv、Michal Shmueli-Scheuer（IBM）、Haoechen Shen 和 Minjia Zhang（伊利諾伊大學厄巴納 - 香檳分校）。

最後，我們要感謝我們的領導對這項工作的支持 - Priya Nagpurkar、David Cox、Sriram Raghavan、Aya Soffer 和 Mukesh Khare。

我們還要感謝社區，特別是 Pablo Montalvo-Leroux 和 Vaibhav Srivastav 來自 Hugging Face，他們為這個博客和對 transformers 的 PR 提供了寶貴的反饋。此外，我們要感謝來自 Neural Magic 的 Tyler Michael Smith，他正在推動與 vLLM 的集成。

🔧 技術細節

我們使用 FSDP 在 此訓練倉庫 中訓練了我們的 Bamba 模型。訓練工作在 FSDP2 之前開始，且在將 Mamba2-Hybrid 貢獻給 HF 之前很久就已開展，因此採用 官方 Mamba 實現 進行 FSDP1 訓練。如今，嘗試復現訓練的用戶可藉助新貢獻的 HF 版本的 Mamba2-Hybrid 獲得更多選擇。

📄 許可證

本項目採用 Apache-2.0 許可證。