gongjingyao
commited on
Commit
•
f26cdb7
1
Parent(s):
e168a47
Upload Transformer
Browse files- LMConfig.py +33 -0
- README.md +199 -0
- adapter_config.json +35 -0
- adapter_model.safetensors +3 -0
- generation_config.json +4 -0
- model.py +416 -0
LMConfig.py
ADDED
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import PretrainedConfig
|
2 |
+
from typing import List
|
3 |
+
|
4 |
+
|
5 |
+
class LMConfig(PretrainedConfig):
|
6 |
+
model_type = "babylm"
|
7 |
+
|
8 |
+
def __init__(
|
9 |
+
self,
|
10 |
+
dim: int = 512,
|
11 |
+
n_layers: int = 8,
|
12 |
+
n_heads: int = 8,
|
13 |
+
n_kv_heads: int = None,
|
14 |
+
vocab_size: int = 64000,
|
15 |
+
hidden_dim: int = None,
|
16 |
+
multiple_of: int = 64,
|
17 |
+
norm_eps: float = 1e-5,
|
18 |
+
max_seq_len: int = 512,
|
19 |
+
dropout: float = 0.0,
|
20 |
+
**kwargs,
|
21 |
+
):
|
22 |
+
|
23 |
+
self.dim = dim
|
24 |
+
self.n_layers = n_layers
|
25 |
+
self.n_heads = n_heads
|
26 |
+
self.n_kv_heads = n_kv_heads
|
27 |
+
self.vocab_size = vocab_size
|
28 |
+
self.hidden_dim = hidden_dim
|
29 |
+
self.multiple_of = multiple_of
|
30 |
+
self.norm_eps = norm_eps
|
31 |
+
self.max_seq_len = max_seq_len
|
32 |
+
self.dropout = dropout
|
33 |
+
super().__init__(**kwargs)
|
README.md
ADDED
@@ -0,0 +1,199 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
### Model Sources [optional]
|
29 |
+
|
30 |
+
<!-- Provide the basic links for the model. -->
|
31 |
+
|
32 |
+
- **Repository:** [More Information Needed]
|
33 |
+
- **Paper [optional]:** [More Information Needed]
|
34 |
+
- **Demo [optional]:** [More Information Needed]
|
35 |
+
|
36 |
+
## Uses
|
37 |
+
|
38 |
+
<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
|
39 |
+
|
40 |
+
### Direct Use
|
41 |
+
|
42 |
+
<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
|
43 |
+
|
44 |
+
[More Information Needed]
|
45 |
+
|
46 |
+
### Downstream Use [optional]
|
47 |
+
|
48 |
+
<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
|
49 |
+
|
50 |
+
[More Information Needed]
|
51 |
+
|
52 |
+
### Out-of-Scope Use
|
53 |
+
|
54 |
+
<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
|
55 |
+
|
56 |
+
[More Information Needed]
|
57 |
+
|
58 |
+
## Bias, Risks, and Limitations
|
59 |
+
|
60 |
+
<!-- This section is meant to convey both technical and sociotechnical limitations. -->
|
61 |
+
|
62 |
+
[More Information Needed]
|
63 |
+
|
64 |
+
### Recommendations
|
65 |
+
|
66 |
+
<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
|
67 |
+
|
68 |
+
Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
|
69 |
+
|
70 |
+
## How to Get Started with the Model
|
71 |
+
|
72 |
+
Use the code below to get started with the model.
|
73 |
+
|
74 |
+
[More Information Needed]
|
75 |
+
|
76 |
+
## Training Details
|
77 |
+
|
78 |
+
### Training Data
|
79 |
+
|
80 |
+
<!-- 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. -->
|
81 |
+
|
82 |
+
[More Information Needed]
|
83 |
+
|
84 |
+
### Training Procedure
|
85 |
+
|
86 |
+
<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
|
87 |
+
|
88 |
+
#### Preprocessing [optional]
|
89 |
+
|
90 |
+
[More Information Needed]
|
91 |
+
|
92 |
+
|
93 |
+
#### Training Hyperparameters
|
94 |
+
|
95 |
+
- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
|
96 |
+
|
97 |
+
#### Speeds, Sizes, Times [optional]
|
98 |
+
|
99 |
+
<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
|
100 |
+
|
101 |
+
[More Information Needed]
|
102 |
+
|
103 |
+
## Evaluation
|
104 |
+
|
105 |
+
<!-- This section describes the evaluation protocols and provides the results. -->
|
106 |
+
|
107 |
+
### Testing Data, Factors & Metrics
|
108 |
+
|
109 |
+
#### Testing Data
|
110 |
+
|
111 |
+
<!-- This should link to a Dataset Card if possible. -->
|
112 |
+
|
113 |
+
[More Information Needed]
|
114 |
+
|
115 |
+
#### Factors
|
116 |
+
|
117 |
+
<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
|
118 |
+
|
119 |
+
[More Information Needed]
|
120 |
+
|
121 |
+
#### Metrics
|
122 |
+
|
123 |
+
<!-- These are the evaluation metrics being used, ideally with a description of why. -->
|
124 |
+
|
125 |
+
[More Information Needed]
|
126 |
+
|
127 |
+
### Results
|
128 |
+
|
129 |
+
[More Information Needed]
|
130 |
+
|
131 |
+
#### Summary
|
132 |
+
|
133 |
+
|
134 |
+
|
135 |
+
## Model Examination [optional]
|
136 |
+
|
137 |
+
<!-- Relevant interpretability work for the model goes here -->
|
138 |
+
|
139 |
+
[More Information Needed]
|
140 |
+
|
141 |
+
## Environmental Impact
|
142 |
+
|
143 |
+
<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
|
144 |
+
|
145 |
+
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).
|
146 |
+
|
147 |
+
- **Hardware Type:** [More Information Needed]
|
148 |
+
- **Hours used:** [More Information Needed]
|
149 |
+
- **Cloud Provider:** [More Information Needed]
|
150 |
+
- **Compute Region:** [More Information Needed]
|
151 |
+
- **Carbon Emitted:** [More Information Needed]
|
152 |
+
|
153 |
+
## Technical Specifications [optional]
|
154 |
+
|
155 |
+
### Model Architecture and Objective
|
156 |
+
|
157 |
+
[More Information Needed]
|
158 |
+
|
159 |
+
### Compute Infrastructure
|
160 |
+
|
161 |
+
[More Information Needed]
|
162 |
+
|
163 |
+
#### Hardware
|
164 |
+
|
165 |
+
[More Information Needed]
|
166 |
+
|
167 |
+
#### Software
|
168 |
+
|
169 |
+
[More Information Needed]
|
170 |
+
|
171 |
+
## Citation [optional]
|
172 |
+
|
173 |
+
<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
|
174 |
+
|
175 |
+
**BibTeX:**
|
176 |
+
|
177 |
+
[More Information Needed]
|
178 |
+
|
179 |
+
**APA:**
|
180 |
+
|
181 |
+
[More Information Needed]
|
182 |
+
|
183 |
+
## Glossary [optional]
|
184 |
+
|
185 |
+
<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
|
186 |
+
|
187 |
+
[More Information Needed]
|
188 |
+
|
189 |
+
## More Information [optional]
|
190 |
+
|
191 |
+
[More Information Needed]
|
192 |
+
|
193 |
+
## Model Card Authors [optional]
|
194 |
+
|
195 |
+
[More Information Needed]
|
196 |
+
|
197 |
+
## Model Card Contact
|
198 |
+
|
199 |
+
[More Information Needed]
|
adapter_config.json
ADDED
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"alpha_pattern": {},
|
3 |
+
"auto_mapping": null,
|
4 |
+
"base_model_name_or_path": "babylm",
|
5 |
+
"bias": "none",
|
6 |
+
"fan_in_fan_out": false,
|
7 |
+
"inference_mode": true,
|
8 |
+
"init_lora_weights": true,
|
9 |
+
"layer_replication": null,
|
10 |
+
"layers_pattern": null,
|
11 |
+
"layers_to_transform": null,
|
12 |
+
"loftq_config": {},
|
13 |
+
"lora_alpha": 16,
|
14 |
+
"lora_dropout": 0.1,
|
15 |
+
"megatron_config": null,
|
16 |
+
"megatron_core": "megatron.core",
|
17 |
+
"modules_to_save": null,
|
18 |
+
"peft_type": "LORA",
|
19 |
+
"r": 8,
|
20 |
+
"rank_pattern": {},
|
21 |
+
"revision": null,
|
22 |
+
"target_modules": [
|
23 |
+
"w2",
|
24 |
+
"wo",
|
25 |
+
"w1",
|
26 |
+
"wq",
|
27 |
+
"output",
|
28 |
+
"w3",
|
29 |
+
"wk",
|
30 |
+
"wv"
|
31 |
+
],
|
32 |
+
"task_type": "CAUSAL_LM",
|
33 |
+
"use_dora": false,
|
34 |
+
"use_rslora": false
|
35 |
+
}
|
adapter_model.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:47c754e66921267847a12d761fa2e408d9ec205459c6c2a30a54f97ba872ba3c
|
3 |
+
size 4601272
|
generation_config.json
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"_from_model_config": true,
|
3 |
+
"transformers_version": "4.41.1"
|
4 |
+
}
|
model.py
ADDED
@@ -0,0 +1,416 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import math
|
2 |
+
import struct
|
3 |
+
import inspect
|
4 |
+
from .LMConfig import LMConfig
|
5 |
+
from typing import Any, Optional, Tuple
|
6 |
+
import numpy as np
|
7 |
+
import torch
|
8 |
+
import torch.nn.functional as F
|
9 |
+
from torch import nn
|
10 |
+
from transformers import PreTrainedModel
|
11 |
+
|
12 |
+
|
13 |
+
class RMSNorm(torch.nn.Module):
|
14 |
+
def __init__(self, dim: int, eps: float):
|
15 |
+
super().__init__()
|
16 |
+
self.eps = eps
|
17 |
+
self.weight = nn.Parameter(torch.ones(dim))
|
18 |
+
|
19 |
+
def _norm(self, x):
|
20 |
+
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
21 |
+
|
22 |
+
def forward(self, x):
|
23 |
+
output = self._norm(x.float()).type_as(x)
|
24 |
+
return output * self.weight
|
25 |
+
|
26 |
+
|
27 |
+
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
|
28 |
+
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
29 |
+
t = torch.arange(end, device=freqs.device) # type: ignore
|
30 |
+
freqs = torch.outer(t, freqs).float() # type: ignore
|
31 |
+
freqs_cos = torch.cos(freqs) # real part
|
32 |
+
freqs_sin = torch.sin(freqs) # imaginary part
|
33 |
+
return freqs_cos, freqs_sin
|
34 |
+
|
35 |
+
|
36 |
+
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
|
37 |
+
ndim = x.ndim
|
38 |
+
assert 0 <= 1 < ndim
|
39 |
+
assert freqs_cis.shape == (x.shape[1], x.shape[-1])
|
40 |
+
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
|
41 |
+
return freqs_cis.view(shape)
|
42 |
+
|
43 |
+
|
44 |
+
def apply_rotary_emb(
|
45 |
+
xq: torch.Tensor,
|
46 |
+
xk: torch.Tensor,
|
47 |
+
freqs_cos: torch.Tensor,
|
48 |
+
freqs_sin: torch.Tensor
|
49 |
+
) -> Tuple[torch.Tensor, torch.Tensor]:
|
50 |
+
# reshape xq and xk to match the complex representation
|
51 |
+
xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
|
52 |
+
xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
|
53 |
+
|
54 |
+
# reshape freqs_cos and freqs_sin for broadcasting
|
55 |
+
freqs_cos = reshape_for_broadcast(freqs_cos, xq_r)
|
56 |
+
freqs_sin = reshape_for_broadcast(freqs_sin, xq_r)
|
57 |
+
|
58 |
+
# apply rotation using real numbers
|
59 |
+
xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
|
60 |
+
xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
|
61 |
+
xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
|
62 |
+
xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos
|
63 |
+
|
64 |
+
# flatten last two dimensions
|
65 |
+
xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3)
|
66 |
+
xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3)
|
67 |
+
|
68 |
+
return xq_out.type_as(xq), xk_out.type_as(xk)
|
69 |
+
|
70 |
+
|
71 |
+
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
72 |
+
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
73 |
+
bs, slen, n_kv_heads, head_dim = x.shape
|
74 |
+
if n_rep == 1:
|
75 |
+
return x
|
76 |
+
return (
|
77 |
+
x[:, :, :, None, :]
|
78 |
+
.expand(bs, slen, n_kv_heads, n_rep, head_dim)
|
79 |
+
.reshape(bs, slen, n_kv_heads * n_rep, head_dim)
|
80 |
+
)
|
81 |
+
|
82 |
+
|
83 |
+
class Attention(nn.Module):
|
84 |
+
def __init__(self, args: LMConfig):
|
85 |
+
super().__init__()
|
86 |
+
self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
|
87 |
+
assert args.n_heads % self.n_kv_heads == 0
|
88 |
+
model_parallel_size = 1
|
89 |
+
self.n_local_heads = args.n_heads // model_parallel_size
|
90 |
+
self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
|
91 |
+
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
92 |
+
self.head_dim = args.dim // args.n_heads
|
93 |
+
self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
|
94 |
+
self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
|
95 |
+
self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
|
96 |
+
self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
|
97 |
+
self.attn_dropout = nn.Dropout(args.dropout)
|
98 |
+
self.resid_dropout = nn.Dropout(args.dropout)
|
99 |
+
self.dropout = args.dropout
|
100 |
+
|
101 |
+
# use flash attention or a manual implementation?
|
102 |
+
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
|
103 |
+
if not self.flash:
|
104 |
+
print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
|
105 |
+
mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
|
106 |
+
mask = torch.triu(mask, diagonal=1)
|
107 |
+
self.register_buffer("mask", mask)
|
108 |
+
|
109 |
+
def forward(
|
110 |
+
self,
|
111 |
+
x: torch.Tensor,
|
112 |
+
freqs_cos: torch.Tensor,
|
113 |
+
freqs_sin: torch.Tensor,
|
114 |
+
):
|
115 |
+
bsz, seqlen, _ = x.shape
|
116 |
+
|
117 |
+
# QKV
|
118 |
+
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
|
119 |
+
xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
|
120 |
+
xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
|
121 |
+
xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
|
122 |
+
|
123 |
+
# RoPE relative positional embeddings
|
124 |
+
xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
|
125 |
+
|
126 |
+
# grouped multiquery attention: expand out keys and values
|
127 |
+
xk = repeat_kv(xk, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)
|
128 |
+
xv = repeat_kv(xv, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)
|
129 |
+
|
130 |
+
# make heads into a batch dimension
|
131 |
+
xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim)
|
132 |
+
xk = xk.transpose(1, 2)
|
133 |
+
xv = xv.transpose(1, 2)
|
134 |
+
|
135 |
+
# flash implementation
|
136 |
+
if self.flash:
|
137 |
+
output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None,
|
138 |
+
dropout_p=self.dropout if self.training else 0.0,
|
139 |
+
is_causal=True)
|
140 |
+
else:
|
141 |
+
# manual implementation
|
142 |
+
scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
|
143 |
+
assert hasattr(self, 'mask')
|
144 |
+
scores = scores + self.mask[:, :, :seqlen, :seqlen] # (bs, n_local_heads, seqlen, cache_len + seqlen)
|
145 |
+
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
|
146 |
+
scores = self.attn_dropout(scores)
|
147 |
+
output = torch.matmul(scores, xv) # (bs, n_local_heads, seqlen, head_dim)
|
148 |
+
|
149 |
+
# restore time as batch dimension and concat heads
|
150 |
+
output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
|
151 |
+
|
152 |
+
# final projection into the residual stream
|
153 |
+
output = self.wo(output)
|
154 |
+
output = self.resid_dropout(output)
|
155 |
+
return output
|
156 |
+
|
157 |
+
|
158 |
+
class FeedForward(nn.Module):
|
159 |
+
def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
|
160 |
+
super().__init__()
|
161 |
+
if hidden_dim is None:
|
162 |
+
hidden_dim = 4 * dim
|
163 |
+
hidden_dim = int(2 * hidden_dim / 3)
|
164 |
+
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
|
165 |
+
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
|
166 |
+
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
|
167 |
+
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
|
168 |
+
self.dropout = nn.Dropout(dropout)
|
169 |
+
|
170 |
+
def forward(self, x):
|
171 |
+
return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
|
172 |
+
|
173 |
+
|
174 |
+
class TransformerBlock(nn.Module):
|
175 |
+
def __init__(self, layer_id: int, args: LMConfig):
|
176 |
+
super().__init__()
|
177 |
+
self.n_heads = args.n_heads
|
178 |
+
self.dim = args.dim
|
179 |
+
self.head_dim = args.dim // args.n_heads
|
180 |
+
self.attention = Attention(args)
|
181 |
+
self.feed_forward = FeedForward(
|
182 |
+
dim=args.dim,
|
183 |
+
hidden_dim=args.hidden_dim,
|
184 |
+
multiple_of=args.multiple_of,
|
185 |
+
dropout=args.dropout,
|
186 |
+
)
|
187 |
+
self.layer_id = layer_id
|
188 |
+
self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
189 |
+
self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
190 |
+
|
191 |
+
def forward(self, x, freqs_cos, freqs_sin):
|
192 |
+
h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin)
|
193 |
+
out = h + self.feed_forward.forward(self.ffn_norm(h))
|
194 |
+
return out
|
195 |
+
|
196 |
+
|
197 |
+
class Transformer(PreTrainedModel):
|
198 |
+
config_class = LMConfig
|
199 |
+
|
200 |
+
last_loss: Optional[torch.Tensor]
|
201 |
+
|
202 |
+
def __init__(self, params: LMConfig = None):
|
203 |
+
super().__init__(params)
|
204 |
+
if not params:
|
205 |
+
params = LMConfig()
|
206 |
+
self.params = params
|
207 |
+
self.vocab_size = params.vocab_size
|
208 |
+
self.n_layers = params.n_layers
|
209 |
+
|
210 |
+
self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
|
211 |
+
self.dropout = nn.Dropout(params.dropout)
|
212 |
+
self.layers = torch.nn.ModuleList()
|
213 |
+
for layer_id in range(params.n_layers):
|
214 |
+
self.layers.append(TransformerBlock(layer_id, params))
|
215 |
+
self.norm = RMSNorm(params.dim, eps=params.norm_eps)
|
216 |
+
self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
|
217 |
+
|
218 |
+
# share the unembedding parameters with the embedding parameters
|
219 |
+
self.tok_embeddings.weight = self.output.weight # https://paperswithcode.com/method/weight-tying
|
220 |
+
|
221 |
+
# some useful precompute for the RoPE relative positional embeddings
|
222 |
+
freqs_cos, freqs_sin = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len)
|
223 |
+
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
|
224 |
+
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
|
225 |
+
|
226 |
+
# init all weights
|
227 |
+
self.apply(self._init_weights)
|
228 |
+
# apply special scaled init to the residual projections, per GPT-2 paper
|
229 |
+
for pn, p in self.named_parameters():
|
230 |
+
if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
|
231 |
+
torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * params.n_layers))
|
232 |
+
|
233 |
+
# Initialize attribute for the loss of the last forward call. This will be set if the forward is called with a targets tensor.
|
234 |
+
self.last_loss = None
|
235 |
+
|
236 |
+
def _init_weights(self, module):
|
237 |
+
if isinstance(module, nn.Linear):
|
238 |
+
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
239 |
+
if module.bias is not None:
|
240 |
+
torch.nn.init.zeros_(module.bias)
|
241 |
+
elif isinstance(module, nn.Embedding):
|
242 |
+
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
243 |
+
|
244 |
+
def forward(self, tokens: Optional[torch.Tensor] = None,
|
245 |
+
targets: Optional[torch.Tensor] = None, **keyargs) -> torch.Tensor:
|
246 |
+
if 'input_ids' in keyargs:
|
247 |
+
tokens = keyargs['input_ids']
|
248 |
+
if 'attention_mask' in keyargs:
|
249 |
+
targets = keyargs['attention_mask']
|
250 |
+
|
251 |
+
_bsz, seqlen = tokens.shape
|
252 |
+
h = self.tok_embeddings(tokens)
|
253 |
+
h = self.dropout(h)
|
254 |
+
freqs_cos = self.freqs_cos[:seqlen]
|
255 |
+
freqs_sin = self.freqs_sin[:seqlen]
|
256 |
+
|
257 |
+
for layer in self.layers:
|
258 |
+
h = layer(h, freqs_cos, freqs_sin)
|
259 |
+
h = self.norm(h)
|
260 |
+
|
261 |
+
if targets is not None:
|
262 |
+
# if we are given some desired targets also calculate the loss
|
263 |
+
logits = self.output(h)
|
264 |
+
self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
|
265 |
+
else:
|
266 |
+
# inference-time mini-optimization: only forward the output on the very last position
|
267 |
+
logits = self.output(h[:, [-1], :]) # note: using list [-1] to preserve the time dim
|
268 |
+
self.last_loss = None
|
269 |
+
|
270 |
+
return logits
|
271 |
+
|
272 |
+
def configure_optimizers(self, weight_decay, learning_rate, device_type):
|
273 |
+
# start with all of the candidate parameters
|
274 |
+
param_dict = {pn: p for pn, p in self.named_parameters()}
|
275 |
+
# filter out those that do not require grad
|
276 |
+
param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
|
277 |
+
# create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
|
278 |
+
# i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
|
279 |
+
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
|
280 |
+
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
|
281 |
+
optim_groups = [
|
282 |
+
{'params': decay_params, 'weight_decay': weight_decay},
|
283 |
+
{'params': nodecay_params, 'weight_decay': 0.0}
|
284 |
+
]
|
285 |
+
num_decay_params = sum(p.numel() for p in decay_params)
|
286 |
+
num_nodecay_params = sum(p.numel() for p in nodecay_params)
|
287 |
+
print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
|
288 |
+
print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
|
289 |
+
# Create AdamW optimizer and use the fused version if it is available
|
290 |
+
fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
|
291 |
+
use_fused = fused_available and device_type == 'cuda'
|
292 |
+
extra_args = dict(fused=True) if use_fused else dict()
|
293 |
+
optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=(0.9, 0.999), **extra_args)
|
294 |
+
print(f"using fused AdamW: {use_fused}")
|
295 |
+
return optimizer
|
296 |
+
|
297 |
+
@torch.inference_mode()
|
298 |
+
def generate(self, idx, max_new_tokens=512, temperature=1.0, top_k=None):
|
299 |
+
"""
|
300 |
+
Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
|
301 |
+
the sequence max_new_tokens times, feeding the predictions back into the model each time.
|
302 |
+
Most likely you'll want to make sure to be in model.eval() mode of operation for this.
|
303 |
+
Also note this is a super inefficient version of sampling with no key/value cache.
|
304 |
+
"""
|
305 |
+
for _ in range(max_new_tokens):
|
306 |
+
# if the sequence context is growing too long we must crop it at block_size
|
307 |
+
idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:]
|
308 |
+
# forward the model to get the logits for the index in the sequence
|
309 |
+
logits = self(idx_cond)
|
310 |
+
logits = logits[:, -1, :] # crop to just the final time step
|
311 |
+
if temperature == 0.0:
|
312 |
+
# "sample" the single most likely index
|
313 |
+
_, idx_next = torch.topk(logits, k=1, dim=-1)
|
314 |
+
else:
|
315 |
+
# pluck the logits at the final step and scale by desired temperature
|
316 |
+
logits = logits / temperature
|
317 |
+
# optionally crop the logits to only the top k options
|
318 |
+
if top_k is not None:
|
319 |
+
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
|
320 |
+
logits[logits < v[:, [-1]]] = -float('Inf')
|
321 |
+
# apply softmax to convert logits to (normalized) probabilities
|
322 |
+
probs = F.softmax(logits, dim=-1)
|
323 |
+
idx_next = torch.multinomial(probs, num_samples=1)
|
324 |
+
# append sampled index to the running sequence and continue
|
325 |
+
idx = torch.cat((idx, idx_next), dim=1)
|
326 |
+
|
327 |
+
return idx
|
328 |
+
|
329 |
+
# @torch.inference_mode()
|
330 |
+
@torch.no_grad()
|
331 |
+
def stream_generate(self, idx, eos, max_new_tokens, temperature=1.0, top_k=None):
|
332 |
+
"""
|
333 |
+
Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
|
334 |
+
the sequence max_new_tokens times, feeding the predictions back into the model each time.
|
335 |
+
Most likely you'll want to make sure to be in model.eval() mode of operation for this.
|
336 |
+
Also note this is a super inefficient version of sampling with no key/value cache.
|
337 |
+
"""
|
338 |
+
idx_ = idx.shape[1]
|
339 |
+
for __ in range(max_new_tokens):
|
340 |
+
# if the sequence context is growing too long we must crop it at block_size
|
341 |
+
idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:]
|
342 |
+
# forward the model to get the logits for the index in the sequence
|
343 |
+
logits = self(idx_cond)
|
344 |
+
logits = logits[:, -1, :] # crop to just the final time step
|
345 |
+
if temperature == 0.0:
|
346 |
+
# "sample" the single most likely index
|
347 |
+
_, idx_next = torch.topk(logits, k=1, dim=-1)
|
348 |
+
else:
|
349 |
+
# pluck the logits at the final step and scale by desired temperature
|
350 |
+
logits = logits / temperature
|
351 |
+
# optionally crop the logits to only the top k options
|
352 |
+
if top_k is not None:
|
353 |
+
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
|
354 |
+
logits[logits < v[:, [-1]]] = -float('Inf')
|
355 |
+
|
356 |
+
# apply softmax to convert logits to (normalized) probabilities
|
357 |
+
probs = F.softmax(logits, dim=-1)
|
358 |
+
idx_next = torch.multinomial(probs, num_samples=1, generator=None)
|
359 |
+
# append sampled index to the running sequence and continue
|
360 |
+
idx = torch.cat((idx, idx_next), dim=1)
|
361 |
+
yield idx[:, idx_:]
|
362 |
+
|
363 |
+
if idx_next == eos:
|
364 |
+
break
|
365 |
+
|
366 |
+
def export(self, filepath='model.bin'):
|
367 |
+
"""export the model weights in fp32 into .bin file to be read from C"""
|
368 |
+
f = open(filepath, 'wb')
|
369 |
+
|
370 |
+
def serialize(t):
|
371 |
+
d = t.detach().cpu().view(-1).numpy().astype(np.float32)
|
372 |
+
b = struct.pack(f'{len(d)}f', *d)
|
373 |
+
f.write(b)
|
374 |
+
|
375 |
+
# first write out the header
|
376 |
+
hidden_dim = self.layers[0].feed_forward.w1.weight.shape[0]
|
377 |
+
p = self.params
|
378 |
+
n_kv_heads = p.n_heads if p.n_kv_heads is None else p.n_kv_heads
|
379 |
+
header = struct.pack('iiiiiii', p.dim, hidden_dim, p.n_layers, p.n_heads,
|
380 |
+
n_kv_heads, p.vocab_size, p.max_seq_len)
|
381 |
+
f.write(header)
|
382 |
+
|
383 |
+
# next write out the embedding weights
|
384 |
+
serialize(self.tok_embeddings.weight)
|
385 |
+
|
386 |
+
# now all the layers
|
387 |
+
# attention weights
|
388 |
+
for layer in self.layers:
|
389 |
+
serialize(layer.attention_norm.weight)
|
390 |
+
for layer in self.layers:
|
391 |
+
serialize(layer.attention.wq.weight)
|
392 |
+
for layer in self.layers:
|
393 |
+
serialize(layer.attention.wk.weight)
|
394 |
+
for layer in self.layers:
|
395 |
+
serialize(layer.attention.wv.weight)
|
396 |
+
for layer in self.layers:
|
397 |
+
serialize(layer.attention.wo.weight)
|
398 |
+
# ffn weights
|
399 |
+
for layer in self.layers:
|
400 |
+
serialize(layer.ffn_norm.weight)
|
401 |
+
for layer in self.layers:
|
402 |
+
serialize(layer.feed_forward.w1.weight)
|
403 |
+
for layer in self.layers:
|
404 |
+
serialize(layer.feed_forward.w2.weight)
|
405 |
+
for layer in self.layers:
|
406 |
+
serialize(layer.feed_forward.w3.weight)
|
407 |
+
# final rmsnorm
|
408 |
+
serialize(self.norm.weight)
|
409 |
+
# note: no need to write final classifier weights due to weight sharing
|
410 |
+
# freqs_cis
|
411 |
+
serialize(self.freqs_cos[:p.max_seq_len])
|
412 |
+
serialize(self.freqs_sin[:p.max_seq_len])
|
413 |
+
|
414 |
+
# write to binary file
|
415 |
+
f.close()
|
416 |
+
print(f"wrote {filepath}")
|