许可证：apache-2.0
库名称：transformers
标签：

• bamba

模型详情

Bamba 9B 模型卡

我们推出 Bamba-9B，这是一款基于 Mamba-2 架构的解码器专用语言模型，旨在处理广泛的文本生成任务。该模型采用两阶段训练方法从头开始训练。在第一阶段，模型在 Dolma v1.7 数据集的 2 万亿 token 上进行训练。在第二阶段，模型额外训练了 2000 亿 token，利用精心筛选的高质量数据混合进一步提升性能并增强输出质量。

当前版本包含以下模型：

原始检查点（dcp 格式）也已上传至公共存储桶：

存储桶：bamba-public  
端点 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/  

安装

除了 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  

推理

您可以使用我们新贡献的 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 是一种具有以下特性的蛇："]  
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])  

训练

我们使用 FSDP 训练了 Bamba 模型，训练代码库位于 此处。
请注意，这项训练工作是在 FSDP2 之前开始的，也远早于我们向 HF 贡献 Mamba2-Hybrid 的时间，因此我们使用的是 官方 Mamba 实现 进行 FSDP1 训练。
对于希望复现训练的用户，您现在有更多选择，可以使用我们新贡献的 HF 版本的 Mamba2-Hybrid。

基准分数

基础预训练模型

*对于 v2 排行榜结果，我们进行了 归一化 并报告归一化后的结果。
有关我们的评估和归一化详情以及运行和分析脚本的更多信息，请参见 此处。

微调

此 示例 展示了如何使用 SFT Trainer 为特定任务微调 Bamba 模型。

量化

我们可以使用 fms-model-optimizer 创建一个（FP8）量化模型，这将使存储和推理更加高效。

python -m fms_mo.run_quant \  
    --model_name_or_path <"原始模型路径"> \  
    --quant_method fp8 \  
    --torch_dtype bfloat16 \  
    --output_dir <"保存新模型的路径">  

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 运行

# 在 CPU 上运行模型（无 GPU 层）  
cd /path/to/llama.cpp  
./bin/llama-cli  -ngl 0 -m /path/to/bamba-model/bamba-model.gguf -p "讲一个关于开发者和他们的狗的故事"  

使用 llama-quantize 量化

您可以使用 llama.cpp 内置的量化工具 llama-quantize 对 GGUF 模型进行量化（可选）。