Zamba2 2.7B

🚀 Model Card for Zamba2-2.7B

Zamba2-2.7B is a hybrid model that combines state-space and transformer blocks, offering high performance with low latency and memory usage.

🚀 Quick Start

Prerequisites

To use Zamba2-2.7B, you need to install transformers from source:

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

To install the dependencies required to run Mamba2 kernels, install mamba-ssm from source (due to compatibility issues with PyTorch) and 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

You can run the model without using the optimized Mamba2 kernels, but it is not recommended as it will lead to significantly higher latency and memory usage.

Inference

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]))

✨ Features

Zamba2-2.7B has three major improvements over Zamba1:

1. Mamba1 blocks have been replaced with Mamba2 blocks.
2. Instead of a single shared attention block, two shared attention blocks are used, interleaved in an ABAB pattern throughout the network.
3. A LoRA projector is applied to each shared MLP block, allowing the network to specialize the MLPs at each invocation of the shared layer across depth with only a minimal increase in total parameter count.

📦 Installation

The installation steps are included in the Quick Start section.

💻 Usage Examples

Basic Usage

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]))

📚 Documentation

Model Details

Zamba2-2.7B utilizes and extends the original Zamba hybrid SSM-attention architecture. The core Zamba architecture consists of a backbone of Mamba layers interleaved with one or more shared attention layers (one shared attention in Zamba1, two in Zamba2). This attention has shared weights to minimize the parameter cost of the model. Concatenating the original model embeddings to the input of this attention block improves performance, likely due to better maintenance of information across depth. The Zamba2 architecture also applies LoRA projection matrices to the shared MLP to gain additional expressivity in each block and allow each shared block to specialize slightly to its own unique position while keeping the additional parameter overhead small.

Performance

Zamba2-2.7B achieves leading and state-of-the-art performance among models of <3B parameters and is competitive with some models of significantly greater size. Due to its unique hybrid SSM architecture, Zamba2-2.7B has extremely low inference latency and rapid generation with a significantly smaller memory footprint than comparable transformer-based models. Its high performance and small inference compute and memory footprint make it an ideal generalist model for on-device applications.

Time to First Token (TTFT)	Output Generation

📄 License

This project is licensed under the Apache-2.0 license.

⚠️ Important Note

Zamba2-2.7B is a pretrained base model and does not have any moderation mechanism, so it may output toxic or otherwise harmful language. Also, do not expect good instruct or chat performance as this model was not fine-tuned for instruction following or chat.