--- license: apache-2.0 --- # Model Card for Zamba v2 2.7B Zamba2-2.7B is a hybrid model between state-space models and transformers. It broadly follows the [Zamba architecture](https://arxiv.org/abs/2405.16712) which consists of a Mamba backbone alternating with shared transformer blocks. Zamba-2-2.7B possesses three major improvements over Zamba1: 1.) Mamba1 blocks have been replaced with Mamba2 blocks. 2.) Instead of a single shared attention block, we utilize two shared attention blocks which are interleaved in an ABAB pattern through the network. 3.) We apply a LoRA projector to each shared MLP block allowing the network to specialize the MLPs at each shared layer with a minimal increase in total parameter count. Zamba was trained using next-token prediction. It uses the Mistral v0.1 tokenizer. Zamba2-2.7B was pre-trained on 3T tokens of text and code data sourced from open web-datasets, including [Zyda](https://arxiv.org/abs/2406.01981). Subsequently, in a second phase, Zamba2-2.7B was annealed on a mixture of 100B high-quality tokens. Note: this is a temporary HuggingFace implementation of Zamba2-2.7B. It may not yet be fully compatible with all frameworks and tools intended to interface with HuggingFace models. A pure Pytorch implementation of Zamba2-2.7B may be found [here](https://github.com/Zyphra/Zamba2). ## Quick start ### Presequities To download Zamba2-2.7B, clone Zyphra's fork of transformers: 1. `git clone https://github.com/Zyphra/transformers_zamba2.git` 2. `cd transformers_zamba2` 3. Install the repository: `pip install -e .` You can run the model without using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly higher latency and memory usage. To run on CPU, please specify `use_mamba_kernels=False` when loading the model using ``AutoModelForCausalLM.from_pretrained``. ### Inference ```python 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 = "A funny prompt would be " input_ids = tokenizer(input_text, return_tensors="pt").to("cuda") outputs = model.generate(**input_ids, max_new_tokens=100) print(tokenizer.decode(outputs[0])) ``` ## Model Details Zamba2-2.7B utilizes and extends our original Zamba hybrid SSM-attention architecture. The core Zamba architecture consists of a backbone of Mamba layers interleaved with a shared attention layer. 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.
Zamba architecture
## 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. Moreover, due to its unique hybrid SSM architecture, Zamba2-2.7B achieves extremely low latency and rapid generation with a significantly smaller memory footprint than comparable transformer based models. Zamba2-2.7B's high performance and small compute and memory footprint renders it an ideal generalist model for on-device applications.
Zamba performance
(-/ TODO All eval figure)
Zamba inference and memory cost
## Notice Zamba2-2.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.