🚀 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 许可证。