Zamba2-7B-Instruct

Model Card for Zamba2-7B-Instruct

Zamba2-7B-Instruct is obtained from Zamba2-7B by fine-tuning on instruction-following and chat datasets.

Zamba2-7B-Instruct is a hybrid model composed of state-space (Mamba2) and transformer blocks.

Zamba2-7B-Instruct long-context has been extended from 4k to 16k context by adjusting the rope frequency in the attention blocks.

Quick start

Prerequisites

To use Zamba2-7B-instruct, 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

# Instantiate model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-7B-instruct")
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-7B-instruct", device_map="cuda", torch_dtype=torch.bfloat16)

# Format the input as a chat template
user_turn_1 = "In one season a flower blooms three times. In one year, there is one blooming season. How many times do two flowers bloom in two years? Please include your logic."
assistant_turn_1 = "In one season, a flower blooms three times. In one year, there is one blooming season. Therefore, in two years, there are two blooming seasons. Since each flower blooms three times in one season, in two blooming seasons, each flower will bloom six times. Since there are two flowers, the total number of times they will bloom in two years is 12."
user_turn_2 = "How many times do the two flowers blossom in three years?"
sample = [{'role': 'user', 'content': user_turn_1}, {'role': 'assistant', 'content': assistant_turn_1}, {'role': 'user', 'content': user_turn_2}]
chat_sample = tokenizer.apply_chat_template(sample, tokenize=False)

# Tokenize input and generate output
input_ids = tokenizer(chat_sample, return_tensors='pt', add_special_tokens=False).to("cuda")
outputs = model.generate(**input_ids, max_new_tokens=150, return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
print((tokenizer.decode(outputs[0])))

To use the context-extended version of Zamba, please load the model with use_long_context=True, i.e.:

model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-7B", device_map="cuda", torch_dtype=torch.bfloat16, use_long_context=True)

Performance

Zamba2-7B-Instruct punches dramatically above its weight, achieving extremely strong instruction-following benchmark scores.

<div style="width: 50%; margin: auto;"> </div>

Moreover, due to its unique hybrid SSM architecture, Zamba2-7B-Instruct achieves extremely low inference latency and rapid generation with a significantly smaller memory footprint than comparable transformer-based models.

And memory overhead

<center> <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/Pgp6o7PFXZVLCY9HQiWYC.png" width="400" alt="Zamba inference and memory cost"> </center>

Zamba2-7B-Instruct's high performance, strong instruction-following and reasoning capabilities for its size makes it an ideal generalist small model for a wide range of applications.

Model Details

Zamba2-7B-Instruct utilizes and extends our original Zamba hybrid SSM-attention architecture. The core Zamba architecture consists of a backbone of Mamba2 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 MLP 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.

<center> <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/XrEIEBxd0fqIgh3LyArAV.png" width="300" alt="Zamba architecture"> </center>

Long Context

Our Zamba2-7B instruct features an experimental long-context mode which extends the context from 4k to 16k context. This was achieved by adjusting the rotation frequency of the rotary position embeddings.

In Needle-In-A-Haystack tests, we observe that Zamba2-7B-Instruct finds the needle with an extremely high success rate up to and slightly beyond 16k context with performance falling off sharply at about 18k context. In future versions we aim to extend this context length significantly.

<center> <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/uP1sbXqF1mZAp1R6NYdCF.png" width="500" alt="Zamba long context performance"> </center>

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.