Upload Llama3.py
Browse files
Llama3.py
ADDED
@@ -0,0 +1,287 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
2 |
+
# This software may be used and distributed in accordance with the terms of the Llama 3 Community License Agreement.
|
3 |
+
import tensorflow as tf
|
4 |
+
from tensorflow.keras.layers import Embedding,Dense
|
5 |
+
from tensorflow.keras import Model
|
6 |
+
|
7 |
+
import math
|
8 |
+
from dataclasses import dataclass
|
9 |
+
from typing import Optional
|
10 |
+
|
11 |
+
|
12 |
+
@dataclass
|
13 |
+
class ModelArgs:
|
14 |
+
dim: int = 4096
|
15 |
+
n_layers: int = 32
|
16 |
+
n_heads: int = 32
|
17 |
+
n_kv_heads: Optional[int] = None
|
18 |
+
vocab_size: int = -1
|
19 |
+
multiple_of: int = 256 # make SwiGLU hidden layer size multiple of large power of 2
|
20 |
+
ffn_dim_multiplier: Optional[float] = None
|
21 |
+
norm_eps: float = 1e-5
|
22 |
+
rope_theta: float = 500000
|
23 |
+
|
24 |
+
max_batch_size: int = 32
|
25 |
+
max_seq_len: int = 2048
|
26 |
+
|
27 |
+
|
28 |
+
class RMSNorm:
|
29 |
+
def __init__(self, dim: int, eps: float = 1e-6):
|
30 |
+
self.eps = eps
|
31 |
+
self.weight = tf.Variable(tf.ones((dim)))
|
32 |
+
|
33 |
+
def _norm(self, x):
|
34 |
+
return x * tf.math.rsqrt(tf.reduce_mean(tf.pow(x, 2), -1, keepdims=True) + self.eps)
|
35 |
+
|
36 |
+
def __call__(self, x):
|
37 |
+
output = tf.cast(self._norm(tf.cast(x, 'float32')), x.dtype)
|
38 |
+
return output * self.weight
|
39 |
+
|
40 |
+
|
41 |
+
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
|
42 |
+
freqs = 1.0 / (theta ** (tf.cast(tf.range(0, dim, 2)[: (dim // 2)], 'float32') / dim))
|
43 |
+
t = tf.range(end, dtype='float32')
|
44 |
+
freqs = tf.experimental.numpy.outer(t, freqs)
|
45 |
+
freqs_cis = tf.complex(tf.ones_like(freqs), freqs)
|
46 |
+
real_part = tf.math.cos(freqs)
|
47 |
+
imag_part = tf.math.sin(freqs)
|
48 |
+
freqs_cis = tf.complex(real_part, imag_part) # complex64
|
49 |
+
return freqs_cis
|
50 |
+
|
51 |
+
|
52 |
+
def reshape_for_broadcast(freqs_cis, x):
|
53 |
+
ndim = x.ndim
|
54 |
+
assert 0 <= 1 < ndim
|
55 |
+
assert freqs_cis.shape == (x.shape[1], x.shape[-1])
|
56 |
+
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
|
57 |
+
return tf.reshape(freqs_cis, shape)
|
58 |
+
|
59 |
+
|
60 |
+
def apply_rotary_emb(
|
61 |
+
xq,
|
62 |
+
xk,
|
63 |
+
freqs_cis,
|
64 |
+
):
|
65 |
+
xq = tf.reshape(tf.cast(xq, 'float32'), (xq.shape[:-1] + (xq.shape[-1] // 2, 2)))
|
66 |
+
real_part = xq[..., 0]
|
67 |
+
imag_part = xq[..., 1]
|
68 |
+
xq_ = tf.complex(real_part, imag_part)
|
69 |
+
xk = tf.reshape(tf.cast(xk, 'float32'), (xk.shape[:-1] + (xk.shape[-1] // 2, 2)))
|
70 |
+
real_part = xk[..., 0]
|
71 |
+
imag_part = xk[..., 1]
|
72 |
+
xk_ = tf.complex(real_part, imag_part)
|
73 |
+
freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
|
74 |
+
xq_freqs_cis = xq_ * freqs_cis
|
75 |
+
xq_out = tf.stack([tf.math.real(xq_freqs_cis), tf.math.imag(xq_freqs_cis)], axis=-1)
|
76 |
+
shape = xq_out.shape
|
77 |
+
xq_out = tf.reshape(xq_out, [-1, shape[1], shape[2], shape[3] * shape[4]])
|
78 |
+
xk_freqs_cis = xk_ * freqs_cis
|
79 |
+
xk_out = tf.stack([tf.math.real(xk_freqs_cis), tf.math.imag(xk_freqs_cis)], axis=-1)
|
80 |
+
shape = xk_out.shape
|
81 |
+
xk_out = tf.reshape(xk_out, [-1, shape[1], shape[2], shape[3] * shape[4]])
|
82 |
+
return tf.cast(xq_out, xq.dtype), tf.cast(xk_out, xk.dtype)
|
83 |
+
|
84 |
+
|
85 |
+
def repeat_kv(x, n_rep: int):
|
86 |
+
bs, slen, n_kv_heads, head_dim = x.shape
|
87 |
+
if n_rep == 1:
|
88 |
+
return x
|
89 |
+
return tf.reshape(tf.tile(x[:, :, :, None, :], [1, 1, 1, n_rep, 1]), (bs, slen, n_kv_heads * n_rep, head_dim))
|
90 |
+
|
91 |
+
|
92 |
+
class Attention:
|
93 |
+
def __init__(self, args: ModelArgs):
|
94 |
+
self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
|
95 |
+
model_parallel_size = 1
|
96 |
+
self.n_local_heads = args.n_heads // model_parallel_size
|
97 |
+
self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
|
98 |
+
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
99 |
+
self.head_dim = args.dim // args.n_heads
|
100 |
+
|
101 |
+
self.wq = Dense(
|
102 |
+
args.n_heads * self.head_dim,
|
103 |
+
use_bias=False,
|
104 |
+
)
|
105 |
+
self.wk = Dense(
|
106 |
+
self.n_kv_heads * self.head_dim,
|
107 |
+
use_bias=False,
|
108 |
+
)
|
109 |
+
self.wv = Dense(
|
110 |
+
self.n_kv_heads * self.head_dim,
|
111 |
+
use_bias=False,
|
112 |
+
)
|
113 |
+
self.wo = Dense(
|
114 |
+
args.dim,
|
115 |
+
use_bias=False,
|
116 |
+
)
|
117 |
+
|
118 |
+
self.cache_k = tf.Variable(tf.zeros(
|
119 |
+
(
|
120 |
+
args.max_batch_size,
|
121 |
+
args.max_seq_len,
|
122 |
+
self.n_local_kv_heads,
|
123 |
+
self.head_dim,
|
124 |
+
)
|
125 |
+
), trainable=False)
|
126 |
+
self.cache_v = tf.Variable(tf.zeros(
|
127 |
+
(
|
128 |
+
args.max_batch_size,
|
129 |
+
args.max_seq_len,
|
130 |
+
self.n_local_kv_heads,
|
131 |
+
self.head_dim,
|
132 |
+
)
|
133 |
+
), trainable=False)
|
134 |
+
|
135 |
+
def __call__(
|
136 |
+
self,
|
137 |
+
x,
|
138 |
+
start_pos: int,
|
139 |
+
freqs_cis,
|
140 |
+
mask,
|
141 |
+
):
|
142 |
+
bsz, seqlen, _ = x.shape
|
143 |
+
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
|
144 |
+
|
145 |
+
xq = tf.reshape(xq, (bsz, seqlen, self.n_local_heads, self.head_dim))
|
146 |
+
xk = tf.reshape(xk, (bsz, seqlen, self.n_local_kv_heads, self.head_dim))
|
147 |
+
xv = tf.reshape(xv, (bsz, seqlen, self.n_local_kv_heads, self.head_dim))
|
148 |
+
|
149 |
+
xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
|
150 |
+
|
151 |
+
self.cache_k = tf.cast(self.cache_k, xq.dtype)
|
152 |
+
self.cache_v = tf.cast(self.cache_v, xq.dtype)
|
153 |
+
|
154 |
+
self.cache_k[:bsz, start_pos : start_pos + seqlen].assign(xk)
|
155 |
+
self.cache_v[:bsz, start_pos : start_pos + seqlen].assign(xv)
|
156 |
+
|
157 |
+
keys = self.cache_k[:bsz, : start_pos + seqlen]
|
158 |
+
values = self.cache_v[:bsz, : start_pos + seqlen]
|
159 |
+
|
160 |
+
# repeat k/v heads if n_kv_heads < n_heads
|
161 |
+
keys = repeat_kv(
|
162 |
+
keys, self.n_rep
|
163 |
+
) # (bs, cache_len + seqlen, n_local_heads, head_dim)
|
164 |
+
values = repeat_kv(
|
165 |
+
values, self.n_rep
|
166 |
+
) # (bs, cache_len + seqlen, n_local_heads, head_dim)
|
167 |
+
|
168 |
+
xq = tf.transpose(xq, (0, 2, 1, 3)) # (bs, n_local_heads, seqlen, head_dim)
|
169 |
+
keys = tf.transpose(keys, (0, 2, 1, 3)) # (bs, n_local_heads, cache_len + seqlen, head_dim)
|
170 |
+
values = tf.transpose(values,
|
171 |
+
(0, 2, 1, 3)
|
172 |
+
) # (bs, n_local_heads, cache_len + seqlen, head_dim)
|
173 |
+
scores = tf.matmul(xq, tf.transpose(keys, (0, 1, 3, 2))) / math.sqrt(self.head_dim)
|
174 |
+
if mask is not None:
|
175 |
+
scores = scores + mask # (bs, n_local_heads, seqlen, cache_len + seqlen)
|
176 |
+
scores = tf.cast(tf.nn.softmax(tf.cast(scores, 'float32')), xq.dtype)
|
177 |
+
output = tf.matmul(scores, values) # (bs, n_local_heads, seqlen, head_dim)
|
178 |
+
output = tf.reshape(tf.transpose(output, (0, 2, 1, 3)), (bsz, seqlen, -1))
|
179 |
+
return self.wo(output)
|
180 |
+
|
181 |
+
|
182 |
+
class FeedForward:
|
183 |
+
def __init__(
|
184 |
+
self,
|
185 |
+
dim: int,
|
186 |
+
hidden_dim: int,
|
187 |
+
multiple_of: int,
|
188 |
+
ffn_dim_multiplier: Optional[float],
|
189 |
+
):
|
190 |
+
hidden_dim = int(2 * hidden_dim / 3)
|
191 |
+
# custom dim factor multiplier
|
192 |
+
if ffn_dim_multiplier is not None:
|
193 |
+
hidden_dim = int(ffn_dim_multiplier * hidden_dim)
|
194 |
+
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
|
195 |
+
|
196 |
+
self.w1 = Dense(
|
197 |
+
hidden_dim, use_bias=False
|
198 |
+
)
|
199 |
+
self.w2 = Dense(
|
200 |
+
dim, use_bias=False
|
201 |
+
)
|
202 |
+
self.w3 = Dense(
|
203 |
+
hidden_dim, use_bias=False
|
204 |
+
)
|
205 |
+
|
206 |
+
def __call__(self, x):
|
207 |
+
return self.w2(tf.nn.silu(self.w1(x)) * self.w3(x))
|
208 |
+
|
209 |
+
|
210 |
+
class TransformerBlock:
|
211 |
+
def __init__(self, layer_id: int, args: ModelArgs):
|
212 |
+
self.n_heads = args.n_heads
|
213 |
+
self.dim = args.dim
|
214 |
+
self.head_dim = args.dim // args.n_heads
|
215 |
+
self.attention = Attention(args)
|
216 |
+
self.feed_forward = FeedForward(
|
217 |
+
dim=args.dim,
|
218 |
+
hidden_dim=4 * args.dim,
|
219 |
+
multiple_of=args.multiple_of,
|
220 |
+
ffn_dim_multiplier=args.ffn_dim_multiplier,
|
221 |
+
)
|
222 |
+
self.layer_id = layer_id
|
223 |
+
self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
224 |
+
self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
225 |
+
|
226 |
+
def __call__(
|
227 |
+
self,
|
228 |
+
x,
|
229 |
+
start_pos,
|
230 |
+
freqs_cis,
|
231 |
+
mask,
|
232 |
+
):
|
233 |
+
h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
|
234 |
+
out = h + self.feed_forward(self.ffn_norm(h))
|
235 |
+
return out
|
236 |
+
|
237 |
+
|
238 |
+
class Llama3(Model):
|
239 |
+
def __init__(self, params: ModelArgs):
|
240 |
+
self.params = params
|
241 |
+
self.vocab_size = params.vocab_size
|
242 |
+
self.n_layers = params.n_layers
|
243 |
+
|
244 |
+
self.tok_embeddings = Embedding(
|
245 |
+
params.vocab_size, params.dim
|
246 |
+
)
|
247 |
+
|
248 |
+
self.layers_ = []
|
249 |
+
for layer_id in range(params.n_layers):
|
250 |
+
self.layers_.append(TransformerBlock(layer_id, params))
|
251 |
+
|
252 |
+
self.norm = RMSNorm(params.dim, eps=params.norm_eps)
|
253 |
+
self.output_ = Dense(
|
254 |
+
params.vocab_size, use_bias=False
|
255 |
+
)
|
256 |
+
|
257 |
+
self.freqs_cis = precompute_freqs_cis(
|
258 |
+
params.dim // params.n_heads,
|
259 |
+
params.max_seq_len * 2,
|
260 |
+
params.rope_theta,
|
261 |
+
)
|
262 |
+
|
263 |
+
def __call__(self, tokens, start_pos: int):
|
264 |
+
_bsz, seqlen = tokens.shape
|
265 |
+
h = self.tok_embeddings(tokens)
|
266 |
+
self.freqs_cis = self.freqs_cis
|
267 |
+
freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
|
268 |
+
|
269 |
+
mask = None
|
270 |
+
if seqlen > 1:
|
271 |
+
mask = tf.fill([seqlen, seqlen], float("-inf"))
|
272 |
+
|
273 |
+
mask = tf.linalg.band_part(mask, 0, -1)
|
274 |
+
mask = mask - tf.linalg.band_part(mask, 0, 0)
|
275 |
+
|
276 |
+
# When performing key-value caching, we compute the attention scores
|
277 |
+
# only for the new sequence. Thus, the matrix of scores is of size
|
278 |
+
# (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for
|
279 |
+
# j > cache_len + i, since row i corresponds to token cache_len + i.
|
280 |
+
mask = tf.linalg.set_diag(mask, tf.zeros(seqlen))
|
281 |
+
mask = tf.cast(mask, h.dtype)
|
282 |
+
|
283 |
+
for layer in self.layers_:
|
284 |
+
h = layer(h, start_pos, freqs_cis, mask)
|
285 |
+
h = self.norm(h)
|
286 |
+
output = tf.cast(self.output_(h), 'float32')
|
287 |
+
return output
|