ssmits commited on
Commit
8ab0af2
·
verified ·
1 Parent(s): ae99d32

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +94 -197
README.md CHANGED
@@ -1,247 +1,144 @@
1
  ---
2
- library_name: transformers
3
- tags: []
4
  ---
5
 
6
- # Model Card for Model ID
7
 
8
- <!-- Provide a quick summary of what the model is/does. -->
9
 
 
10
 
 
11
 
12
- ## Model Details
13
-
14
- ### Model Description
15
-
16
- <!-- Provide a longer summary of what this model is. -->
17
-
18
- This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
19
-
20
- - **Developed by:** [More Information Needed]
21
- - **Funded by [optional]:** [More Information Needed]
22
- - **Shared by [optional]:** [More Information Needed]
23
- - **Model type:** [More Information Needed]
24
- - **Language(s) (NLP):** [More Information Needed]
25
- - **License:** [More Information Needed]
26
- - **Finetuned from model [optional]:** [More Information Needed]
27
-
28
-
29
-
30
- # Learning Rate Optimization for Language Model Fine-tuning
31
-
32
- This script implements an advanced learning rate optimization strategy for fine-tuning large language models, combining Bayesian optimization with Gaussian Process Regression (GPR) for precise learning rate selection.
33
-
34
- ## Key Features
35
-
36
- ### 1. Bayesian Optimization
37
- * Uses Optuna framework to perform systematic learning rate search
38
- * Implements Tree-structured Parzen Estimators (TPE) for efficient hyperparameter optimization
39
- * Automatically explores learning rates between 1e-6 and 1e-4 in log space
40
-
41
- ### 2. Advanced Loss Tracking
42
- * Evaluates model performance using mean loss from the final 20% of training steps
43
- * Handles training failures gracefully with proper memory management
44
-
45
- ### 3. Sophisticated Post-processing
46
- * Applies Gaussian Process Regression to model the learning rate-loss relationship
47
- * Calculates uncertainty estimates for each prediction
48
- * Implements Expected Improvement (EI) acquisition function for optimal learning rate selection
49
-
50
- ### 4. Memory Optimization
51
- * Implements gradient checkpointing for efficient memory usage
52
- * Includes automatic memory clearing between trials
53
-
54
- ## Technical Details
55
-
56
- The optimization process consists of three main phases:
57
- 1. Initial exploration using Bayesian optimization
58
- 2. Refinement using Gaussian Process Regression
59
- 3. Final selection using Expected Improvement criterion
60
-
61
- The script was designed this way because:
62
- * Bayesian optimization provides efficient exploration of the learning rate space
63
- * GPR adds uncertainty quantification and smooth interpolation between observed points
64
- * The combination allows for both exploration and exploitation of the learning rate space
65
-
66
- ## Advantages
67
-
68
- * More reliable than manual learning rate selection
69
- * Provides uncertainty estimates for each prediction
70
- * Automatically adapts to different model sizes and datasets
71
- * Generates visualizations for analysis
72
- * Saves comprehensive results for reproducibility
73
-
74
- This approach is particularly valuable for fine-tuning large language models where training costs are high and optimal learning rate selection is crucial for model performance.
75
-
76
- ### Model Sources [optional]
77
-
78
- <!-- Provide the basic links for the model. -->
79
-
80
- - **Repository:** [More Information Needed]
81
- - **Paper [optional]:** [More Information Needed]
82
- - **Demo [optional]:** [More Information Needed]
83
-
84
- ## Uses
85
-
86
- <!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
87
-
88
- ### Direct Use
89
-
90
- <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
91
-
92
- [More Information Needed]
93
-
94
- ### Downstream Use [optional]
95
-
96
- <!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
97
-
98
- [More Information Needed]
99
-
100
- ### Out-of-Scope Use
101
 
102
- <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
103
 
104
- [More Information Needed]
105
 
106
- ## Bias, Risks, and Limitations
107
 
108
- <!-- This section is meant to convey both technical and sociotechnical limitations. -->
109
 
110
- [More Information Needed]
111
 
112
- ### Recommendations
113
 
114
- <!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
115
 
116
- Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
117
 
118
- ## How to Get Started with the Model
119
 
120
- Use the code below to get started with the model.
121
 
122
- [More Information Needed]
 
 
 
 
123
 
124
- ## Training Details
125
 
126
- ### Training Data
127
 
128
- <!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
129
 
130
- [More Information Needed]
131
 
132
- ### Training Procedure
 
 
 
133
 
134
- <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
 
135
 
136
- #### Preprocessing [optional]
 
137
 
138
- [More Information Needed]
 
 
139
 
 
140
 
141
- #### Training Hyperparameters
142
 
143
- - **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
 
 
144
 
145
- #### Speeds, Sizes, Times [optional]
 
 
 
 
 
 
146
 
147
- <!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
 
 
 
 
 
 
 
 
 
148
 
149
- [More Information Needed]
 
 
150
 
151
- ## Evaluation
 
 
 
 
152
 
153
- <!-- This section describes the evaluation protocols and provides the results. -->
154
 
155
- ### Testing Data, Factors & Metrics
156
-
157
- #### Testing Data
158
-
159
- <!-- This should link to a Dataset Card if possible. -->
160
-
161
- [More Information Needed]
162
-
163
- #### Factors
164
-
165
- <!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
166
-
167
- [More Information Needed]
168
-
169
- #### Metrics
170
-
171
- <!-- These are the evaluation metrics being used, ideally with a description of why. -->
172
-
173
- [More Information Needed]
174
-
175
- ### Results
176
-
177
- [More Information Needed]
178
-
179
- #### Summary
180
-
181
-
182
-
183
- ## Model Examination [optional]
184
-
185
- <!-- Relevant interpretability work for the model goes here -->
186
-
187
- [More Information Needed]
188
-
189
- ## Environmental Impact
190
-
191
- <!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
192
-
193
- Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
194
-
195
- - **Hardware Type:** [More Information Needed]
196
- - **Hours used:** [More Information Needed]
197
- - **Cloud Provider:** [More Information Needed]
198
- - **Compute Region:** [More Information Needed]
199
- - **Carbon Emitted:** [More Information Needed]
200
-
201
- ## Technical Specifications [optional]
202
-
203
- ### Model Architecture and Objective
204
-
205
- [More Information Needed]
206
-
207
- ### Compute Infrastructure
208
-
209
- [More Information Needed]
210
-
211
- #### Hardware
212
-
213
- [More Information Needed]
214
-
215
- #### Software
216
-
217
- [More Information Needed]
218
 
219
- ## Citation [optional]
220
 
221
- <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
 
 
222
 
223
- **BibTeX:**
224
 
225
- [More Information Needed]
226
 
227
- **APA:**
228
 
229
- [More Information Needed]
 
 
230
 
231
- ## Glossary [optional]
 
 
232
 
233
- <!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
 
 
 
 
234
 
235
- [More Information Needed]
 
 
236
 
237
- ## More Information [optional]
238
 
239
- [More Information Needed]
240
 
241
- ## Model Card Authors [optional]
 
 
242
 
243
- [More Information Needed]
244
 
245
- ## Model Card Contact
246
 
247
- [More Information Needed]
 
1
  ---
2
+ license: apache-2.0
3
+ library_name: transformers_zamba2
4
  ---
5
 
6
+ # Model Card for Zamba2-1.2B
7
 
8
+ Zamba2-1.2B is a hybrid model composed of state-space ([Mamba](https://github.com/state-spaces/mamba)) 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-1.2B possesses three major improvements over Zamba1:
9
 
10
+ 1.) Mamba1 blocks have been replaced with Mamba2 blocks.
11
 
12
+ 2.) We apply a LoRA projector to each shared MLP and attention block, which allows the network to specialize at each invocation of the shared transformer layer across depth. LoRA enables us to add depth-specialization for only a minimal increase in total parameter count.
13
 
14
+ 3.) We utilize rotary position embeddings in the shared attention layer.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
 
16
+ Zamba2-1.2B differs from our [2.7B model](https://huggingface.co/Zyphra/Zamba2-2.7B) in three ways:
17
 
18
+ 1.) We have added rotary position embeddings
19
 
20
+ 2.) A single shared transformer block (instead of two that we alternate between)
21
 
22
+ 3.) Added LoRA projectors to attention blocks (instead of just a LoRA on the MLP block)
23
 
24
+ We found that while hybrid SSM-transformer models are perfectly capable of performing well without position embeddings, adding rotary embeddings to the shared attention block slightly improved performance. Secondly, we utilize a single attention block (instead of alternating between two independent transformer blocks) because this enables a higher flop count for the model at a given parameter budget and at smaller scales this becomes more important than the slightly faster latency.
25
 
26
+ Zamba2-1.2B 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-1.2B was annealed on a mixture of 100B high-quality tokens.
27
 
28
+ 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.
29
 
30
+ A standalone Pytorch implementation of Zamba2-1.2B may be found [here](https://github.com/Zyphra/Zamba2).
31
 
32
+ ## Quick start
33
 
34
+ ### Prerequisites
35
 
36
+ To download Zamba2-1.2B, clone Zyphra's fork of transformers:
37
+ 1. `git clone https://github.com/Zyphra/transformers_zamba2.git`
38
+ 2. `cd transformers_zamba2`
39
+ 3. Install the repository: `pip install -e .`
40
+ 4. `pip install accelerate`
41
 
 
42
 
43
+ 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.
44
 
45
+ To run on CPU, please specify `use_mamba_kernels=False` when loading the model using ``AutoModelForCausalLM.from_pretrained``.
46
 
 
47
 
48
+ ### Inference
49
+ ```python
50
+ from transformers import AutoTokenizer, AutoModelForCausalLM
51
+ import torch
52
 
53
+ tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-1.2B")
54
+ model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-1.2B", device_map="cuda", torch_dtype=torch.bfloat16)
55
 
56
+ input_text = "What factors contributed to the fall of the Roman Empire?"
57
+ input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
58
 
59
+ outputs = model.generate(**input_ids, max_new_tokens=100)
60
+ print(tokenizer.decode(outputs[0]))
61
+ ```
62
 
63
+ ### Fine-tuning with Learning Rate Optimization
64
 
65
+ The model includes an advanced learning rate optimization system for fine-tuning, implemented through the `LROptimizerCallback` class. This callback automatically handles learning rate optimization during training. Here's how to use it:
66
 
67
+ ```python
68
+ from transformers import AutoTokenizer, Trainer
69
+ from lr_optimizer import setup_training, LROptimizerCallback
70
 
71
+ # Method 1: Using the complete setup function
72
+ training_setup = setup_training(
73
+ model_name="Zyphra/Zamba2-1.2B",
74
+ dataset_name="your/dataset",
75
+ num_trials=10
76
+ )
77
+ trainer = training_setup['trainer']
78
 
79
+ # Method 2: Using the callback directly
80
+ callback = LROptimizerCallback(
81
+ num_trials=10,
82
+ lr_range=(1e-6, 1e-4)
83
+ )
84
+ trainer = Trainer(
85
+ model=model,
86
+ args=training_args,
87
+ callbacks=[callback]
88
+ )
89
 
90
+ # Start training with optimized configuration
91
+ trainer.train()
92
+ ```
93
 
94
+ The optimization process automatically:
95
+ - Explores learning rates between 1e-6 and 1e-4 using Bayesian optimization
96
+ - Applies Gaussian Process Regression for precise LR selection
97
+ - Implements memory optimization through gradient checkpointing
98
+ - Supports both fp16 and bf16 training
99
 
100
+ For detailed configuration options, see the [fine-tuning documentation](link-to-docs).
101
 
102
+ ## Model Details
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
103
 
104
+ Zamba2-1.2B 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. 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.
105
 
106
+ <center>
107
+ <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/Vay6htbnBcySR3Z6NEgwj.png" width="300" alt="Zamba architecture">
108
+ </center>
109
 
110
+ ## Performance
111
 
112
+ Zamba2-1.2B achieves leading and state-of-the-art performance among models of <2B parameters and is competitive with some models of significantly greater size. Moreover, due to its unique hybrid SSM architecture, Zamba2-1.2B achieves extremely low inference latency and rapid generation with a significantly smaller memory footprint than comparable transformer based models.
113
 
114
+ Zamba2-1.2B's high performance and small inference compute and memory footprint renders it an ideal generalist model for on-device applications.
115
 
116
+ <center>
117
+ <img src="https://cdn-uploads.huggingface.co/production/uploads/65bc13717c6ad1994b6619e9/7Japy8VaJzKaFEjJgtWBp.png" width="700" alt="Zamba performance">
118
+ </center>
119
 
120
+ <center>
121
+ <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/Viwo3-bpYLFUu7cLIUFVv.png" width="800" alt="Zamba performance">
122
+ </center>
123
 
124
+ <!--
125
+ <center>
126
+ <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/JVZUvVMPIpIJy9RDyohMJ.png" width="800" alt="Zamba performance">
127
+ </center>
128
+ -->
129
 
130
+ Time to First Token (TTFT) | Output Generation
131
+ :-------------------------:|:-------------------------:
132
+ ![image/png](https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/5lpWDLdtPPVAk8COJq7gZ.png) | ![image/png](https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/V2tS6eCOGbpKybEoZmOB7.png)
133
 
 
134
 
135
+ And memory overhead
136
 
137
+ <center>
138
+ <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/m0YUmAmiVnRg6l9m10CEt.png" width="400" alt="Zamba inference and memory cost">
139
+ </center>
140
 
 
141
 
142
+ ## Notice
143
 
144
+ Zamba2-1.2B 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.