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许可证。