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。