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是一個預訓練的基礎模型，因此沒有任何審核機制。此外，由於該模型未針對對話進行微調，因此不應期望它有良好的對話性能。