OpenNLPLab commited on
Commit
3b1902b
1 Parent(s): 8f01081

Publish 3B2-300B

Browse files
config.json ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "TransnormerForCausalLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "configuration_transnormer.TransnormerConfig",
7
+ "AutoModelForCausalLM": "modeling_transnormer.TransnormerForCausalLM"
8
+ },
9
+ "pad_token_id": 0,
10
+ "bos_token_id": 1,
11
+ "eos_token_id": 2,
12
+ "vocab_size": 64000,
13
+ "use_cache": true,
14
+ "init_std": 0.02,
15
+ "decoder_embed_dim": 4096,
16
+ "decoder_layers": 30,
17
+ "decoder_attention_heads": 32,
18
+ "no_scale_embedding": true,
19
+ "add_bos_token": false,
20
+ "norm_type": "simplermsnorm",
21
+ "linear_use_lrpe_list": [
22
+ 1,
23
+ 0,
24
+ 0,
25
+ 0,
26
+ 0,
27
+ 0,
28
+ 0,
29
+ 0,
30
+ 0,
31
+ 0,
32
+ 0,
33
+ 0,
34
+ 0,
35
+ 0,
36
+ 0,
37
+ 0,
38
+ 0,
39
+ 0,
40
+ 0,
41
+ 0,
42
+ 0,
43
+ 0,
44
+ 0,
45
+ 0,
46
+ 0,
47
+ 0,
48
+ 0,
49
+ 0,
50
+ 0,
51
+ 0
52
+ ],
53
+ "hidden_dim": 4096,
54
+ "linear_act_fun": "swish",
55
+ "glu_dim": 11008,
56
+ "bias": false,
57
+ "torch_dtype": "bfloat16",
58
+ "transformers_version": "4.32.0"
59
+ }
configuration_transnormer.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 OpenNLPLab
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # coding=utf-8
16
+ """ Transnormer configuration"""
17
+
18
+ from transformers.configuration_utils import PretrainedConfig
19
+ from transformers.utils import logging
20
+
21
+ logger = logging.get_logger(__name__)
22
+
23
+
24
+ class TransnormerConfig(PretrainedConfig):
25
+ model_type = "transnormer"
26
+ keys_to_ignore_at_inference = ["past_key_values"]
27
+
28
+ def __init__(
29
+ self,
30
+ pad_token_id=0,
31
+ bos_token_id=1,
32
+ eos_token_id=2,
33
+ vocab_size=64000,
34
+ use_cache=True,
35
+ init_std=0.02,
36
+ # model config
37
+ decoder_embed_dim=1024,
38
+ decoder_layers=24,
39
+ decoder_attention_heads=8,
40
+ no_scale_embedding=False,
41
+ add_bos_token=False,
42
+ norm_type="simplermsnorm",
43
+ linear_use_lrpe_list=[],
44
+ hidden_dim=1024,
45
+ linear_act_fun="silu",
46
+ glu_dim=2816,
47
+ bias=False,
48
+ **kwargs,
49
+ ):
50
+ super().__init__(
51
+ pad_token_id=pad_token_id,
52
+ bos_token_id=bos_token_id,
53
+ eos_token_id=eos_token_id,
54
+ **kwargs,
55
+ )
56
+ # hf origin
57
+ self.vocab_size = vocab_size
58
+ self.use_cache = use_cache
59
+ self.init_std = init_std
60
+ # add
61
+ self.decoder_embed_dim = decoder_embed_dim
62
+ self.decoder_layers = decoder_layers
63
+ self.decoder_attention_heads = decoder_attention_heads
64
+ self.no_scale_embedding = no_scale_embedding
65
+ self.add_bos_token = add_bos_token
66
+ self.norm_type = norm_type
67
+ self.linear_use_lrpe_list = linear_use_lrpe_list
68
+ self.hidden_dim = hidden_dim
69
+ self.linear_act_fun = linear_act_fun
70
+ self.glu_dim = glu_dim
71
+ self.bias = bias
generation_config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "pad_token_id": 0,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "max_new_tokens": 2048,
6
+ "temperature": 1.0,
7
+ "repetition_penalty": 1.03,
8
+ "do_sample": true
9
+ }
lightning_attention.py ADDED
@@ -0,0 +1,540 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 OpenNLPLab
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # coding=utf-8
16
+ import torch
17
+ import triton
18
+ import triton.language as tl
19
+
20
+
21
+ @triton.jit
22
+ def _fwd_kernel(
23
+ Q,
24
+ K,
25
+ V,
26
+ Out,
27
+ S,
28
+ stride_qz,
29
+ stride_qh,
30
+ stride_qm,
31
+ stride_qk,
32
+ stride_kz,
33
+ stride_kh,
34
+ stride_kn,
35
+ stride_kk,
36
+ stride_vz,
37
+ stride_vh,
38
+ stride_vn,
39
+ stride_ve,
40
+ stride_oz,
41
+ stride_oh,
42
+ stride_om,
43
+ stride_oe,
44
+ stride_sh,
45
+ Z,
46
+ H,
47
+ N_CTX,
48
+ BLOCK_M: tl.constexpr,
49
+ BLOCK_DMODEL_QK: tl.constexpr,
50
+ BLOCK_N: tl.constexpr,
51
+ BLOCK_DMODEL_V: tl.constexpr,
52
+ IS_CAUSAL: tl.constexpr,
53
+ USE_DECAY: tl.constexpr,
54
+ ):
55
+ start_m = tl.program_id(0)
56
+ off_hz = tl.program_id(1)
57
+ off_h = off_hz % H
58
+ # initialize offsets
59
+ offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
60
+ offs_n = tl.arange(0, BLOCK_N)
61
+ offs_k = tl.arange(0, BLOCK_DMODEL_QK)
62
+ offs_e = tl.arange(0, BLOCK_DMODEL_V)
63
+ # get current offset of q k v
64
+ off_q = (off_hz * stride_qh + offs_m[:, None] * stride_qm +
65
+ offs_k[None, :] * stride_qk)
66
+ off_k = (off_hz * stride_kh + offs_n[:, None] * stride_kn +
67
+ offs_k[None, :] * stride_kk)
68
+ off_v = (off_hz * stride_vh + offs_n[:, None] * stride_vn +
69
+ offs_e[None, :] * stride_ve)
70
+ off_o = (off_hz * stride_oh + offs_m[:, None] * stride_om +
71
+ offs_e[None, :] * stride_oe)
72
+
73
+ # Initialize pointers to Q, K, V
74
+ q_ptrs = Q + off_q
75
+ k_ptrs = K + off_k
76
+ v_ptrs = V + off_v
77
+
78
+ # initialize pointer to m and l
79
+ acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_V], dtype=tl.float32)
80
+ # load q: it will stay in SRAM throughout
81
+ q = tl.load(q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0)
82
+ # loop over k, v and update accumulator
83
+ lo = 0
84
+ # print(start_m)
85
+ hi = (start_m + 1) * BLOCK_M if IS_CAUSAL else N_CTX
86
+ for start_n in range(lo, hi, BLOCK_N):
87
+ # -- load k, v --
88
+ k = tl.load(
89
+ k_ptrs + start_n * stride_kn,
90
+ mask=(start_n + offs_n)[:, None] < N_CTX,
91
+ other=0.0,
92
+ )
93
+ v = tl.load(
94
+ v_ptrs + start_n * stride_vn,
95
+ mask=(start_n + offs_n)[:, None] < N_CTX,
96
+ other=0.0,
97
+ )
98
+ # -- compute qk ---
99
+ # qk = tl.dot(q, k)
100
+ qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
101
+ # qk += tl.dot(q, k, trans_b=True)
102
+ qk += tl.dot(q, tl.trans(k))
103
+ if IS_CAUSAL:
104
+ index = offs_m[:, None] - (start_n + offs_n[None, :])
105
+ if USE_DECAY:
106
+ S_block_ptr = S + off_h * stride_sh
107
+ s = tl.load(S_block_ptr)
108
+ s_index = s * index
109
+ s_index = tl.where(s_index >= 0, -s_index, float("-inf"))
110
+ qk = tl.exp(s_index) * qk
111
+ else:
112
+ qk = tl.where(index >= 0, qk, 0)
113
+ acc += tl.dot(qk, v.to(qk.dtype))
114
+
115
+ out_ptrs = Out + off_o
116
+ tl.store(out_ptrs, acc.to(q.dtype), mask=offs_m[:, None] < N_CTX)
117
+
118
+
119
+ @triton.jit
120
+ def _bwd_kernel_kv(
121
+ Q,
122
+ K,
123
+ V,
124
+ S,
125
+ DO,
126
+ DQ,
127
+ DK,
128
+ DV,
129
+ stride_qz,
130
+ stride_qh,
131
+ stride_qm,
132
+ stride_qk,
133
+ stride_kz,
134
+ stride_kh,
135
+ stride_kn,
136
+ stride_kk,
137
+ stride_vz,
138
+ stride_vh,
139
+ stride_vn,
140
+ stride_ve,
141
+ stride_oz,
142
+ stride_oh,
143
+ stride_om,
144
+ stride_oe,
145
+ stride_sh,
146
+ Z,
147
+ H,
148
+ N_CTX,
149
+ num_block,
150
+ BLOCK_M: tl.constexpr,
151
+ BLOCK_DMODEL_QK: tl.constexpr,
152
+ BLOCK_N: tl.constexpr,
153
+ BLOCK_DMODEL_V: tl.constexpr,
154
+ CAUSAL: tl.constexpr,
155
+ USE_DECAY: tl.constexpr,
156
+ ):
157
+ start_n = tl.program_id(0)
158
+ off_hz = tl.program_id(1)
159
+
160
+ off_z = off_hz // H
161
+ off_h = off_hz % H
162
+ # offset pointers for batch/head
163
+ Q += off_z * stride_qz + off_h * stride_qh
164
+ K += off_z * stride_kz + off_h * stride_kh
165
+ V += off_z * stride_vz + off_h * stride_vh
166
+ DO += off_z * stride_oz + off_h * stride_oh
167
+ DQ += off_z * stride_qz + off_h * stride_qh
168
+ DK += off_z * stride_kz + off_h * stride_kh
169
+ DV += off_z * stride_vz + off_h * stride_vh
170
+
171
+ # start of q
172
+ if CAUSAL:
173
+ lo = start_n * BLOCK_M
174
+ else:
175
+ lo = 0
176
+ # initialize row/col offsets
177
+ # seqlence offset
178
+ offs_qm = lo + tl.arange(0, BLOCK_M)
179
+ offs_kvn = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
180
+ # feature offset
181
+ offs_qkk = tl.arange(0, BLOCK_DMODEL_QK)
182
+ offs_ve = tl.arange(0, BLOCK_DMODEL_V)
183
+ # row block index
184
+ offs_m = tl.arange(0, BLOCK_M)
185
+ # initialize pointers to value-like data
186
+ q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_qkk[None, :] * stride_qk)
187
+ k_ptrs = K + (offs_kvn[:, None] * stride_kn +
188
+ offs_qkk[None, :] * stride_kk)
189
+ v_ptrs = V + (offs_kvn[:, None] * stride_vn + offs_ve[None, :] * stride_ve)
190
+ do_ptrs = DO + (offs_qm[:, None] * stride_om +
191
+ offs_ve[None, :] * stride_oe)
192
+ dq_ptrs = DQ + (offs_qm[:, None] * stride_qm +
193
+ offs_qkk[None, :] * stride_qk)
194
+ # initialize dv amd dk
195
+ dv = tl.zeros([BLOCK_N, BLOCK_DMODEL_V], dtype=tl.float32)
196
+ dk = tl.zeros([BLOCK_N, BLOCK_DMODEL_QK], dtype=tl.float32)
197
+ # k and v stay in SRAM throughout
198
+ k = tl.load(k_ptrs, mask=offs_kvn[:, None] < N_CTX, other=0.0)
199
+ v = tl.load(v_ptrs, mask=offs_kvn[:, None] < N_CTX, other=0.0)
200
+ # loop over rows
201
+ for start_m in range(lo, num_block * BLOCK_M, BLOCK_M):
202
+ offs_m_curr = start_m + offs_m
203
+ # load q, k, v, do on-chip
204
+ q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < N_CTX, other=0.0)
205
+ qk = tl.dot(q, tl.trans(k))
206
+ # qk = tl.dot(q, k, trans_b=True)
207
+ if CAUSAL:
208
+ index = offs_m_curr[:, None] - offs_kvn[None, :]
209
+ if USE_DECAY:
210
+ S_block_ptr = S + off_h * stride_sh
211
+ s = tl.load(S_block_ptr)
212
+ s_index = s * index
213
+ s_index = tl.where(s_index >= 0, -s_index, float("-inf"))
214
+ s = tl.exp(s_index)
215
+ qk = qk * s
216
+ else:
217
+ qk = tl.where(index >= 0, qk, 0)
218
+
219
+ p = qk
220
+ # compute dv
221
+ do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < N_CTX, other=0.0)
222
+ dv += tl.dot(tl.trans(p.to(do.dtype)), do)
223
+ dp = tl.dot(do, tl.trans(v).to(do.dtype))
224
+ if CAUSAL:
225
+ if USE_DECAY:
226
+ dp = dp * s
227
+ else:
228
+ dp = tl.where(index >= 0, dp, 0)
229
+
230
+ dk += tl.dot(tl.trans(dp.to(q.dtype)), q).to(tl.float32)
231
+
232
+ # increment pointers
233
+ q_ptrs += BLOCK_M * stride_qm
234
+ do_ptrs += BLOCK_M * stride_om
235
+ # write-back
236
+ dv_ptrs = DV + (offs_kvn[:, None] * stride_vn +
237
+ offs_ve[None, :] * stride_ve)
238
+ dk_ptrs = DK + (offs_kvn[:, None] * stride_kn +
239
+ offs_qkk[None, :] * stride_kk)
240
+ tl.store(dv_ptrs, dv, mask=offs_kvn[:, None] < N_CTX)
241
+ tl.store(dk_ptrs, dk, mask=offs_kvn[:, None] < N_CTX)
242
+
243
+
244
+ @triton.jit
245
+ def _bwd_kernel_q(
246
+ Q,
247
+ K,
248
+ V,
249
+ S,
250
+ DO,
251
+ DQ,
252
+ DK,
253
+ DV,
254
+ stride_qz,
255
+ stride_qh,
256
+ stride_qm,
257
+ stride_qk,
258
+ stride_kz,
259
+ stride_kh,
260
+ stride_kn,
261
+ stride_kk,
262
+ stride_vz,
263
+ stride_vh,
264
+ stride_vn,
265
+ stride_ve,
266
+ stride_oz,
267
+ stride_oh,
268
+ stride_om,
269
+ stride_oe,
270
+ stride_sh,
271
+ Z,
272
+ H,
273
+ N_CTX,
274
+ num_block,
275
+ BLOCK_M: tl.constexpr,
276
+ BLOCK_DMODEL_QK: tl.constexpr,
277
+ BLOCK_N: tl.constexpr,
278
+ BLOCK_DMODEL_V: tl.constexpr,
279
+ CAUSAL: tl.constexpr,
280
+ USE_DECAY: tl.constexpr,
281
+ ):
282
+ start_m = tl.program_id(0)
283
+ off_hz = tl.program_id(1)
284
+ off_z = off_hz // H
285
+ off_h = off_hz % H
286
+ # offset pointers for batch/head
287
+ K += off_z * stride_kz + off_h * stride_kh
288
+ V += off_z * stride_vz + off_h * stride_vh
289
+ DO += off_z * stride_oz + off_h * stride_oh
290
+ DQ += off_z * stride_qz + off_h * stride_qh
291
+ # feature offset
292
+ offs_qkk = tl.arange(0, BLOCK_DMODEL_QK)
293
+ offs_ve = tl.arange(0, BLOCK_DMODEL_V)
294
+ # row block index
295
+ offs_m = tl.arange(0, BLOCK_M)
296
+ # row block index
297
+ offs_qm = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
298
+ # do
299
+ do_ptrs = DO + (offs_qm[:, None] * stride_om +
300
+ offs_ve[None, :] * stride_oe)
301
+ dq_ptrs = DQ + (offs_qm[:, None] * stride_qm +
302
+ offs_qkk[None, :] * stride_qk)
303
+
304
+ do = tl.load(do_ptrs, mask=offs_qm[:, None] < N_CTX, other=0.0)
305
+
306
+ dq = tl.zeros([BLOCK_M, BLOCK_DMODEL_QK], dtype=tl.float32)
307
+ lo = 0
308
+ hi = (start_m + 1) * BLOCK_M if CAUSAL else N_CTX
309
+
310
+ offs_m_curr = start_m * BLOCK_M + offs_m
311
+
312
+ for start_n in range(0, num_block):
313
+ offs_kvn = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
314
+ k_ptrs = K + (offs_kvn[:, None] * stride_kn +
315
+ offs_qkk[None, :] * stride_kk)
316
+ v_ptrs = V + (offs_kvn[:, None] * stride_vn +
317
+ offs_ve[None, :] * stride_ve)
318
+ # k and v stay in SRAM throughout
319
+ k = tl.load(k_ptrs, mask=offs_kvn[:, None] < N_CTX, other=0.0)
320
+ v = tl.load(v_ptrs, mask=offs_kvn[:, None] < N_CTX, other=0.0)
321
+ # dp = do vT
322
+ dp = tl.dot(do, tl.trans(v).to(do.dtype))
323
+ if CAUSAL:
324
+ index = offs_m_curr[:, None] - offs_kvn[None, :]
325
+ if USE_DECAY:
326
+ S_block_ptr = S + off_h * stride_sh
327
+ s = tl.load(S_block_ptr)
328
+ s_index = s * index
329
+ s_index = tl.where(s_index >= 0, -s_index, float("-inf"))
330
+ s = tl.exp(s_index)
331
+ dp = dp * s
332
+ else:
333
+ dp = tl.where(index >= 0, dp, 0)
334
+ # dq = dq + dp k
335
+ dq += tl.dot(dp.to(k.dtype), k)
336
+
337
+ tl.store(dq_ptrs, dq, mask=offs_qm[:, None] < N_CTX)
338
+
339
+
340
+ class _attention(torch.autograd.Function):
341
+
342
+ @staticmethod
343
+ def forward(ctx, q, k, v, causal, s):
344
+ q = q.contiguous()
345
+ k = k.contiguous()
346
+ v = v.contiguous()
347
+ s = s.contiguous()
348
+ # only support for Ampere now
349
+ capability = torch.cuda.get_device_capability()
350
+ if capability[0] < 8:
351
+ raise RuntimeError(
352
+ "Lightning attention currently only supported for compute capability >= 80"
353
+ )
354
+ # shape constraints
355
+ Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
356
+ # right
357
+ o = torch.empty(
358
+ (q.shape[0], q.shape[1], q.shape[2], v.shape[-1]),
359
+ dtype=q.dtype,
360
+ device=q.device,
361
+ )
362
+
363
+ BLOCK_M = 128
364
+ BLOCK_N = 64
365
+ num_warps = 4 if Lk <= 64 else 8
366
+ num_stages = 1
367
+
368
+ grid = (triton.cdiv(q.shape[2], BLOCK_M), q.shape[0] * q.shape[1], 1)
369
+ use_decay = s.shape[0] > 0
370
+ _fwd_kernel[grid](
371
+ q,
372
+ k,
373
+ v,
374
+ o,
375
+ s,
376
+ q.stride(0),
377
+ q.stride(1),
378
+ q.stride(2),
379
+ q.stride(3),
380
+ k.stride(0),
381
+ k.stride(1),
382
+ k.stride(2),
383
+ k.stride(3),
384
+ v.stride(0),
385
+ v.stride(1),
386
+ v.stride(2),
387
+ v.stride(3),
388
+ o.stride(0),
389
+ o.stride(1),
390
+ o.stride(2),
391
+ o.stride(3),
392
+ s.stride(0),
393
+ q.shape[0],
394
+ q.shape[1],
395
+ q.shape[2],
396
+ BLOCK_M=BLOCK_M,
397
+ BLOCK_DMODEL_QK=Lk,
398
+ BLOCK_N=BLOCK_N,
399
+ BLOCK_DMODEL_V=Lv,
400
+ IS_CAUSAL=causal,
401
+ USE_DECAY=use_decay,
402
+ num_warps=num_warps,
403
+ num_stages=num_stages,
404
+ )
405
+
406
+ ctx.save_for_backward(q, k, v, s)
407
+ ctx.grid = grid
408
+ ctx.BLOCK_M = BLOCK_M
409
+ ctx.BLOCK_DMODEL_QK = Lk
410
+ ctx.BLOCK_N = BLOCK_N
411
+ ctx.BLOCK_DMODEL_V = Lv
412
+ ctx.causal = causal
413
+ ctx.use_decay = use_decay
414
+ return o
415
+
416
+ @staticmethod
417
+ def backward(ctx, do):
418
+ q, k, v, s = ctx.saved_tensors
419
+ BLOCK_M = 32
420
+ BLOCK_N = 32
421
+ num_warps = 4
422
+ num_stages = 1
423
+
424
+ do = do.contiguous()
425
+ dq = torch.zeros_like(q, dtype=torch.float32)
426
+ dk = torch.empty_like(k)
427
+ dv = torch.empty_like(v)
428
+
429
+ grid_kv = (triton.cdiv(k.shape[2],
430
+ BLOCK_N), k.shape[0] * k.shape[1], 1)
431
+ _bwd_kernel_kv[grid_kv](
432
+ q,
433
+ k,
434
+ v,
435
+ s,
436
+ do,
437
+ dq,
438
+ dk,
439
+ dv,
440
+ q.stride(0),
441
+ q.stride(1),
442
+ q.stride(2),
443
+ q.stride(3),
444
+ k.stride(0),
445
+ k.stride(1),
446
+ k.stride(2),
447
+ k.stride(3),
448
+ v.stride(0),
449
+ v.stride(1),
450
+ v.stride(2),
451
+ v.stride(3),
452
+ do.stride(0),
453
+ do.stride(1),
454
+ do.stride(2),
455
+ do.stride(3),
456
+ s.stride(0),
457
+ q.shape[0],
458
+ q.shape[1],
459
+ q.shape[2],
460
+ grid_kv[0],
461
+ BLOCK_M=BLOCK_M,
462
+ BLOCK_DMODEL_QK=ctx.BLOCK_DMODEL_QK,
463
+ BLOCK_N=BLOCK_N,
464
+ BLOCK_DMODEL_V=ctx.BLOCK_DMODEL_V,
465
+ CAUSAL=ctx.causal,
466
+ USE_DECAY=ctx.use_decay,
467
+ num_warps=num_warps,
468
+ num_stages=num_stages,
469
+ )
470
+
471
+ grid_q = (triton.cdiv(q.shape[2], BLOCK_M), q.shape[0] * q.shape[1], 1)
472
+
473
+ _bwd_kernel_q[grid_q](
474
+ q,
475
+ k,
476
+ v,
477
+ s,
478
+ do,
479
+ dq,
480
+ dk,
481
+ dv,
482
+ q.stride(0),
483
+ q.stride(1),
484
+ q.stride(2),
485
+ q.stride(3),
486
+ k.stride(0),
487
+ k.stride(1),
488
+ k.stride(2),
489
+ k.stride(3),
490
+ v.stride(0),
491
+ v.stride(1),
492
+ v.stride(2),
493
+ v.stride(3),
494
+ do.stride(0),
495
+ do.stride(1),
496
+ do.stride(2),
497
+ do.stride(3),
498
+ s.stride(0),
499
+ q.shape[0],
500
+ q.shape[1],
501
+ q.shape[2],
502
+ grid_q[0],
503
+ BLOCK_M=BLOCK_M,
504
+ BLOCK_DMODEL_QK=ctx.BLOCK_DMODEL_QK,
505
+ BLOCK_N=BLOCK_N,
506
+ BLOCK_DMODEL_V=ctx.BLOCK_DMODEL_V,
507
+ CAUSAL=ctx.causal,
508
+ USE_DECAY=ctx.use_decay,
509
+ num_warps=num_warps,
510
+ num_stages=num_stages,
511
+ )
512
+
513
+ return dq.to(q.dtype), dk, dv, None, None
514
+
515
+
516
+ attention = _attention.apply
517
+
518
+
519
+ def lightning_attention(q, k, v, causal, ed):
520
+ d = q.shape[-1]
521
+ e = v.shape[-1]
522
+ # arr = f(d)
523
+ if d >= 128:
524
+ m = 128
525
+ else:
526
+ m = 64
527
+ arr = [m * i for i in range(d // m + 1)]
528
+ if arr[-1] != d:
529
+ arr.append(d)
530
+ n = len(arr)
531
+ output = 0
532
+ for i in range(n - 1):
533
+ s = arr[i]
534
+ e = arr[i + 1]
535
+ q1 = q[..., s:e]
536
+ k1 = k[..., s:e]
537
+ o = attention(q1, k1, v, causal, ed)
538
+ output = output + o
539
+
540
+ return output
modeling_transnormer.py ADDED
@@ -0,0 +1,939 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 OpenNLPLab
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # coding=utf-8
16
+ """ PyTorch Transnormer model."""
17
+ import math
18
+ import os
19
+ from typing import List, Optional, Tuple, Union
20
+
21
+ from einops import rearrange
22
+ import numpy as np
23
+ import torch
24
+ from torch import nn
25
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
26
+ import torch.nn.functional as F
27
+ import torch.utils.checkpoint
28
+ from transformers.activations import ACT2FN
29
+ from transformers.modeling_outputs import (
30
+ BaseModelOutputWithPast,
31
+ CausalLMOutputWithPast,
32
+ )
33
+ from transformers.modeling_utils import PreTrainedModel
34
+ from transformers.utils import (
35
+ add_start_docstrings,
36
+ add_start_docstrings_to_model_forward,
37
+ logging,
38
+ replace_return_docstrings,
39
+ )
40
+
41
+ from .configuration_transnormer import TransnormerConfig
42
+ from .norm import SimpleRMSNorm as SimpleRMSNorm_torch
43
+ from .srmsnorm_triton import SimpleRMSNorm as SimpleRMSNorm_triton
44
+ from .utils import (
45
+ get_activation_fn,
46
+ get_norm_fn,
47
+ logging_info,
48
+ print_module,
49
+ print_params,
50
+ )
51
+
52
+ logger = logging.get_logger(__name__)
53
+
54
+ _CONFIG_FOR_DOC = "TransnormerConfig"
55
+
56
+ use_triton = eval(os.environ.get("use_triton", default="True"))
57
+ debug = eval(os.environ.get("debug", default="False"))
58
+
59
+ if use_triton:
60
+ try:
61
+ from .lightning_attention import lightning_attention
62
+
63
+ has_lightning_attention = True
64
+ except (ImportError, ModuleNotFoundError):
65
+ has_lightning_attention = False
66
+ else:
67
+ has_lightning_attention = False
68
+
69
+ if debug:
70
+ logger.info(f"Use triton: {use_triton}")
71
+ logger.info(f"Use lightning attention: {has_lightning_attention}")
72
+ logger.info(f"Debug mode: {debug}, {type(debug)}")
73
+
74
+ if not has_lightning_attention:
75
+
76
+ def linear_attention(q, k, v, attn_mask):
77
+ energy = torch.einsum("... n d, ... m d -> ... n m", q, k)
78
+ energy = energy * attn_mask
79
+ output = torch.einsum("... n m, ... m d -> ... n d", energy, v)
80
+
81
+ return output
82
+
83
+ ########## start Transnormer
84
+ ##### Linearized Relative Positional Encoding: https://openreview.net/forum?id=xoLyps2qWc&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DTMLR%2FAuthors%23your-submissions)
85
+ class Lrpe(nn.Module):
86
+ def __init__(
87
+ self,
88
+ num_heads=8,
89
+ embed_dim=64,
90
+ ):
91
+ super().__init__()
92
+ d = num_heads * embed_dim
93
+
94
+ self.index = torch.empty(0)
95
+ self.theta = nn.Parameter(
96
+ 10000 ** (-2 / d * torch.arange(d)).reshape(num_heads, 1, -1)
97
+ )
98
+
99
+ def extra_repr(self):
100
+ return print_module(self)
101
+
102
+ def forward(self, x, offset=0):
103
+ # x: b, h, n, d
104
+ # offset: for k, v cache
105
+ n = x.shape[-2]
106
+ if self.index.shape[0] < n:
107
+ self.index = torch.arange(n).reshape(1, -1, 1).to(x)
108
+ index = self.index[:, :n] + offset
109
+ theta = self.theta * index
110
+ x = torch.concat([x * torch.cos(theta), x * torch.sin(theta)], dim=-1)
111
+
112
+ return x
113
+
114
+
115
+ class GLU(nn.Module):
116
+ def __init__(self, d1, d2, bias=False):
117
+ super().__init__()
118
+ if debug:
119
+ # get local varables
120
+ params = locals()
121
+ # print params
122
+ print_params(**params)
123
+
124
+ self.l1 = nn.Linear(d1, d2, bias=bias)
125
+ self.l2 = nn.Linear(d1, d2, bias=bias)
126
+ self.l3 = nn.Linear(d2, d1, bias=bias)
127
+
128
+ def forward(self, x):
129
+ o1 = self.l1(x)
130
+ o2 = self.l2(x)
131
+ output = o1 * o2
132
+ output = self.l3(output)
133
+
134
+ return output
135
+
136
+
137
+ class NormLinearAttention(nn.Module):
138
+ def __init__(
139
+ self,
140
+ embed_dim,
141
+ hidden_dim,
142
+ num_heads,
143
+ linear_act_fun="silu",
144
+ norm_type="simplermsnorm",
145
+ linear_use_lrpe=False,
146
+ bias=False,
147
+ ):
148
+ super().__init__()
149
+ if debug:
150
+ # get local varables
151
+ params = locals()
152
+ # print params
153
+ print_params(**params)
154
+
155
+ self.out_proj = nn.Linear(hidden_dim, embed_dim, bias=bias)
156
+ self.act = get_activation_fn(linear_act_fun)
157
+ self.num_heads = num_heads
158
+ self.embed_dim = embed_dim
159
+ self.head_dim = self.embed_dim // self.num_heads
160
+ self.norm = get_norm_fn(norm_type)(hidden_dim)
161
+
162
+ self.linear_use_lrpe = linear_use_lrpe
163
+ if self.linear_use_lrpe:
164
+ self.lrpe = Lrpe(
165
+ num_heads=self.num_heads,
166
+ embed_dim=self.head_dim,
167
+ )
168
+
169
+ self.qkvu_proj = nn.Linear(embed_dim, 4 * hidden_dim, bias=bias)
170
+
171
+ # for inference only
172
+ self.offset = 0
173
+
174
+ def forward(
175
+ self,
176
+ x,
177
+ attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m)
178
+ attn_padding_mask: Optional[torch.Tensor] = None, # (b, m)
179
+ output_attentions: bool = False,
180
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
181
+ use_cache: bool = False,
182
+ slope_rate: Optional[torch.Tensor] = None,
183
+ ):
184
+ do_eval = eval(os.environ.get("do_eval", default="False"))
185
+ if (not self.training) and (not do_eval):
186
+ return self.inference(
187
+ x,
188
+ attn_mask,
189
+ attn_padding_mask,
190
+ output_attentions,
191
+ past_key_value,
192
+ use_cache,
193
+ slope_rate,
194
+ )
195
+ # x: b n d
196
+ n = x.shape[-2]
197
+ # linear map
198
+ q, k, v, u = self.qkvu_proj(x).chunk(4, dim=-1)
199
+ # reshape
200
+ q, k, v = map(
201
+ lambda x: rearrange(x, "b n (h d) -> b h n d", h=self.num_heads), [q, k, v]
202
+ )
203
+ # act
204
+ q = self.act(q)
205
+ k = self.act(k)
206
+
207
+ q_offset = 0
208
+ # lrpe relys on position, get cache first
209
+ if past_key_value is not None:
210
+ # reuse k, v, for evaluation only
211
+ k = torch.cat([past_key_value[0], k], dim=-2)
212
+ v = torch.cat([past_key_value[1], v], dim=-2)
213
+ q_offset = past_key_value[0].shape[-2]
214
+
215
+ past_key_value = (k, v) if use_cache else None
216
+
217
+ # lrpe
218
+ if self.linear_use_lrpe:
219
+ q = self.lrpe(q, offset=q_offset)
220
+ k = self.lrpe(k)
221
+
222
+ if attn_mask == None:
223
+ attn_mask = (torch.tril(torch.ones(n, n))).to(q)
224
+
225
+ if attn_padding_mask is not None:
226
+ v = v.masked_fill(
227
+ (1 - attn_padding_mask).unsqueeze(1).unsqueeze(-1).to(torch.bool), 0
228
+ )
229
+
230
+ if not has_lightning_attention:
231
+ if slope_rate != None:
232
+ attn_mask = torch.exp(slope_rate * attn_mask)
233
+ output = linear_attention(q, k, v, attn_mask)
234
+ else:
235
+ output = lightning_attention(
236
+ q, k, v, True, slope_rate.squeeze(-1).squeeze(-1)
237
+ )
238
+
239
+ # reshape
240
+ output = rearrange(output, "b h n d -> b n (h d)")
241
+ # normalize
242
+ output = self.norm(output)
243
+ # gate
244
+ output = u * output
245
+ # outproj
246
+ output = self.out_proj(output)
247
+
248
+ if not output_attentions:
249
+ attn_weights = None
250
+ else:
251
+ attn_weights = torch.einsum("... n d, ... m d -> ... n m", q, k)
252
+
253
+ return output, attn_weights, past_key_value
254
+
255
+ def inference(
256
+ self,
257
+ x,
258
+ attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m)
259
+ attn_padding_mask: Optional[torch.Tensor] = None, # (b, m)
260
+ output_attentions: bool = False,
261
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
262
+ use_cache: bool = False,
263
+ slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1)
264
+ ):
265
+ # x: b n d
266
+ n = x.shape[-2]
267
+ # linear map
268
+ q, k, v, u = self.qkvu_proj(x).chunk(4, dim=-1)
269
+ # reshape
270
+ q, k, v = map(
271
+ lambda x: rearrange(x, "b n (h d) -> b h n d", h=self.num_heads), [q, k, v]
272
+ )
273
+ # act
274
+ q = self.act(q)
275
+ k = self.act(k)
276
+
277
+ # rpe
278
+ if self.linear_use_lrpe:
279
+ q = self.lrpe(q, offset=self.offset)
280
+ k = self.lrpe(k)
281
+
282
+ if past_key_value == None:
283
+ self.offset = q.shape[-2]
284
+ else:
285
+ self.offset += 1
286
+
287
+ ratio = torch.exp(-slope_rate)
288
+
289
+ # only use for the first time
290
+ if past_key_value == None:
291
+ if attn_mask == None:
292
+ attn_mask = (torch.tril(torch.ones(n, n))).to(q)
293
+ if slope_rate != None:
294
+ attn_mask = torch.exp(slope_rate * attn_mask)
295
+
296
+ if attn_padding_mask is not None:
297
+ attn_mask = attn_mask.masked_fill(
298
+ (1 - attn_padding_mask).unsqueeze(1).unsqueeze(2).to(torch.bool),
299
+ 0,
300
+ )
301
+ energy = torch.einsum("... n d, ... m d -> ... n m", q, k)
302
+
303
+ if attn_mask != None:
304
+ energy = energy * attn_mask
305
+
306
+ output = torch.einsum("... n m, ... m d -> ... n d", energy, v)
307
+
308
+ eval_and_not_generate = eval(
309
+ os.environ.get("eval_and_not_generate", default="False")
310
+ )
311
+ if eval_and_not_generate:
312
+ kv = None
313
+ else:
314
+ # b, h, n, e, d
315
+ kv_outproduct = torch.einsum("... n e, ... n d -> ... n e d", k, v)
316
+ # 1, 1, n, 1, 1
317
+ index = torch.arange(n - 1, -1, -1).reshape(1, 1, -1, 1, 1).to(x)
318
+ # (h, 1, 1) -> (1, h, 1, 1, 1); (1, h, 1, 1, 1), (1, 1, n, 1, 1) -> (1, h, n, 1, 1)
319
+ decay = ratio.unsqueeze(0).unsqueeze(-1) ** index
320
+
321
+ kv_outproduct_with_decay = kv_outproduct * decay
322
+ kv = torch.sum(kv_outproduct_with_decay, dim=-3)
323
+ else:
324
+ kv = past_key_value
325
+
326
+ output = []
327
+ for i in range(n):
328
+ kv = ratio * kv + torch.einsum(
329
+ "... n d, ... n e -> ... d e",
330
+ k[:, :, i : i + 1],
331
+ v[:, :, i : i + 1],
332
+ )
333
+ qkv = torch.einsum(
334
+ "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv
335
+ )
336
+ output.append(qkv)
337
+ output = torch.concat(output, dim=-2)
338
+
339
+ # reshape
340
+ output = rearrange(output, "b h n d -> b n (h d)")
341
+ # normalize
342
+ output = self.norm(output)
343
+ # gate
344
+ output = u * output
345
+ # outproj
346
+ output = self.out_proj(output)
347
+
348
+ attn_weights = None
349
+
350
+ return output, attn_weights, kv
351
+
352
+
353
+ class TransnormerDecoderLayer(nn.Module):
354
+ def __init__(self, config: TransnormerConfig):
355
+ super().__init__()
356
+ self.embed_dim = config.decoder_embed_dim
357
+ ##### normalize
358
+ norm_type = config.norm_type
359
+ if debug:
360
+ logging_info(f"Decoder Norm Type: {norm_type}")
361
+ self.token_norm = get_norm_fn(norm_type)(self.embed_dim)
362
+ self.channel_norm = get_norm_fn(norm_type)(self.embed_dim)
363
+
364
+ ##### token mixer
365
+ self.token_mixer = self.build_token_mixer(
366
+ self.embed_dim,
367
+ config,
368
+ )
369
+
370
+ ##### channel mixer
371
+ self.glu_dim = config.glu_dim
372
+ if self.glu_dim == -1:
373
+ self.glu_dim = self.embed_dim
374
+ bias = config.bias
375
+ self.channel_mixer = GLU(self.embed_dim, self.glu_dim, bias)
376
+
377
+ def build_token_mixer(self, embed_dim, config):
378
+ return NormLinearAttention(
379
+ embed_dim=embed_dim,
380
+ hidden_dim=config.hidden_dim,
381
+ num_heads=config.decoder_attention_heads,
382
+ linear_act_fun=config.linear_act_fun,
383
+ norm_type=config.norm_type,
384
+ linear_use_lrpe=config.linear_use_lrpe,
385
+ bias=config.bias,
386
+ )
387
+
388
+ def residual_connection(self, x, residual):
389
+ return residual + x
390
+
391
+ def forward(
392
+ self,
393
+ x,
394
+ attn_mask: Optional[torch.Tensor] = None,
395
+ attn_padding_mask: Optional[torch.Tensor] = None,
396
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
397
+ output_attentions: Optional[bool] = False,
398
+ use_cache: Optional[bool] = False,
399
+ slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1)
400
+ ):
401
+ residual = x
402
+ input = x
403
+
404
+ o1, self_attn_weights, present_key_value = self.token_mixer(
405
+ x=self.token_norm(input),
406
+ attn_mask=attn_mask,
407
+ attn_padding_mask=attn_padding_mask,
408
+ past_key_value=past_key_value,
409
+ output_attentions=output_attentions,
410
+ use_cache=use_cache,
411
+ slope_rate=slope_rate,
412
+ )
413
+
414
+ o2 = self.channel_mixer(self.channel_norm(input))
415
+ o = o1 + o2
416
+ o = self.residual_connection(o, residual)
417
+
418
+ outputs = (o, )
419
+
420
+ if output_attentions:
421
+ outputs += (self_attn_weights,)
422
+
423
+ if use_cache:
424
+ outputs += (present_key_value,)
425
+
426
+ return outputs
427
+
428
+
429
+ TRANSNORMER_START_DOCSTRING = r"""
430
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
431
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
432
+ etc.)
433
+
434
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
435
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
436
+ and behavior.
437
+
438
+ Parameters:
439
+ config ([`TransnormerConfig`]):
440
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
441
+ load the weights associated with the model, only the configuration. Check out the
442
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
443
+ """
444
+
445
+
446
+ @add_start_docstrings(
447
+ TRANSNORMER_START_DOCSTRING,
448
+ )
449
+ class TransnormerPreTrainedModel(PreTrainedModel):
450
+ config_class = TransnormerConfig
451
+ base_model_prefix = "model"
452
+ supports_gradient_checkpointing = True
453
+ _no_split_modules = ["TransnormerDecoderLayer"]
454
+ _skip_keys_device_placement = "past_key_values"
455
+ _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
456
+
457
+ def _init_weights(self, module):
458
+ std = self.config.init_std
459
+ if isinstance(module, nn.Linear):
460
+ module.weight.data.normal_(mean=0.0, std=std)
461
+ if module.bias is not None:
462
+ module.bias.data.zero_()
463
+ elif isinstance(module, nn.Embedding):
464
+ module.weight.data.normal_(mean=0.0, std=std)
465
+ if module.padding_idx is not None:
466
+ module.weight.data[module.padding_idx].zero_()
467
+
468
+ def _set_gradient_checkpointing(self, module, value=False):
469
+ if isinstance(module, TransnormerModel):
470
+ module.gradient_checkpointing = value
471
+
472
+
473
+ TRANSNORMER_INPUTS_DOCSTRING = r"""
474
+ Args:
475
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
476
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
477
+ it.
478
+
479
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
480
+ [`PreTrainedTokenizer.__call__`] for details.
481
+
482
+ [What are input IDs?](../glossary#input-ids)
483
+ attn_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
484
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
485
+
486
+ - 1 for tokens that are **not masked**,
487
+ - 0 for tokens that are **masked**.
488
+
489
+ [What are attention masks?](../glossary#attention-mask)
490
+
491
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
492
+ [`PreTrainedTokenizer.__call__`] for details.
493
+
494
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
495
+ `past_key_values`).
496
+
497
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attn_mask`]
498
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
499
+ information on the default strategy.
500
+
501
+ - 1 indicates the head is **not masked**,
502
+ - 0 indicates the head is **masked**.
503
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
504
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
505
+ config.n_positions - 1]`.
506
+
507
+ [What are position IDs?](../glossary#position-ids)
508
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
509
+ Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
510
+ `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
511
+ `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
512
+
513
+ Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
514
+ blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
515
+
516
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
517
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
518
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
519
+ use_cache (`bool`, *optional*):
520
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
521
+ `past_key_values`).
522
+ output_attentions (`bool`, *optional*):
523
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
524
+ tensors for more detail.
525
+ output_hidden_states (`bool`, *optional*):
526
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
527
+ more detail.
528
+ return_dict (`bool`, *optional*):
529
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
530
+ """
531
+
532
+
533
+ @add_start_docstrings(
534
+ TRANSNORMER_START_DOCSTRING,
535
+ )
536
+ class TransnormerModel(TransnormerPreTrainedModel):
537
+ """
538
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`TransnormerDecoderLayer`]
539
+
540
+ Args:
541
+ config: TransnormerConfig
542
+ """
543
+
544
+ def __init__(self, config: TransnormerConfig):
545
+ super().__init__(config)
546
+ # hf origin
547
+ self.padding_idx = config.pad_token_id
548
+ self.vocab_size = config.vocab_size
549
+ self.gradient_checkpointing = False
550
+ # mask
551
+ self._linear_attn_mask = torch.empty(0)
552
+ # config
553
+ self.linear_use_lrpe_list = config.linear_use_lrpe_list
554
+ self.num_layers = config.decoder_layers
555
+ # h, 1, 1
556
+ self.slopes = self._build_slope_tensor(config.decoder_attention_heads)
557
+
558
+ # params
559
+ self.embed_tokens = nn.Embedding(
560
+ config.vocab_size, config.decoder_embed_dim, self.padding_idx
561
+ )
562
+ self.layers = nn.ModuleList([])
563
+ for i in range(config.decoder_layers):
564
+ if len(self.linear_use_lrpe_list) > 0:
565
+ config.linear_use_lrpe = self.linear_use_lrpe_list[i]
566
+ self.layers.append(TransnormerDecoderLayer(config))
567
+
568
+ self.final_norm = get_norm_fn(config.norm_type)(config.decoder_embed_dim)
569
+ self.embed_dim = config.decoder_embed_dim
570
+ self.embed_scale = (
571
+ 1.0 if config.no_scale_embedding else math.sqrt(self.embed_dim)
572
+ )
573
+
574
+ # Initialize weights and apply final processing
575
+ self.post_init()
576
+
577
+ @staticmethod
578
+ def _build_slope_tensor(n_attention_heads: int):
579
+ def get_slopes(n):
580
+ def get_slopes_power_of_2(n):
581
+ start = 2 ** (-(2 ** -(math.log2(n) - 3)))
582
+ ratio = start
583
+ return [start * ratio**i for i in range(n)]
584
+
585
+ if math.log2(n).is_integer():
586
+ return get_slopes_power_of_2(
587
+ n
588
+ ) # In the paper, we only train models that have 2^a heads for some a. This function has
589
+ else: # some good properties that only occur when the input is a power of 2. To maintain that even
590
+ closest_power_of_2 = 2 ** math.floor(
591
+ math.log2(n)
592
+ ) # when the number of heads is not a power of 2, we use this workaround.
593
+ return (
594
+ get_slopes_power_of_2(closest_power_of_2)
595
+ + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
596
+ )
597
+
598
+ # h, 1, 1
599
+ slopes = torch.tensor(get_slopes(n_attention_heads)).reshape(
600
+ n_attention_heads, 1, 1
601
+ )
602
+
603
+ return slopes
604
+
605
+ def extra_repr(self):
606
+ return print_module(self)
607
+
608
+ def get_input_embeddings(self):
609
+ return self.embed_tokens
610
+
611
+ def set_input_embeddings(self, value):
612
+ self.embed_tokens = value
613
+
614
+ def _prepare_decoder_linear_attn_mask(
615
+ self, input_shape, inputs_embeds, past_key_values_length
616
+ ):
617
+ bsz, tgt_len = input_shape
618
+ src_len = tgt_len + past_key_values_length
619
+
620
+ def power_log(x):
621
+ return 2 ** (math.ceil(math.log(x, 2)))
622
+
623
+ n = power_log(max(tgt_len, src_len))
624
+ if self._linear_attn_mask.shape[-1] < n:
625
+
626
+ def get_mask(n):
627
+ mask = torch.triu(torch.zeros(n, n).float().fill_(float("-inf")), 1)
628
+ # no slope version
629
+ # -n, ..., -2, -1, 0
630
+ for i in range(n):
631
+ x = torch.arange(i + 1)
632
+ y = x
633
+ mask[i, : i + 1] = -torch.flip(y, [0])
634
+
635
+ return mask
636
+
637
+ arr = []
638
+ for slope in self.slopes:
639
+ arr.append(get_mask(n))
640
+ self._linear_attn_mask = torch.stack(arr, dim=0).to(inputs_embeds)
641
+
642
+ linear_attn_mask = self._linear_attn_mask[:, -tgt_len:, -src_len:]
643
+ num_heads = linear_attn_mask.shape[0]
644
+
645
+ return linear_attn_mask[None, :, :, :].expand(bsz, num_heads, tgt_len, src_len)
646
+
647
+ @add_start_docstrings_to_model_forward(TRANSNORMER_INPUTS_DOCSTRING)
648
+ def forward(
649
+ self,
650
+ input_ids: torch.LongTensor = None,
651
+ attn_padding_mask: Optional[torch.Tensor] = None,
652
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
653
+ inputs_embeds: Optional[torch.FloatTensor] = None,
654
+ use_cache: Optional[bool] = None,
655
+ output_attentions: Optional[bool] = None,
656
+ output_hidden_states: Optional[bool] = None,
657
+ return_dict: Optional[bool] = None,
658
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
659
+ output_attentions = (
660
+ output_attentions
661
+ if output_attentions is not None
662
+ else self.config.output_attentions
663
+ )
664
+ output_hidden_states = (
665
+ output_hidden_states
666
+ if output_hidden_states is not None
667
+ else self.config.output_hidden_states
668
+ )
669
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
670
+
671
+ return_dict = (
672
+ return_dict if return_dict is not None else self.config.use_return_dict
673
+ )
674
+
675
+ # retrieve input_ids and inputs_embeds
676
+ if input_ids is not None and inputs_embeds is not None:
677
+ raise ValueError(
678
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
679
+ )
680
+ elif input_ids is not None:
681
+ batch_size, seq_length = input_ids.shape
682
+ elif inputs_embeds is not None:
683
+ batch_size, seq_length, _ = inputs_embeds.shape
684
+ else:
685
+ raise ValueError(
686
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
687
+ )
688
+
689
+ seq_length_with_past = seq_length
690
+ past_key_values_length = 0
691
+
692
+ if past_key_values is not None:
693
+ past_key_values_length = past_key_values[0][0].shape[-2]
694
+ seq_length_with_past = seq_length_with_past + past_key_values_length
695
+
696
+ if inputs_embeds is None:
697
+ # !!! use embed_scale
698
+ inputs_embeds = self.embed_scale * self.embed_tokens(input_ids)
699
+
700
+ hidden_states = inputs_embeds
701
+
702
+ if self.gradient_checkpointing and self.training:
703
+ if use_cache:
704
+ logger.warning_once(
705
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
706
+ )
707
+ use_cache = False
708
+
709
+ # decoder layers
710
+ all_hidden_states = () if output_hidden_states else None
711
+ all_self_attns = () if output_attentions else None
712
+ next_decoder_cache = () if use_cache else None
713
+
714
+ ##### norm linear layers
715
+ linear_attn_padding_mask = attn_padding_mask
716
+ linear_attn_mask = self._prepare_decoder_linear_attn_mask(
717
+ (batch_size, seq_length), inputs_embeds, past_key_values_length
718
+ )
719
+
720
+ slope_rates = [self.slopes.to(input_ids.device) for _ in range(self.num_layers)]
721
+
722
+ for idx, layer in enumerate(self.layers):
723
+ if output_hidden_states:
724
+ all_hidden_states += (hidden_states,)
725
+
726
+ past_key_value = (
727
+ past_key_values[idx] if past_key_values is not None else None
728
+ )
729
+
730
+ slope_rate = slope_rates[idx]
731
+ slope_rate = slope_rate * (1 - idx / (self.num_layers - 1) + 1e-5)
732
+ mask = linear_attn_mask
733
+
734
+ layer_outputs = layer(
735
+ hidden_states,
736
+ attn_mask=mask,
737
+ attn_padding_mask=linear_attn_padding_mask,
738
+ past_key_value=past_key_value,
739
+ output_attentions=output_attentions,
740
+ use_cache=use_cache,
741
+ slope_rate=slope_rate,
742
+ )
743
+
744
+ hidden_states = layer_outputs[0]
745
+
746
+ if use_cache:
747
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
748
+
749
+ if output_attentions:
750
+ all_self_attns += (layer_outputs[1],)
751
+
752
+ # if idx == 0:
753
+ # break
754
+
755
+ hidden_states = self.final_norm(hidden_states)
756
+
757
+ # add hidden states from the last decoder layer
758
+ if output_hidden_states:
759
+ all_hidden_states += (hidden_states,)
760
+
761
+ next_cache = next_decoder_cache if use_cache else None
762
+ if not return_dict:
763
+ return tuple(
764
+ v
765
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
766
+ if v is not None
767
+ )
768
+ return BaseModelOutputWithPast(
769
+ last_hidden_state=hidden_states,
770
+ past_key_values=next_cache,
771
+ hidden_states=all_hidden_states,
772
+ attentions=all_self_attns,
773
+ )
774
+
775
+
776
+ class TransnormerForCausalLM(TransnormerPreTrainedModel):
777
+ def __init__(self, config):
778
+ super().__init__(config)
779
+ self.model = TransnormerModel(config)
780
+ if debug:
781
+ logging_info(self.model)
782
+
783
+ # the lm_head weight is automatically tied to the embed tokens weight
784
+ self.lm_head = nn.Linear(
785
+ config.decoder_embed_dim, config.vocab_size, bias=False
786
+ )
787
+
788
+ # Initialize weights and apply final processing
789
+ self.post_init()
790
+
791
+ def get_input_embeddings(self):
792
+ return self.model.embed_tokens
793
+
794
+ def set_input_embeddings(self, value):
795
+ self.model.embed_tokens = value
796
+
797
+ def get_output_embeddings(self):
798
+ return self.lm_head
799
+
800
+ def set_output_embeddings(self, new_embeddings):
801
+ self.lm_head = new_embeddings
802
+
803
+ def set_decoder(self, decoder):
804
+ self.model = decoder
805
+
806
+ def get_decoder(self):
807
+ return self.model
808
+
809
+ @add_start_docstrings_to_model_forward(TRANSNORMER_INPUTS_DOCSTRING)
810
+ @replace_return_docstrings(
811
+ output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
812
+ )
813
+ def forward(
814
+ self,
815
+ input_ids: torch.LongTensor = None,
816
+ attention_mask: Optional[torch.Tensor] = None,
817
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
818
+ inputs_embeds: Optional[torch.FloatTensor] = None,
819
+ labels: Optional[torch.LongTensor] = None,
820
+ use_cache: Optional[bool] = None,
821
+ output_attentions: Optional[bool] = None,
822
+ output_hidden_states: Optional[bool] = None,
823
+ return_dict: Optional[bool] = None,
824
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
825
+ r"""
826
+ Args:
827
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
828
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
829
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
830
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
831
+
832
+ Returns:
833
+
834
+ Example:
835
+
836
+ ```python
837
+ >>> from transformers import AutoTokenizer, TransnormerForCausalLM
838
+
839
+ >>> model = TransnormerForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
840
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
841
+
842
+ >>> prompt = "Hey, are you consciours? Can you talk to me?"
843
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
844
+
845
+ >>> # Generate
846
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
847
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
848
+ "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you."
849
+ ```"""
850
+ output_attentions = (
851
+ output_attentions
852
+ if output_attentions is not None
853
+ else self.config.output_attentions
854
+ )
855
+ output_hidden_states = (
856
+ output_hidden_states
857
+ if output_hidden_states is not None
858
+ else self.config.output_hidden_states
859
+ )
860
+ return_dict = (
861
+ return_dict if return_dict is not None else self.config.use_return_dict
862
+ )
863
+
864
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
865
+ outputs = self.model(
866
+ input_ids=input_ids,
867
+ attn_padding_mask=attention_mask,
868
+ past_key_values=past_key_values,
869
+ inputs_embeds=inputs_embeds,
870
+ use_cache=use_cache,
871
+ output_attentions=output_attentions,
872
+ output_hidden_states=output_hidden_states,
873
+ return_dict=return_dict,
874
+ )
875
+
876
+ hidden_states = outputs[0]
877
+ logits = self.lm_head(hidden_states)
878
+
879
+ loss = None
880
+ if labels is not None:
881
+ # Shift so that tokens < n predict n
882
+ shift_logits = logits[..., :-1, :].contiguous()
883
+ shift_labels = labels[..., 1:].contiguous()
884
+ # Flatten the tokens
885
+ loss_fct = CrossEntropyLoss()
886
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
887
+ shift_labels = shift_labels.view(-1)
888
+ # Enable model parallelism
889
+ shift_labels = shift_labels.to(shift_logits.device)
890
+ loss = loss_fct(shift_logits, shift_labels)
891
+
892
+ if not return_dict:
893
+ output = (logits,) + outputs[1:]
894
+ return (loss,) + output if loss is not None else output
895
+
896
+ return CausalLMOutputWithPast(
897
+ loss=loss,
898
+ logits=logits,
899
+ past_key_values=outputs.past_key_values,
900
+ hidden_states=outputs.hidden_states,
901
+ attentions=outputs.attentions,
902
+ )
903
+
904
+ def prepare_inputs_for_generation(
905
+ self,
906
+ input_ids,
907
+ past_key_values=None,
908
+ attention_mask=None,
909
+ inputs_embeds=None,
910
+ **kwargs,
911
+ ):
912
+ if past_key_values:
913
+ input_ids = input_ids[:, -1:]
914
+
915
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
916
+ if inputs_embeds is not None and past_key_values is None:
917
+ model_inputs = {"inputs_embeds": inputs_embeds}
918
+ else:
919
+ model_inputs = {"input_ids": input_ids}
920
+
921
+ model_inputs.update(
922
+ {
923
+ "past_key_values": past_key_values,
924
+ "use_cache": kwargs.get("use_cache"),
925
+ "attention_mask": attention_mask,
926
+ }
927
+ )
928
+ return model_inputs
929
+
930
+ @staticmethod
931
+ def _reorder_cache(past_key_values, beam_idx):
932
+ reordered_past = ()
933
+ for layer_past in past_key_values:
934
+ reordered_past += (
935
+ tuple(
936
+ past_state.index_select(0, beam_idx) for past_state in layer_past
937
+ ),
938
+ )
939
+ return reordered_past
norm.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 OpenNLPLab
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # coding=utf-8
16
+ import logging
17
+ import os
18
+ import sys
19
+
20
+ import torch
21
+ from torch import nn
22
+
23
+ logging.basicConfig(
24
+ format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
25
+ datefmt="%Y-%m-%d %H:%M:%S",
26
+ level=os.environ.get("LOGLEVEL", "INFO").upper(),
27
+ stream=sys.stdout,
28
+ )
29
+ logger = logging.getLogger("srmsnorm")
30
+
31
+
32
+ class SimpleRMSNorm(nn.Module):
33
+
34
+ def __init__(self, dim: int, eps: float = 1e-6):
35
+ super().__init__()
36
+ self.eps = eps
37
+
38
+ def _norm(self, x):
39
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
40
+
41
+ def forward(self, x):
42
+ output = self._norm(x.float()).type_as(x)
43
+
44
+ return output
pytorch_model-00001-of-00003.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:548b1f7cf50fc6bfe95b9112312ee2f6c56e6e42e456a5c5d3c4da7309b07779
3
+ size 1973966419
pytorch_model-00002-of-00003.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0084c1d699eb9a9fc758ae283910b007c5f040504c4dafe2ac3df463f4210cd0
3
+ size 1972653283
pytorch_model-00003-of-00003.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ce1ccb72dc65cd5e928bb8f329a82380ceb4cdde7db1923af7e256caf2ddbfea
3
+ size 1875658651
pytorch_model.bin.index.json ADDED
@@ -0,0 +1,170 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 5822223360
4
+ },
5
+ "weight_map": {
6
+ "lm_head.weight": "pytorch_model-00001-of-00003.bin",
7
+ "model.embed_tokens.weight": "pytorch_model-00001-of-00003.bin",
8
+ "model.layers.0.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
9
+ "model.layers.0.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
10
+ "model.layers.0.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
11
+ "model.layers.0.token_mixer.lrpe.theta": "pytorch_model-00001-of-00003.bin",
12
+ "model.layers.0.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
13
+ "model.layers.0.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
14
+ "model.layers.1.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
15
+ "model.layers.1.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
16
+ "model.layers.1.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
17
+ "model.layers.1.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
18
+ "model.layers.1.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
19
+ "model.layers.10.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
20
+ "model.layers.10.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
21
+ "model.layers.10.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
22
+ "model.layers.10.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
23
+ "model.layers.10.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
24
+ "model.layers.11.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
25
+ "model.layers.11.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
26
+ "model.layers.11.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
27
+ "model.layers.11.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
28
+ "model.layers.11.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
29
+ "model.layers.12.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
30
+ "model.layers.12.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
31
+ "model.layers.12.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
32
+ "model.layers.12.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
33
+ "model.layers.12.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
34
+ "model.layers.13.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
35
+ "model.layers.13.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
36
+ "model.layers.13.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
37
+ "model.layers.13.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
38
+ "model.layers.13.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
39
+ "model.layers.14.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
40
+ "model.layers.14.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
41
+ "model.layers.14.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
42
+ "model.layers.14.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
43
+ "model.layers.14.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
44
+ "model.layers.15.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
45
+ "model.layers.15.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
46
+ "model.layers.15.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
47
+ "model.layers.15.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
48
+ "model.layers.15.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
49
+ "model.layers.16.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
50
+ "model.layers.16.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
51
+ "model.layers.16.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
52
+ "model.layers.16.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
53
+ "model.layers.16.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
54
+ "model.layers.17.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
55
+ "model.layers.17.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
56
+ "model.layers.17.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
57
+ "model.layers.17.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
58
+ "model.layers.17.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
59
+ "model.layers.18.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
60
+ "model.layers.18.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
61
+ "model.layers.18.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
62
+ "model.layers.18.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
63
+ "model.layers.18.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
64
+ "model.layers.19.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
65
+ "model.layers.19.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
66
+ "model.layers.19.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
67
+ "model.layers.19.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
68
+ "model.layers.19.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
69
+ "model.layers.2.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
70
+ "model.layers.2.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
71
+ "model.layers.2.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
72
+ "model.layers.2.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
73
+ "model.layers.2.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
74
+ "model.layers.20.channel_mixer.l1.weight": "pytorch_model-00002-of-00003.bin",
75
+ "model.layers.20.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
76
+ "model.layers.20.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
77
+ "model.layers.20.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
78
+ "model.layers.20.token_mixer.qkvu_proj.weight": "pytorch_model-00002-of-00003.bin",
79
+ "model.layers.21.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
80
+ "model.layers.21.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
81
+ "model.layers.21.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
82
+ "model.layers.21.token_mixer.out_proj.weight": "pytorch_model-00002-of-00003.bin",
83
+ "model.layers.21.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
84
+ "model.layers.22.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
85
+ "model.layers.22.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
86
+ "model.layers.22.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
87
+ "model.layers.22.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
88
+ "model.layers.22.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
89
+ "model.layers.23.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
90
+ "model.layers.23.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
91
+ "model.layers.23.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
92
+ "model.layers.23.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
93
+ "model.layers.23.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
94
+ "model.layers.24.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
95
+ "model.layers.24.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
96
+ "model.layers.24.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
97
+ "model.layers.24.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
98
+ "model.layers.24.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
99
+ "model.layers.25.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
100
+ "model.layers.25.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
101
+ "model.layers.25.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
102
+ "model.layers.25.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
103
+ "model.layers.25.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
104
+ "model.layers.26.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
105
+ "model.layers.26.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
106
+ "model.layers.26.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
107
+ "model.layers.26.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
108
+ "model.layers.26.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
109
+ "model.layers.27.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
110
+ "model.layers.27.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
111
+ "model.layers.27.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
112
+ "model.layers.27.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
113
+ "model.layers.27.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
114
+ "model.layers.28.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
115
+ "model.layers.28.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
116
+ "model.layers.28.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
117
+ "model.layers.28.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
118
+ "model.layers.28.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
119
+ "model.layers.29.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
120
+ "model.layers.29.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
121
+ "model.layers.29.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
122
+ "model.layers.29.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
123
+ "model.layers.29.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
124
+ "model.layers.3.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
125
+ "model.layers.3.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
126
+ "model.layers.3.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
127
+ "model.layers.3.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
128
+ "model.layers.3.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
129
+ "model.layers.30.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
130
+ "model.layers.30.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
131
+ "model.layers.30.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
132
+ "model.layers.30.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
133
+ "model.layers.30.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
134
+ "model.layers.31.channel_mixer.l1.weight": "pytorch_model-00003-of-00003.bin",
135
+ "model.layers.31.channel_mixer.l2.weight": "pytorch_model-00003-of-00003.bin",
136
+ "model.layers.31.channel_mixer.l3.weight": "pytorch_model-00003-of-00003.bin",
137
+ "model.layers.31.token_mixer.out_proj.weight": "pytorch_model-00003-of-00003.bin",
138
+ "model.layers.31.token_mixer.qkvu_proj.weight": "pytorch_model-00003-of-00003.bin",
139
+ "model.layers.4.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
140
+ "model.layers.4.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
141
+ "model.layers.4.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
142
+ "model.layers.4.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
143
+ "model.layers.4.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
144
+ "model.layers.5.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
145
+ "model.layers.5.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
146
+ "model.layers.5.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
147
+ "model.layers.5.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
148
+ "model.layers.5.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
149
+ "model.layers.6.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
150
+ "model.layers.6.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
151
+ "model.layers.6.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
152
+ "model.layers.6.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
153
+ "model.layers.6.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
154
+ "model.layers.7.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
155
+ "model.layers.7.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
156
+ "model.layers.7.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
157
+ "model.layers.7.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
158
+ "model.layers.7.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
159
+ "model.layers.8.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
160
+ "model.layers.8.channel_mixer.l2.weight": "pytorch_model-00001-of-00003.bin",
161
+ "model.layers.8.channel_mixer.l3.weight": "pytorch_model-00001-of-00003.bin",
162
+ "model.layers.8.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
163
+ "model.layers.8.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin",
164
+ "model.layers.9.channel_mixer.l1.weight": "pytorch_model-00001-of-00003.bin",
165
+ "model.layers.9.channel_mixer.l2.weight": "pytorch_model-00002-of-00003.bin",
166
+ "model.layers.9.channel_mixer.l3.weight": "pytorch_model-00002-of-00003.bin",
167
+ "model.layers.9.token_mixer.out_proj.weight": "pytorch_model-00001-of-00003.bin",
168
+ "model.layers.9.token_mixer.qkvu_proj.weight": "pytorch_model-00001-of-00003.bin"
169
+ }
170
+ }
special_tokens_map.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": true,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "</s>",
11
+ "lstrip": false,
12
+ "normalized": true,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "unk_token": {
17
+ "content": "<unk>",
18
+ "lstrip": false,
19
+ "normalized": true,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ }
23
+ }
srmsnorm_triton.py ADDED
@@ -0,0 +1,202 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # CREDITS: This comes almost as-is from the Triton layer norm tutorial
2
+ # https://github.com/openai/triton/blob/master/python/tutorials/05-layer-norm.py
3
+ # Copyright 2024 OpenNLPLab
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+
17
+ # coding=utf-8
18
+ import torch
19
+ import torch.nn.functional as F
20
+ import triton
21
+ import triton.language as tl
22
+
23
+
24
+ # fmt: off
25
+ @triton.jit
26
+ def srms_norm_fw(X, Y, V, stride, N, eps, BLOCK_SIZE_N: tl.constexpr):
27
+ # fmt: on
28
+ row = tl.program_id(0)
29
+ cols = tl.arange(0, BLOCK_SIZE_N)
30
+ mask = cols < N
31
+
32
+ # Move to this row
33
+ x_ptrs = X + row * stride + cols
34
+ x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
35
+
36
+ x_zm = tl.where(mask, x, 0.0)
37
+
38
+ x_var = tl.sum(x_zm * x_zm, axis=0) / N
39
+ rstd = 1.0 / tl.sqrt(x_var + eps)
40
+
41
+ # Normalize, optionally affine
42
+ y = x_zm * rstd
43
+ tl.store(V + row, rstd)
44
+
45
+ y_ptrs = Y + row * stride + cols
46
+ tl.store(y_ptrs, y, mask=mask)
47
+
48
+
49
+ # Backward pass (DX + partial DW + partial DB)
50
+ # fmt: off
51
+ @triton.jit
52
+ def srms_norm_bwd_dx_fused(
53
+ DX, DY,
54
+ X, V,
55
+ stride, N,
56
+ # META-parameters
57
+ BLOCK_SIZE_N: tl.constexpr,
58
+ ):
59
+ # fmt: on
60
+
61
+ # position of elements processed by this program
62
+ row = tl.program_id(0)
63
+ cols = tl.arange(0, BLOCK_SIZE_N)
64
+ mask = cols < N
65
+
66
+ # offset data pointers to start at the row of interest
67
+ x_ptrs = X + row * stride + cols
68
+ dy_ptrs = DY + row * stride + cols
69
+
70
+ # load data to SRAM
71
+ x = tl.load(x_ptrs, mask=mask, other=0)
72
+ dy = tl.load(dy_ptrs, mask=mask, other=0)
73
+ rstd = tl.load(V + row)
74
+
75
+ # compute dx
76
+ xhat = x * rstd
77
+ wdy = dy
78
+
79
+ xhat = tl.where(mask, xhat, 0.)
80
+ wdy = tl.where(mask, wdy, 0.)
81
+ mean1 = tl.sum(xhat * wdy, axis=0) / N
82
+ dx = (wdy - (xhat * mean1)) * rstd
83
+
84
+ # write-back dx
85
+ mask = cols < N # re-materialize the mask to save registers
86
+ dx_ptrs = DX + row * stride + cols
87
+ tl.store(dx_ptrs, dx, mask=mask)
88
+
89
+
90
+ class _SrmsNorm(torch.autograd.Function):
91
+
92
+ @staticmethod
93
+ def forward(ctx, x, eps):
94
+ # catch eps being too small if the tensors are fp16
95
+ if x.dtype == torch.float16:
96
+ eps = max(eps, 1.6e-5)
97
+
98
+ # allocate output
99
+ y = torch.empty_like(x)
100
+
101
+ # reshape input data into 2D tensor
102
+ x_arg = x.reshape(-1, x.shape[-1])
103
+ M, N = x_arg.shape
104
+
105
+ # allocate mean and std, they'll be used in the backward pass
106
+ rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
107
+
108
+ # Less than 64KB per feature: enqueue fused kernel
109
+ MAX_FUSED_SIZE = 65536 // x.element_size()
110
+ BLOCK_SIZE_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
111
+ if N > BLOCK_SIZE_N:
112
+ raise RuntimeError(
113
+ "This layer norm doesn't support feature dim >= 64KB.")
114
+
115
+ if not x_arg.is_contiguous() or not y.is_contiguous():
116
+ x_arg = x_arg.contiguous()
117
+ y = y.contiguous()
118
+
119
+ # heuristics for number of warps.
120
+ num_warps = min(max(BLOCK_SIZE_N // 256, 1), 16)
121
+
122
+ # enqueue kernel
123
+ # fmt: off
124
+ srms_norm_fw[(M,)](
125
+ x_arg, y, rstd,
126
+ x_arg.stride(0),
127
+ N,
128
+ eps,
129
+ num_warps=num_warps,
130
+ BLOCK_SIZE_N=BLOCK_SIZE_N,
131
+ )
132
+ # fmt: on
133
+
134
+ ctx.save_for_backward(x, rstd)
135
+ ctx.BLOCK_SIZE_N = BLOCK_SIZE_N
136
+ ctx.num_warps = num_warps
137
+
138
+ return y.reshape_as(x)
139
+
140
+ @staticmethod
141
+ def backward(
142
+ ctx, dy
143
+ ): # pragma: no cover # this is covered, but called directly from C++
144
+ x, rstd = ctx.saved_tensors
145
+
146
+ # flatten the batch dimension, if any.
147
+ # We're interested in 'samples' x norm_dimension
148
+ x = x.reshape(-1, x.size(-1))
149
+ M, N = x.size()
150
+
151
+ # heuristics for amount of parallel reduction stream for DG/DB
152
+ GROUP_SIZE_M = 32
153
+ if N <= 8192:
154
+ GROUP_SIZE_M = 64
155
+ if N <= 4096:
156
+ GROUP_SIZE_M = 96
157
+ if N <= 2048:
158
+ GROUP_SIZE_M = 128
159
+ if N <= 1024:
160
+ GROUP_SIZE_M = 256
161
+
162
+ if dy.dtype == torch.float32:
163
+ GROUP_SIZE_M = GROUP_SIZE_M // 2
164
+
165
+ # allocate output
166
+ dy = dy.contiguous()
167
+ dx = torch.empty_like(dy)
168
+
169
+ # Check the tensor shapes and layouts
170
+ # we suppose in the kernel that they have the same size and are contiguous
171
+ assert (
172
+ dy.numel() == x.numel()
173
+ ), "Something is wrong in the backward graph, possibly because of an inplace operation after the layernorm"
174
+
175
+ # enqueue kernel using forward pass heuristics
176
+ # also compute partial sums for DW and DB
177
+ num_warps = min(max(ctx.BLOCK_SIZE_N // 256, 1), 16)
178
+
179
+ # fmt: off
180
+ srms_norm_bwd_dx_fused[(M,)](
181
+ dx, dy, x,
182
+ rstd,
183
+ x.stride(0),
184
+ N,
185
+ BLOCK_SIZE_N=ctx.BLOCK_SIZE_N,
186
+ num_warps=num_warps
187
+ )
188
+ # fmt: on
189
+
190
+ dx = dx.reshape_as(dy)
191
+ return dx, None, None
192
+
193
+
194
+ class SimpleRMSNorm(torch.nn.Module):
195
+
196
+ def __init__(self, dim: int, eps: float = 1e-6):
197
+ super().__init__()
198
+ self.eps = eps
199
+ self.dim = dim
200
+
201
+ def forward(self, x):
202
+ return _SrmsNorm.apply(x, self.eps)
tokenization_baichuan.py ADDED
@@ -0,0 +1,248 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+
21
+ import os
22
+ from shutil import copyfile
23
+ from typing import Any, Dict, List, Optional, Tuple
24
+
25
+ import sentencepiece as spm
26
+ from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
27
+ from transformers.utils import logging
28
+
29
+ logger = logging.get_logger(__name__)
30
+
31
+ VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
32
+
33
+ PRETRAINED_VOCAB_FILES_MAP = {
34
+ "vocab_file": {},
35
+ "tokenizer_file": {},
36
+ }
37
+ PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {}
38
+
39
+
40
+ class BaiChuanTokenizer(PreTrainedTokenizer):
41
+ """
42
+ Construct a BaiChuan tokenizer. Based on byte-level Byte-Pair-Encoding.
43
+
44
+ Args:
45
+ vocab_file (`str`):
46
+ Path to the vocabulary file.
47
+ """
48
+
49
+ vocab_files_names = VOCAB_FILES_NAMES
50
+ pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
51
+ max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
52
+ model_input_names = ["input_ids", "attention_mask"]
53
+
54
+ def __init__(
55
+ self,
56
+ vocab_file,
57
+ unk_token="<unk>",
58
+ bos_token="<s>",
59
+ eos_token="</s>",
60
+ pad_token=None,
61
+ sp_model_kwargs: Optional[Dict[str, Any]] = None,
62
+ add_bos_token=True,
63
+ add_eos_token=False,
64
+ clean_up_tokenization_spaces=False,
65
+ **kwargs,
66
+ ):
67
+ self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
68
+ bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
69
+ eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
70
+ unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
71
+ pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
72
+ self.vocab_file = vocab_file
73
+ self.add_bos_token = add_bos_token
74
+ self.add_eos_token = add_eos_token
75
+ self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
76
+ self.sp_model.Load(vocab_file)
77
+ super().__init__(
78
+ bos_token=bos_token,
79
+ eos_token=eos_token,
80
+ unk_token=unk_token,
81
+ pad_token=pad_token,
82
+ add_bos_token=add_bos_token,
83
+ add_eos_token=add_eos_token,
84
+ sp_model_kwargs=self.sp_model_kwargs,
85
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
86
+ **kwargs,
87
+ )
88
+
89
+ def __getstate__(self):
90
+ state = self.__dict__.copy()
91
+ state["sp_model"] = None
92
+ return state
93
+
94
+ def __setstate__(self, d):
95
+ self.__dict__ = d
96
+ self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
97
+ self.sp_model.Load(self.vocab_file)
98
+
99
+ @property
100
+ def vocab_size(self):
101
+ """Returns vocab size"""
102
+ return self.sp_model.get_piece_size()
103
+
104
+ def get_vocab(self):
105
+ """Returns vocab as a dict"""
106
+ vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
107
+ vocab.update(self.added_tokens_encoder)
108
+ return vocab
109
+
110
+ def _tokenize(self, text):
111
+ """Returns a tokenized string."""
112
+ return self.sp_model.encode(text, out_type=str)
113
+
114
+ def _convert_token_to_id(self, token):
115
+ """Converts a token (str) in an id using the vocab."""
116
+ return self.sp_model.piece_to_id(token)
117
+
118
+ def _convert_id_to_token(self, index):
119
+ """Converts an index (integer) in a token (str) using the vocab."""
120
+ token = self.sp_model.IdToPiece(index)
121
+ return token
122
+
123
+ def convert_tokens_to_string(self, tokens):
124
+ """Converts a sequence of tokens (string) in a single string."""
125
+ current_sub_tokens = []
126
+ out_string = ""
127
+ prev_is_special = False
128
+ for i, token in enumerate(tokens):
129
+ # make sure that special tokens are not decoded using sentencepiece model
130
+ if token in self.all_special_tokens:
131
+ if not prev_is_special and i != 0:
132
+ out_string += " "
133
+ out_string += self.sp_model.decode(current_sub_tokens) + token
134
+ prev_is_special = True
135
+ current_sub_tokens = []
136
+ else:
137
+ current_sub_tokens.append(token)
138
+ prev_is_special = False
139
+ out_string += self.sp_model.decode(current_sub_tokens)
140
+ return out_string
141
+
142
+ def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
143
+ """
144
+ Save the vocabulary and special tokens file to a directory.
145
+
146
+ Args:
147
+ save_directory (`str`):
148
+ The directory in which to save the vocabulary.
149
+
150
+ Returns:
151
+ `Tuple(str)`: Paths to the files saved.
152
+ """
153
+ if not os.path.isdir(save_directory):
154
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
155
+ return
156
+ out_vocab_file = os.path.join(
157
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
158
+ )
159
+
160
+ if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
161
+ copyfile(self.vocab_file, out_vocab_file)
162
+ elif not os.path.isfile(self.vocab_file):
163
+ with open(out_vocab_file, "wb") as fi:
164
+ content_spiece_model = self.sp_model.serialized_model_proto()
165
+ fi.write(content_spiece_model)
166
+
167
+ return (out_vocab_file,)
168
+
169
+ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
170
+ bos_token_id = [self.bos_token_id] if self.add_bos_token else []
171
+ eos_token_id = [self.eos_token_id] if self.add_eos_token else []
172
+
173
+ output = bos_token_id + token_ids_0 + eos_token_id
174
+
175
+ if token_ids_1 is not None:
176
+ output = output + bos_token_id + token_ids_1 + eos_token_id
177
+
178
+ return output
179
+
180
+ def get_special_tokens_mask(
181
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
182
+ ) -> List[int]:
183
+ """
184
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
185
+ special tokens using the tokenizer `prepare_for_model` method.
186
+
187
+ Args:
188
+ token_ids_0 (`List[int]`):
189
+ List of IDs.
190
+ token_ids_1 (`List[int]`, *optional*):
191
+ Optional second list of IDs for sequence pairs.
192
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
193
+ Whether or not the token list is already formatted with special tokens for the model.
194
+
195
+ Returns:
196
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
197
+ """
198
+ if already_has_special_tokens:
199
+ return super().get_special_tokens_mask(
200
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
201
+ )
202
+
203
+ bos_token_id = [1] if self.add_bos_token else []
204
+ eos_token_id = [1] if self.add_eos_token else []
205
+
206
+ if token_ids_1 is None:
207
+ return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
208
+ return (
209
+ bos_token_id
210
+ + ([0] * len(token_ids_0))
211
+ + eos_token_id
212
+ + bos_token_id
213
+ + ([0] * len(token_ids_1))
214
+ + eos_token_id
215
+ )
216
+
217
+ def create_token_type_ids_from_sequences(
218
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
219
+ ) -> List[int]:
220
+ """
221
+ Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
222
+ sequence pair mask has the following format:
223
+
224
+ ```
225
+ 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
226
+ | first sequence | second sequence |
227
+ ```
228
+
229
+ if token_ids_1 is None, only returns the first portion of the mask (0s).
230
+
231
+ Args:
232
+ token_ids_0 (`List[int]`):
233
+ List of ids.
234
+ token_ids_1 (`List[int]`, *optional*):
235
+ Optional second list of IDs for sequence pairs.
236
+
237
+ Returns:
238
+ `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
239
+ """
240
+ bos_token_id = [self.bos_token_id] if self.add_bos_token else []
241
+ eos_token_id = [self.eos_token_id] if self.add_eos_token else []
242
+
243
+ output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
244
+
245
+ if token_ids_1 is not None:
246
+ output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
247
+
248
+ return output
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4be54af290d93c113bcbf421115ae9eed9d6340408f564898f1e966dc738ef01
3
+ size 1136699
tokenizer_config.json ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoTokenizer": [
4
+ "tokenization_baichuan.BaiChuanTokenizer",
5
+ null
6
+ ]
7
+ },
8
+ "add_bos_token": false,
9
+ "add_eos_token": false,
10
+ "bos_token": {
11
+ "__type": "AddedToken",
12
+ "content": "<s>",
13
+ "lstrip": false,
14
+ "normalized": true,
15
+ "rstrip": false,
16
+ "single_word": false
17
+ },
18
+ "clean_up_tokenization_spaces": false,
19
+ "eos_token": {
20
+ "__type": "AddedToken",
21
+ "content": "</s>",
22
+ "lstrip": false,
23
+ "normalized": true,
24
+ "rstrip": false,
25
+ "single_word": false
26
+ },
27
+ "model_max_length": 1000000000000000019884624838656,
28
+ "sp_model_kwargs": {},
29
+ "tokenizer_class": "BaiChuanTokenizer",
30
+ "unk_token": {
31
+ "__type": "AddedToken",
32
+ "content": "<unk>",
33
+ "lstrip": false,
34
+ "normalized": true,
35
+ "rstrip": false,
36
+ "single_word": false
37
+ }
38
+ }
utils.py ADDED
@@ -0,0 +1,166 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 OpenNLPLab
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # coding=utf-8
16
+ import logging
17
+ import os
18
+ import sys
19
+
20
+ import torch
21
+ from torch import nn
22
+ import torch.distributed as dist
23
+ import torch.nn.functional as F
24
+
25
+ from .norm import SimpleRMSNorm as SimpleRMSNormTorch
26
+ from .srmsnorm_triton import SimpleRMSNorm as SimpleRMSNormTriton
27
+
28
+ use_triton = eval(os.environ.get("use_triton", default="True"))
29
+ debug = eval(os.environ.get("debug", default="False"))
30
+
31
+ if use_triton:
32
+ SimpleRMSNorm = SimpleRMSNormTriton
33
+ else:
34
+ SimpleRMSNorm = SimpleRMSNormTorch
35
+
36
+ logging.basicConfig(
37
+ format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
38
+ datefmt="%Y-%m-%d %H:%M:%S",
39
+ level=os.environ.get("LOGLEVEL", "INFO").upper(),
40
+ stream=sys.stdout,
41
+ )
42
+ logger = logging.getLogger("print_config")
43
+
44
+ BASE_DIM = 256
45
+
46
+
47
+ def is_dist_avail_and_initialized():
48
+ if not dist.is_available():
49
+ return False
50
+ if not dist.is_initialized():
51
+ return False
52
+ return True
53
+
54
+
55
+ def get_world_size():
56
+ if not is_dist_avail_and_initialized():
57
+ return 1
58
+ return dist.get_world_size()
59
+
60
+
61
+ def get_rank():
62
+ if not is_dist_avail_and_initialized():
63
+ return 0
64
+ return dist.get_rank()
65
+
66
+
67
+ def is_main_process():
68
+ return get_rank() == 0
69
+
70
+
71
+ def logging_info(string):
72
+ if is_main_process():
73
+ logger.info(string)
74
+
75
+
76
+ def print_params(**kwargs):
77
+ if is_main_process():
78
+ logger.info(f"start print config of {kwargs['__class__']}")
79
+ for key in kwargs:
80
+ if key in ["__class__", "self"]:
81
+ continue
82
+ logger.info(f"{key}: {kwargs[key]}")
83
+ logger.info(f"end print config of {kwargs['__class__']}")
84
+
85
+
86
+ def print_config(config):
87
+ if is_main_process():
88
+ logger.info(f"start print config of {config['__class__']}")
89
+ for key in config:
90
+ if key in ["__class__", "self"]:
91
+ continue
92
+ logger.info(f"{key}: {config[key]}")
93
+ logger.info(f"end print config of {config['__class__']}")
94
+
95
+
96
+ def print_module(module):
97
+ named_modules = set()
98
+ for p in module.named_modules():
99
+ named_modules.update([p[0]])
100
+ named_modules = list(named_modules)
101
+
102
+ string_repr = ""
103
+ for p in module.named_parameters():
104
+ name = p[0].split(".")[0]
105
+ if name not in named_modules:
106
+ string_repr = (string_repr + "(" + name + "): " + "Tensor(" +
107
+ str(tuple(p[1].shape)) + ", requires_grad=" +
108
+ str(p[1].requires_grad) + ")\n")
109
+
110
+ return string_repr.rstrip("\n")
111
+
112
+
113
+ def get_activation_fn(activation):
114
+ if debug:
115
+ logger.info(f"activation: {activation}")
116
+ if activation == "gelu":
117
+ return F.gelu
118
+ elif activation == "relu":
119
+ return F.relu
120
+ elif activation == "elu":
121
+ return F.elu
122
+ elif activation == "sigmoid":
123
+ return F.sigmoid
124
+ elif activation == "exp":
125
+
126
+ def f(x):
127
+ with torch.no_grad():
128
+ x_max = torch.max(x, dim=-1, keepdims=True).values
129
+ y = torch.exp(x - x_max)
130
+
131
+ return y
132
+
133
+ return f
134
+ elif activation == "leak":
135
+ return F.leaky_relu
136
+ elif activation == "1+elu":
137
+
138
+ def f(x):
139
+ return 1 + F.elu(x)
140
+
141
+ return f
142
+ elif activation == "2+elu":
143
+
144
+ def f(x):
145
+ return 2 + F.elu(x)
146
+
147
+ return f
148
+ elif activation == "silu" or activation == "swish":
149
+ return F.silu
150
+ elif activation == "sine":
151
+ return torch.sin
152
+ else:
153
+ logger.info(
154
+ f"activation: does not support {activation}, use Identity!!!")
155
+ return lambda x: x
156
+
157
+
158
+ def get_norm_fn(norm_type):
159
+ if norm_type == "simplermsnorm":
160
+ return SimpleRMSNorm
161
+ else:
162
+ return nn.LayerNorm
163
+
164
+
165
+ def convert_to_multiple_of_base(x):
166
+ return BASE_DIM * ((x + BASE_DIM - 1) // BASE_DIM)