Zamba 7B V1 Phase1

🚀 Zamba模型卡片

Zamba-7B-v1-phase1是一个结合了状态空间模型Mamba和Transformer的混合模型。它采用Mamba作为主干架构，并每6个模块插入一个共享的Transformer层。Zamba通过下一令牌预测进行训练，使用Mistral v0.1分词器。我们经过一系列小规模的消融实验后确定了这一架构。Zamba-7B-v1-phase-1在从开放网络数据集获取的1T文本和代码数据令牌上进行了预训练。与Zamba-v1不同，此模型仅代表在网络数据集上进行纯预训练后的检查点。我们主要将其用作比较工具，以探索退火过程的效果。

⚠️ 重要提示

当前Huggingface上的Zamba实现比我们内部的实现速度慢。我们正在与Huggingface团队合作解决这个问题。

我们描述Zamba训练过程的技术报告可点击此处查看。

🚀 快速开始

前置依赖

要下载Zamba，请克隆Zyphra的transformers分支：

1. git clone https://github.com/Zyphra/transformers_zamba
2. cd transformers_zamba
3. 安装该仓库：pip install -e .

为了在CUDA设备上运行优化后的Mamba实现，你需要安装mamba-ssm和causal-conv1d：

pip install mamba-ssm causal-conv1d>=1.2.0

你可以在不使用优化后的Mamba内核的情况下运行模型，但不建议这样做，因为这会显著增加延迟。

若要在CPU上运行，请在使用AutoModelForCausalLM.from_pretrained加载模型时指定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 = "What factors contributed to the fall of the Roman Empire?"
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")

默认迭代是完全训练后的第一阶段模型，对应迭代462070次。这是从随机初始化开始训练模型所进行的迭代次数。有关训练的更多详细信息，请参阅arXiv:2405.16712。

📚 详细文档

模型细节

Zamba采用了独特的混合SSM架构。该架构由Mamba层作为主干，并穿插着一个共享的注意力层。这个注意力层使用共享权重，以最小化模型的参数成本。我们发现，将原始模型嵌入连接到这个注意力块的输入可以提高性能，这可能是因为在深度上更好地保留了信息。

性能表现

我们发现，在这个规模下，Zamba的性能明显优于现有的开放模型（具有开放数据集和训练细节）。然而，它在7B规模上的表现略逊于领先的开放权重模型。这种差异主要源于MMLU和推理评估。不过，Zamba的训练令牌数量明显少于这些模型，并且在每个训练令牌的性能方面是最节省样本的模型。

由于其SSM架构，Zamba在推理方面极其高效，由于其显著减少的KV缓存，在推理延迟和生成的内存成本方面大幅优于可比的7B和8B模型。

📄 许可证

本项目采用Apache-2.0许可证。

📖 引用

如果您在工作中发现Zamba很有用，请按以下方式引用它：

@article{glorioso2024zamba,
  title={Zamba: A Compact 7B SSM Hybrid Model},
  author={Glorioso, Paolo and Anthony, Quentin and Tokpanov, Yury and Whittington, James and Pilault, Jonathan and Ibrahim, Adam and Millidge, Beren},
  journal={arXiv preprint arXiv:2405.16712},
  year={2024}
}

⚠️ 注意事项

Zamba是一个预训练的基础模型，因此没有任何审核机制。此外，由于该模型未针对对话进行微调，因此不应期望它有良好的对话性能。