Bamba-9B-v1開源語言模型 - 免費部署助力各類文本生成任務

首頁

Bamba 9B V1

由ibm-ai-platform開發

Bamba-9B 是基於 Mamba-2 架構的解碼器專用語言模型，經過兩階段訓練，擅長處理廣泛的文本生成任務。

大型語言模型開源協議:Apache-2.0 #Mamba2架構 #長文本生成 #兩階段訓練

下載量 16.19k

發布時間 : 12/3/2024

模型概述

Bamba-9B 是一款高效的語言模型，採用兩階段訓練方法，第一階段在 Dolma v1.7 數據集的 2 萬億 token 上訓練，第二階段額外訓練了 2000 億 token 以提升性能。

模型特點

兩階段訓練

第一階段在 2 萬億 token 上訓練，第二階段在 2000 億高質量 token 上進一步優化。

高效架構

基於 Mamba-2 架構，具有 32 層和 4096 隱藏維度，支持 4096 上下文長度。

量化支持

提供 FP8 量化版本，顯著減少內存佔用，提升推理效率。

模型能力

文本生成

語言理解

上下文推理

使用案例

通用文本生成

內容創作

生成文章、故事或其他創意文本內容。

問答系統

回答用戶提出的各種問題。

教育

學習輔助

幫助學生理解複雜概念或生成學習材料。

🚀 Bamba 9B模型

Bamba-9B是一款基於Mamba-2架構的僅解碼器語言模型，可處理多種文本生成任務。它採用兩階段訓練方法從零開始訓練，第一階段在Dolma v1.7數據集的2萬億個標記上訓練，第二階段在精心挑選的2000億個標記上進一步訓練，以優化性能和提升輸出質量。

🚀 快速開始

安裝

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

在使用pip install時，部分依賴對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版本的模型時，需要安裝包含Bamba模型新合併實現的最新transformers庫：

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

推理

可以使用新貢獻的HF集成對Bamba模型進行推理：

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ibm-fms/Bamba-9B-v1")
tokenizer = AutoTokenizer.from_pretrained("ibm-fms/Bamba-9B-v1")

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架構，專為文本生成任務設計。
兩階段訓練：第一階段在Dolma v1.7數據集的2萬億個標記上訓練，第二階段在2000億個標記上進一步優化。
多場景支持：可處理多種文本生成任務。

📦 安裝指南

依賴安裝

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

在使用pip install時，部分依賴對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版本模型安裝

使用HF版本的模型時，需要安裝包含Bamba模型新合併實現的最新transformers庫：

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

💻 使用示例

基礎用法

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ibm-fms/Bamba-9B-v1")
tokenizer = AutoTokenizer.from_pretrained("ibm-fms/Bamba-9B-v1")

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

📚 詳細文檔

模型詳情

Bamba

模型參數

模型	參數數量	層數	隱藏維度	注意力頭數	GQA	KV頭數	上下文長度	綁定嵌入
Bamba	9B (9.78B)	32	4096	32	是	8	4096	否

當前發佈的模型

階段	Bamba 9B	量化版本	說明
基礎模型	ibm-fms/Bamba-9B-v1	ibm-fms/Bamba-9B-fp8	第二階段預訓練
基礎模型	ibm-fms/Bamba-9B-2T	ibm-fms/Bamba-9B-fp8	第一階段預訓練
基礎模型	ibm-fms/Bamba-9B-1.8T	ibm-fms/Bamba-9B-fp8	第一階段中間檢查點，後續還有更多
SFT	即將推出	即將推出	下一次發佈
DPO	即將推出	即將推出	下一次發佈

原始檢查點

原始檢查點（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/2_2t/step_140000_ckp/

訓練

使用FSDP在訓練倉庫中訓練Bamba模型。注意，此訓練工作在FSDP2之前開始，在向HF貢獻Mamba2-Hybrid之前很久就開始了，因此使用官方Mamba實現進行FSDP1訓練。現在，使用新貢獻的HF版本的Mamba2-Hybrid，用戶有更多選擇來複現訓練。

基準測試分數

基礎預訓練模型

類別	基準測試	Bamba 9B (2.2T)
通用	MMLU (5-shot)	60.77
通用	ARC-C (25-shot)	63.23
通用	GSM8K (5-shot)	36.77
通用	Hellaswag (10-shot)	81.8
通用	OpenbookQA (5-shot)	47.6
通用	Piqa (5-shot)	82.26
通用	TruthfulQA (0-shot)	49.21
通用	Winogrande (5-shot)	76.87
HF OpenLLM- V2*	MMLU-PRO (5-shot)	17.53
HF OpenLLM- V2*	BBH (3-shot)	17.4
HF OpenLLM- V2*	GPQA (0-shot)	4.14
HF OpenLLM- V2*	IFEval (0-shot)	15.16
HF OpenLLM- V2*	MATH Lvl 5 (4-shot)	1.66
HF OpenLLM- V2*	MuSR (0-shot)	9.59
安全任務	PopQA (5-shot)	20.5
安全任務	Toxigen (5-shot)	57.4
安全任務	BBQ (5-shot)	44.2
安全任務	Crows-pairs english (5-shot)	70.78

*對於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量化前後模型大小對比：

	原始模型	量化模型
內存（總計）	39.12 GB	10.83 GB
內存（細分）	`torch.float32` 39.12 GB	`torch.bfloat16` 2.10 GB `torch.float8_e4m3fn` 8.73 GB

有關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

🔧 技術細節

模型架構

模型架構設計由普林斯頓大學、卡內基梅隆大學、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。