File size: 6,860 Bytes
189b8e7 accd6ba 189b8e7 5bde2e1 189b8e7 8037865 189b8e7 443ad2e 189b8e7 4678449 189b8e7 4678449 189b8e7 4678449 189b8e7 f945251 189b8e7 9b2853e 17f05b8 9b2853e 189b8e7 96c3f4a 189b8e7 17b544f 189b8e7 |
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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 |
---
license: apache-2.0
datasets:
- tiiuae/falcon-refinedweb
pipeline_tag: text-generation
library_name: openlm
tags:
- mamba
- linear
language:
- en
model-index:
- name: mamba-7b
results:
- task:
type: text-generation
dataset:
type: MMLU
name: MMLU
metrics:
- name: accuracy
type: accuracy
value: 33.3
verified: false
- task:
type: text-generation
dataset:
type: HellaSwag
name: HellaSwag
metrics:
- name: accuracy
type: accuracy
value: 77.9
verified: false
- task:
type: text-generation
dataset:
type: PIQA
name: PIQA
metrics:
- name: accuracy
type: accuracy
value: 81.0
verified: false
- task:
type: text-generation
dataset:
type: Winogrande
name: Winogrande
metrics:
- name: accuracy
type: accuracy
value: 71.8
verified: false
- task:
type: text-generation
dataset:
type: ai2_arc
name: ARC-E
metrics:
- name: accuracy
type: accuracy
value: 77.5
verified: false
- task:
type: text-generation
dataset:
type: ai2_arc
name: ARC-C
metrics:
- name: accuracy
type: accuracy
value: 46.7
verified: false
---
# Mamba-7B
This is a 7B parameter model with the [Mamba](https://arxiv.org/abs/2312.00752) architecture, trained on multiple epochs (1.2T tokens) of the [RefinedWeb](https://huggingface.co/datasets/tiiuae/falcon-refinedweb) dataset.
Mamba is a state-space model that does not use self-attention unlike the standard transformer architecture. It has shown strong performance on various natural language benchmarks. To date, the largest publicly released pure-Mamba pretrain is [Mamba-2.8B](https://huggingface.co/state-spaces/mamba-2.8b).
We follow their training recipe and release our version of Mamba-7B.
This model was trained as a baseline for our paper [Linearizing Large Language Models](https://arxiv.org/abs/2405.06640).
## Model Details
- **Developed by**: [Toyota Research Institute](https://www.tri.global/our-work/robotics)
- **Model Type**: This is an auto-regressive language model based on the [Mamba](https://arxiv.org/abs/2312.00752) architecture.
- **Dataset**: Trained on 1.2T tokens of [RefinedWeb](https://huggingface.co/datasets/tiiuae/falcon-refinedweb)
- **Tokenizer**: `EleutherAI/gpt-neox-20b`
- **Library**: [OpenLM](https://github.com/mlfoundations/open_lm/)
- **License**: This model is licensed under [Apache License, Version 2.0](https://www.apache.org/licenses/LICENSE-2.0).
| Parameters | Hidden Size | Layers | Vocab Size | Sequence Length |
|------------|-------------|--------| ---------- | --------------- |
| 7B | 4096 | 64 | 50432 | 2048 |
## Training Details
- Mamba-7B was trained using AWS SageMaker on 128 H100 80GB GPUs.
- Training began in March 2024 and lasted three weeks.
| **Hyperparameter** | **Value** |
|--------------------|------------|
| Precision | `bfloat16` |
| Optimizer | AdamW |
| Learning rate | 3e-4 |
| LR cooldown end | 1e-5 |
| Warmup steps | 2000 |
| Z-loss | 1e-4 |
| Batch size | 2M |
## Usage
This model was trained using [OpenLM](https://github.com/mlfoundations/open_lm/). The weights have been converted to be compatible with HuggingFace.
```python
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("tri-ml/mamba-7b-rw")
model = AutoModelForCausalLM.from_pretrained("tri-ml/mamba-7b-rw")
inputs = tokenizer(["The Toyota Supra"], return_tensors="pt")
gen_kwargs = {"max_new_tokens": 50, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
output = model.generate(inputs['input_ids'], **gen_kwargs)
output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
print(output)
# The Toyota Supra is a sports car that has been in production since 1978. The car was discontinued in 2002, but it has recently been revived and will be available again in 2020. The Supra has always been known for its powerful engines and agile handling.
```
## Performance Evaluation
Our evaluations were done using the [Eleuther LM Eval Harness](https://github.com/EleutherAI/lm-evaluation-harness) repo.
Below we report the performance of Mamba 7B compared to other base models.
<style>
.evalTable th { background: white; }
.evalTable tr:nth-child(1) { background: #f3f3f3; }
.evalTable tr:nth-child(2) { background: #f3f3f3; }
.evalTable tr:nth-child(7) { background: #f3f3f3; }
</style>
<div class="evalTable">
| | HellaSwag | PIQA | Winogrande | ARC-E | ARC-C | MMLU (5-shot) |
| ----------------- | ------------- | -------- | -------------- | --------- | --------- | ---------------- |
| Mamba-1.4B | 59.0 | 73.9 | 61.4 | 65.5 | 32.9 | 25.2 |
| Mamba-2.8B | 71.0 | 78.1 | 65.9 | 68.2 | 41.7 | 26.2 |
| RWKV5-1.7T-7B | 73.0 | 78.6 | 72.9 | 75.8 | 45.6 | 34.9 |
| Llama2-7B | 76.0 | 79.1 | 69.1 | 76.3 | 46.3 | 45.9 |
| Gemma-7B | 80.7 | 81.9 | 73.7 | 81.1 | 53.2 | 62.9 |
| Mistral-7B | 81.0 | 82.1 | 74.0 | 80.9 | 53.8 | 62.4 |
| **Mamba-7B** | 77.9 | 81.0 | 71.8 | 77.5 | 46.7 | 33.3 |
</div>
## How to Cite
If you use this model, please cite our paper on [Linearizing Large Language Models](https://arxiv.org/abs/2405.06640).
```
@article{Mercat2024Linearizing,
title={Linearizing Large Language Models},
author={Jean Mercat and Igor Vasiljevic and Sedrick Keh and Kushal Arora and Achal Dave and Adrien Gaidon and Thomas Kollar},
journal={arXiv preprint arXiv:2405.06640},
year={2024}
}
```
## Citations
Mamba
```
@article{mamba,
title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
author={Gu, Albert and Dao, Tri},
journal={arXiv preprint arXiv:2312.00752},
year={2023}
}
```
OpenLM
```
@misc{open_lm,
author = {Gururangan, Suchin and Wortsman, Mitchell and Gadre, Samir Yitzhak and Dave, Achal and Kilian, Maciej and Shi, Weijia and Mercat, Jean and Smyrnis, Georgios and Ilharco, Gabriel and Jordan, Matt and Heckel, Reinhard and Dimakis, Alex and Farhadi, Ali and Shankar, Vaishaal and Schmidt, Ludwig},
title = {{open_lm}: a minimal but performative language modeling (LM) repository},
year = {2023},
note = {GitHub repository},
url = {https://github.com/mlfoundations/open_lm/}
}
``` |