Zamba2 2.7B

🚀 Zamba2-2.7B模型卡片

Zamba2-2.7B是一個由狀態空間和Transformer模塊組成的混合模型。它大致遵循Zamba架構，該架構由一個Mamba主幹與共享的Transformer模塊交替組成（詳見模型細節中的圖示）。與Zamba1相比，Zamba2-2.7B有三大改進：

1. 用Mamba2模塊取代了Mamba1模塊。
2. 不再使用單個共享注意力模塊，而是採用兩個共享注意力模塊，它們以ABAB模式貫穿整個網絡。
3. 對每個共享MLP模塊應用LoRA投影器，這使得網絡能夠在跨深度調用共享層時對MLP進行專門化。LoRA使我們只需少量增加總參數數量，就能實現深度專門化。

Zamba2-2.7B使用Mistral v0.1分詞器，在從開放網絡數據集（包括Zyda）獲取的3T文本和代碼數據上進行了預訓練。隨後，在第二階段，Zamba2-2.7B在100B高質量令牌的混合數據上進行了退火處理。

注意：這是Zamba2-2.7B在HuggingFace上的臨時實現。它可能尚未完全兼容所有旨在與HuggingFace模型交互的框架和工具。

Zamba2-2.7B的獨立PyTorch實現可在此處找到。

🚀 快速開始

前提條件

要使用Zamba2-2.7B，需從源代碼安裝transformers：

1. git clone https://github.com/huggingface/transformers.git
2. cd transformers && pip install .

要安裝運行Mamba2內核所需的依賴項，需從源代碼安裝mamba-ssm（由於與PyTorch存在兼容性問題）以及causal-conv1d：

1. git clone https://github.com/state-spaces/mamba.git
2. cd mamba && git checkout v2.1.0 && pip install .
3. pip install causal-conv1d

你可以在不使用優化的Mamba2內核的情況下運行模型，但不建議這樣做，因為這將導致顯著更高的延遲和內存使用。

推理

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-2.7B")
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-2.7B", device_map="cuda", 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]))

✨ 主要特性

• 架構創新：採用混合的狀態空間和Transformer架構，結合Mamba模塊與共享注意力模塊，提升性能。
• 參數優化：通過共享注意力權重和應用LoRA投影器，在控制參數數量的同時實現深度專門化。
• 高性能：在參數少於3B的模型中取得領先的性能，與一些更大規模的模型具有競爭力。
• 低延遲和小內存佔用：獨特的混合SSM架構使其具有極低的推理延遲、快速的生成速度和較小的內存佔用。

📦 安裝指南

安裝transformers

git clone https://github.com/huggingface/transformers.git
cd transformers && pip install .

安裝Mamba2內核依賴

git clone https://github.com/state-spaces/mamba.git
cd mamba && git checkout v2.1.0 && pip install .
pip install causal-conv1d

💻 使用示例

基礎用法

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-2.7B")
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-2.7B", device_map="cuda", 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]))

📚 詳細文檔

模型細節

Zamba2-2.7B利用並擴展了我們最初的Zamba混合SSM-注意力架構。核心Zamba架構由Mamba層的主幹與一個或多個共享注意力層交替組成（Zamba1中有一個共享注意力層，Zamba2中有兩個）。這種注意力具有共享權重，以最小化模型的參數成本。我們發現，將原始模型嵌入連接到該注意力模塊的輸入可以提高性能，這可能是因為跨深度更好地維護了信息。Zamba2架構還將LoRA投影矩陣應用於共享MLP，以在每個模塊中獲得一些額外的表達能力，並允許每個共享模塊根據其獨特位置進行微調，同時保持額外的參數開銷較小。

性能表現

Zamba2-2.7B在參數少於3B的模型中取得了領先的、最先進的性能，並且與一些規模大得多的模型具有競爭力。此外，由於其獨特的混合SSM架構，Zamba2-2.7B實現了極低的推理延遲和快速生成，與可比的基於Transformer的模型相比，內存佔用顯著更小。

Zamba2-2.7B的高性能以及較小的推理計算和內存佔用使其成為設備端應用的理想通用模型。

首次令牌時間 (TTFT)	輸出生成

🔧 技術細節

Zamba2-2.7B的技術實現細節包括：

• 架構設計：採用混合的狀態空間和Transformer架構，結合Mamba模塊與共享注意力模塊，以提高性能和效率。
• 參數優化：通過共享注意力權重和應用LoRA投影器，在控制參數數量的同時實現深度專門化。
• 數據處理：使用Mistral v0.1分詞器，在3T文本和代碼數據上進行預訓練，並在100B高質量令牌的混合數據上進行退火處理。

📄 許可證

本項目採用Apache-2.0許可證。