Text Generation
Transformers
Safetensors
zamba2
File size: 5,306 Bytes
de74f0f
 
 
be85c15
de74f0f
74d6591
de74f0f
b274e6c
4ac715e
356b2e3
4ac715e
356b2e3
b274e6c
871ddca
de74f0f
4074c1c
 
9f8442b
de74f0f
 
 
be85c15
de74f0f
687e5a7
de74f0f
687e5a7
 
14d866d
 
de74f0f
14d866d
 
 
 
 
 
de74f0f
 
 
 
 
 
 
 
51d1471
019014a
de74f0f
3f5315e
de74f0f
 
 
 
 
 
4074c1c
de74f0f
356b2e3
de74f0f
 
4074c1c
de74f0f
 
 
9570f91
de74f0f
356b2e3
de74f0f
356b2e3
de74f0f
 
9570f91
de74f0f
 
3e804a1
07a6082
3e804a1
de74f0f
356b2e3
 
 
 
 
de74f0f
356b2e3
de74f0f
 
 
 
9570f91
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
---
license: apache-2.0
---
# Model Card for Zamba2-2.7B

Zamba2-2.7B is a hybrid model composed of state-space and transformer blocks. It broadly follows the [Zamba architecture](https://arxiv.org/abs/2405.16712) which consists of a Mamba backbone alternating with shared transformer blocks (see diagram in [Model Details](#model-details)). Zamba2-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 throughout the network.

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

Zamba2-2.7B uses the Mistral v0.1 tokenizer and 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 standalone Pytorch implementation of Zamba2-2.7B may be found [here](https://github.com/Zyphra/Zamba2).

## Quick start

### Prerequisites

To use Zamba2-2.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

```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 = "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]))
```

## 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 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. 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>


## 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 inference latency and rapid generation with a significantly smaller memory footprint than comparable transformer based models. 

Zamba2-2.7B's high performance and small inference compute and memory footprint renders it an ideal generalist model for on-device applications.

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

<center>
<img src="https://cdn-uploads.huggingface.co/production/uploads/65bc13717c6ad1994b6619e9/3u8k7tcRi-oC_ltGhdHAk.png" width="800" alt="Zamba performance">
</center>

Time to First Token (TTFT)             |  Output Generation
:-------------------------:|:-------------------------:
![](https://cdn-uploads.huggingface.co/production/uploads/65bc13717c6ad1994b6619e9/BmE8X6tDNVw5OJcbZt8sZ.png)  |  ![](https://cdn-uploads.huggingface.co/production/uploads/65bc13717c6ad1994b6619e9/wECc9cItK1FW1MOMGSLrp.png)


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

## 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.