gongjingyao commited on
Commit
f26cdb7
1 Parent(s): e168a47

Upload Transformer

Browse files
Files changed (6) hide show
  1. LMConfig.py +33 -0
  2. README.md +199 -0
  3. adapter_config.json +35 -0
  4. adapter_model.safetensors +3 -0
  5. generation_config.json +4 -0
  6. 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}")