license: apache-2.0

Zamba模型卡

Zamba-7B-v1-phase1是状态空间模型Mamba与Transformer的混合架构。该模型以Mamba为主干网络，每6个模块共享一个Transformer层，通过下一词预测进行训练，并采用Mistral v0.1的分词器。此架构是我们经过小规模消融实验后确定的成果。该模型在开放网络数据集提供的1万亿文本和代码标记上进行了预训练。与Zamba-v1不同，此版本仅基于网络数据集的纯预训练检查点，主要用途是作为对比工具来研究退火处理效果。

注：当前Huggingface实现的Zamba运行速度低于我们内部版本，正与Huggingface团队协同优化。

技术报告详见arXiv论文。

快速开始

环境准备

下载步骤：

1. 克隆Zyphra定制版transformers：git clone https://github.com/Zyphra/transformers_zamba
2. 进入目录：cd transformers_zamba
3. 安装仓库：pip install -e .

CUDA设备需安装优化组件：

pip install mamba-ssm causal-conv1d>=1.2.0

非优化模式虽可运行但会显著增加延迟（不推荐）。CPU运行时需在加载模型时指定use_mamba_kernels=False。

推理示例

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba-7B-v1-phase1")
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba-7B-v1-phase1", device_map="auto", torch_dtype=torch.bfloat16)

input_text = "罗马帝国衰亡的因素有哪些？"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

加载特定检查点（如第2500次迭代）：

model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba-7B-v1-phase1", device_map="auto", torch_dtype=torch.bfloat16, revision="iter2500")

默认检查点为完整训练的phase1模型（对应462070次迭代），训练细节参见arXiv论文。

模型架构

Zamba采用独特的混合SSM架构：Mamba主干层间插入共享权重的注意力层，通过拼接原始嵌入提升信息跨层保持能力。

性能表现

在同等规模开源模型中，Zamba展现出显著优势，但略逊于7B量级的顶尖开源模型（差异主要来自MMLU和推理评估）。值得注意的是，Zamba的训练标记量更少，具有最优的样本效率。

得益于SSM架构，Zamba在推理效率和生成内存开销上大幅领先同类7B/8B模型。

引用

引用格式：

@article{glorioso2024zamba,
  title={Zamba: 紧凑型70亿参数SSM混合模型},
  author={Glorioso, Paolo and Anthony, Quentin and Tokpanov, Yury and Whittington, James and Pilault, Jonathan and Ibrahim, Adam and Millidge, Beren},
  journal={arXiv预印本 arXiv:2405.16712},
  year={2024}
}

注意事项

Zamba是基础预训练模型，未包含内容审核机制。由于未针对对话微调，不适用于聊天场景。