Model Card for Zamba2-7B

Zamba2-7B is a hybrid model composed of state-space (Mamba) and transformer blocks. It broadly follows the Zamba architecture which consists of a Mamba backbone alternating with shared transformer blocks (see diagram in Model Details). Zamba2-7B possesses four major improvements over Zamba1:

1.) Mamba1 blocks have been replaced with Mamba2 blocks.

2.) We apply a LoRA projector to each shared MLP and attention block, which allows the network to specialize at each invocation of the shared transformer layer across depth. LoRA enables us to add depth-specialization for only a minimal increase in total parameter count.

3.) We utilize two alternating shared attention blocks.

4.) We utilize rotary position embeddings in the shared attention layer.

We found that while hybrid SSM-transformer models are perfectly capable of performing well without position embeddings, adding rotary embeddings to the shared attention block slightly improved performance. Secondly, we utilize two alternating shared attention blocks. We find that this improves performance slightly over a single shared block in terms of performance at fixed parameter budget.

Zamba2-7B uses the Mistral v0.1 tokenizer and was pre-trained on 2T tokens of text and code data sourced from open web-datasets, including Zyda. Subsequently, in a second phase, Zamba2-7B was annealed on a mixture of approximately 100B high-quality tokens.

Note: this is a temporary HuggingFace implementation of Zamba2-7B. It may not yet be fully compatible with all frameworks and tools intended to interface with HuggingFace models.

A standalone Pytorch implementation of Zamba2-7B may be found here.

Quick start

Prerequisites

To use Zamba2-7B, install transformers from source:

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

To install dependencies necessary to run Mamba2 kernels, install mamba-ssm from source (due to compatibility issues with PyTorch) as well as 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 result in significantly higher latency and memory usage.

Inference

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

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

Performance

Zamba2-7B achieves leading and state-of-the-art performance among models ≤8B parameters, outperforming several extremely strong baselines such as Meta's Llama3 series, Google's Gemma series and Mistral-7B. Moreover, due to its unique hybrid SSM architecture, Zamba2-7B achieves extremely low inference latency and rapid generation with a significantly smaller memory footprint than comparable transformer based models.

We believe Zamba2-7B is an ideal generalist model which is cheap and fast to run and fits on the majority of consumer hardware but possesses a powerful intelligence.

And memory overhead

Model Details

Zamba2-7B utilizes and extends our 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. This attention has shared weights to minimize the parameter cost of the model. We find that concatenating the original model embeddings to the input to 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 transformer blocks to gain some 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.

Notice

Zamba2-7B is a pretrained base model and therefore does not have any moderation mechanism and may output toxic or otherwise harmful language. In addition, one should not expect good instruct or chat performance, as this model was not fine-tuned for instruction following or chat.