sedrickkeh commited on
Commit
189b8e7
1 Parent(s): c0202d0

Create README.md

Browse files
Files changed (1) hide show
  1. README.md +194 -0
README.md ADDED
@@ -0,0 +1,194 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ datasets:
4
+ - tiiuae/falcon-refinedweb
5
+ pipeline_tag: text-generation
6
+ library_name: openlm
7
+ tags:
8
+ - mamba
9
+ - linear
10
+ language:
11
+ - en
12
+ model-index:
13
+ - name: mamba-7b
14
+ results:
15
+ - task:
16
+ type: text-generation
17
+ dataset:
18
+ type: MMLU
19
+ name: MMLU
20
+ metrics:
21
+ - name: accuracy
22
+ type: accuracy
23
+ value: 33.3
24
+ verified: false
25
+ - task:
26
+ type: text-generation
27
+ dataset:
28
+ type: HellaSwag
29
+ name: HellaSwag
30
+ metrics:
31
+ - name: accuracy
32
+ type: accuracy
33
+ value: 77.9
34
+ verified: false
35
+ - task:
36
+ type: text-generation
37
+ dataset:
38
+ type: PIQA
39
+ name: PIQA
40
+ metrics:
41
+ - name: accuracy
42
+ type: accuracy
43
+ value: 81.0
44
+ verified: false
45
+ - task:
46
+ type: text-generation
47
+ dataset:
48
+ type: Winogrande
49
+ name: Winogrande
50
+ metrics:
51
+ - name: accuracy
52
+ type: accuracy
53
+ value: 71.8
54
+ verified: false
55
+ - task:
56
+ type: text-generation
57
+ dataset:
58
+ type: ai2_arc
59
+ name: ARC-E
60
+ metrics:
61
+ - name: accuracy
62
+ type: accuracy
63
+ value: 77.5
64
+ verified: false
65
+ - task:
66
+ type: text-generation
67
+ dataset:
68
+ type: ai2_arc
69
+ name: ARC-C
70
+ metrics:
71
+ - name: accuracy
72
+ type: accuracy
73
+ value: 46.7
74
+ verified: false
75
+ ---
76
+
77
+ # Mamba-7B
78
+ (insert cool midjourney pic here?)<br>
79
+ This is a 7B parameter model with the [Mamba](https://arxiv.org/abs/2312.00752) architecture, trained on 1.2T tokens of the [RefinedWeb](https://huggingface.co/datasets/tiiuae/falcon-refinedweb) dataset.
80
+ 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).
81
+ We follow their training recipe and release our version of Mamba-7B.
82
+
83
+ ## Model Details
84
+ - **Developed by**: [Toyota Research Institute](https://www.tri.global/our-work/robotics)
85
+ - **Model Type**: This is an auto-regressive language model based on the [Mamba](https://arxiv.org/abs/2312.00752) architecture.
86
+ - **Dataset**: Trained on 1.2T tokens of [RefinedWeb](https://huggingface.co/datasets/tiiuae/falcon-refinedweb)
87
+ - **Tokenizer**: `EleutherAI/gpt-neox-20b`
88
+ - **Library**: [OpenLM](https://github.com/mlfoundations/open_lm/)
89
+ - **License**: This model is licensed under [Apache License, Version 2.0](https://www.apache.org/licenses/LICENSE-2.0).
90
+
91
+ | Parameters | Hidden Size | Layers | Vocab Size | Sequence Length |
92
+ |------------|-------------|--------| ---------- | --------------- |
93
+ | 7B | 4096 | 64 | 50432 | 2048 |
94
+
95
+ ## Training Details
96
+ - Mamba-7B was trained using AWS SageMaker on 128 H100 80GB GPUs.
97
+ - Training began in March 2023 and lasted around 3 weeks (some down time due to crashes and loss spikes)
98
+ | **Hyperparameter** | **Value** |
99
+ |--------------------|------------|
100
+ | Precision | `bfloat16` |
101
+ | Optimizer | AdamW |
102
+ | Learning rate | 3e-4 |
103
+ | LR cooldown end | 1e-5 |
104
+ | QK-norm | False |
105
+ | Warmup steps | 2000 |
106
+ | Z-loss | 1e-4 |
107
+ | Batch size | 2M |
108
+
109
+
110
+ ## Usage
111
+ (Not functional yet. Not sure if this is the right flow. Will work on this.)<br>
112
+ This model was trained using [OpenLM](https://github.com/mlfoundations/open_lm/).
113
+
114
+ To use HuggingFace models trained with OpenLM, first install the OpenLM package
115
+ ```bash
116
+ pip install openlm
117
+ ```
118
+
119
+ Importing from `openlm_hf` will automatically import the necessary classes.
120
+
121
+ ```python
122
+ from openlm_hf import * # registers the Auto* classes
123
+ from transformers import AutoTokenizer, AutoModelForCausalLM
124
+ tokenizer = AutoTokenizer.from_pretrained("tri-ml/mamba-7b-rw")
125
+ model = AutoModelForCausalLM.from_pretrained("tri-ml/mamba-7b-rw").cuda()
126
+
127
+ inputs = tokenizer(["A beautiful flower"], return_tensors="pt")
128
+ gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
129
+ output = model.generate(inputs['input_ids'], **gen_kwargs)
130
+ output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
131
+ print(output)
132
+ # <s> A beautiful flower box made of white rose wood. It is a perfect gift for weddings, birthdays and anniversaries.
133
+ # All the roses are from our farm Roses Flanders. Therefor you know that these flowers last much longer than those in store or online!</s>
134
+ ```
135
+
136
+
137
+ ## Performance Evaluation
138
+ Our evaluations were done using the [Eleuther LM Eval Harness](https://github.com/EleutherAI/lm-evaluation-harness) repo.
139
+
140
+ Below we report the performance of Mamba 7B compared to other base models.
141
+
142
+ <style>
143
+ .evalTable th { background: white; }
144
+ .evalTable tr:nth-child(1) { background: #f3f3f3; }
145
+ .evalTable tr:nth-child(2) { background: #f3f3f3; }
146
+ .evalTable tr:nth-child(6) { background: #f3f3f3; }
147
+ </style>
148
+
149
+ <div class="evalTable">
150
+
151
+ | | MMLU (5-shot) | HellaSwag | PIQA | Winogrande | ARC-E | ARC-C |
152
+ | ----------------- | ---------------- | ------------- | -------- | -------------- | --------- | --------- |
153
+ | Mamba-1.4B | 25.2 | 59.0 | 73.9 | 61.4 | 65.5 | 32.9 |
154
+ | Mamba-2.8B | 26.3 | 66.2 | 75.8 | 63.4 | 69.7 | 36.3 |
155
+ | Llama2-7B | 45.9 | 76.0 | 79.1 | 69.1 | 76.3 | 46.3 |
156
+ | Gemma-7B | 62.9 | 80.7 | 81.9 | 73.7 | 81.1 | 53.2 |
157
+ | Mistral-7B | 62.4 | 81.0 | 82.1 | 74.0 | 80.9 | 53.8 |
158
+ | **Mamba-7B** | 33.3 | 77.9 | 81.0 | 71.8 | 77.5 | 46.7 |
159
+
160
+ </div>
161
+
162
+
163
+ ## How to Cite
164
+ If you use this model, please cite our paper on Linearizing Large Language Models.
165
+ ```
166
+ @article{Mercat2024Linearizing,
167
+ title={Linearizing Large Language Models},
168
+ author={Jean Mercat and Igor Vasiljevic and Sedrick Keh and Kushal Arora and Achal Dave and Adrien Gaidon and Thomas Kollar},
169
+ journal={ArXiv},
170
+ year={2024},
171
+ volume={},
172
+ }
173
+ ```
174
+
175
+ ## Citations
176
+ Mamba
177
+ ```
178
+ @article{mamba,
179
+ title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
180
+ author={Gu, Albert and Dao, Tri},
181
+ journal={arXiv preprint arXiv:2312.00752},
182
+ year={2023}
183
+ }
184
+ ```
185
+ OpenLM
186
+ ```
187
+ @misc{open_lm,
188
+ 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},
189
+ title = {{open_lm}: a minimal but performative language modeling (LM) repository},
190
+ year = {2023},
191
+ note = {GitHub repository},
192
+ url = {https://github.com/mlfoundations/open_lm/}
193
+ }
194
+ ```