File size: 5,672 Bytes
2362e4c
 
0bdb478
 
 
 
 
 
 
 
 
2362e4c
 
 
 
 
c62d006
 
9acfe2c
c62d006
2362e4c
c91f41e
2362e4c
 
 
 
 
c06b789
7834475
c06b789
 
7834475
 
 
 
 
 
2362e4c
 
0b9ddaa
2362e4c
 
 
 
 
 
 
 
 
 
 
 
 
6fba98f
 
2362e4c
 
 
 
 
 
 
 
 
 
7542ad8
2362e4c
1cb4251
7274961
1cb4251
dc811fc
7274961
 
 
7542ad8
 
aedddb4
7542ad8
 
 
dc811fc
0b9ddaa
2362e4c
 
7274961
2362e4c
 
 
 
 
 
 
 
 
 
 
 
 
a59589a
 
 
 
 
0b9ddaa
a59589a
 
 
 
 
 
 
 
 
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
---
license: apache-2.0
datasets:
- HuggingFaceH4/ultrachat_200k
- BAAI/Infinity-Instruct
- HuggingFaceH4/ultrafeedback_binarized
- Intel/orca_dpo_pairs
- argilla/OpenHermesPreferences
base_model:
- Zyphra/Zamba2-1.2B
library_name: transformers
---


# Model Card for Zamba2-1.2B

Zamba2-1.2B-instruct is obtained from Zamba2-1.2B by fine-tuning on instruction-following and chat datasets. Specifically:

1. SFT of the base [Zamba2-1.2B](https://huggingface.co/Zyphra/Zamba2-1.2B) model on [ultrachat_200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) and [Infinity-Instruct](https://huggingface.co/datasets/BAAI/Infinity-Instruct)
2. DPO of the SFT checkpoint on [ultrafeedback_binarized](https://huggingface.co/datasets/HuggingFaceH4/ultrafeedback_binarized), [orca_dpo_pairs](https://huggingface.co/datasets/Intel/orca_dpo_pairs), and [OpenHermesPreferences](https://huggingface.co/datasets/argilla/OpenHermesPreferences)

Zamba2-1.2B-Instruct is a hybrid model composed of state-space ([Mamba2](https://github.com/state-spaces/mamba)) and transformer blocks.

## Quick start

### Prerequisites

To download Zamba2-1.2B-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

```python
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

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

# Format the input as a chat template
prompt = "What factors contributed to the fall of the Roman Empire?"
sample = [{'role': 'user', 'content': prompt}]
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])))
```

## Performance

Zamba2-1.2B-Instruct achieves leading instruction-following and multi-turn chat performance for a model of its size and matches strong models significantly larger. For instance, Zamba2-1.2B-Instruct outperforms Gemma2-2B-Instruct, a very strong model over 2x its size. 

<center>
<img src="https://cdn-uploads.huggingface.co/production/uploads/65bc13717c6ad1994b6619e9/ceOUHVeJPhBgwTDCsR9Y6.png" width="900"/>
</center>

  
| Model | Size | Aggregate MT-Bench | IFEval |
|:-------------:|:----:|:-------------:|:----:|
| **Zamba2-1.2B-Instruct** | 1.2B | **59.53** | **41.45** |
| Gemma2-2B-Instruct | 2.7B | 51.69 | 42.20 |
| H2O-Danube-1.8B-Chat | 1.6B | 49.78 | 27.95 |
| StableLM-1.6B-Chat | 1.6B | 49.87 | 33.77 |
| SmolLM-1.7B-Instruct | 1.7B | 43.37 | 16.53 |
| Qwen2-1.5B-Instruct | 1.5B | N/A | 34.68 |

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

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


Time to First Token (TTFT)             |  Output Generation
:-------------------------:|:-------------------------:
![image/png](https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/5lpWDLdtPPVAk8COJq7gZ.png)  |  ![image/png](https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/V2tS6eCOGbpKybEoZmOB7.png)


And memory overhead

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



## Model Details

Zamba2-1.2B 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 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. 

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


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

A standalone Pytorch implementation of Zamba2-1.2B may be found [here](https://github.com/Zyphra/Zamba2).