Initial commit.
[picoclvr.git] / mygpt.py
1 #!/usr/bin/env python
2
3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
5
6 # Written by Francois Fleuret <francois@fleuret.org>
7
8 import math
9
10 import torch
11
12 from torch import nn
13 from torch.nn import functional as F
14
15 ######################################################################
16
17
18 class WithResidual(nn.Module):
19     def __init__(self, *f):
20         super().__init__()
21         self.f = f[0] if len(f) == 1 else nn.Sequential(*f)
22
23     def forward(self, bs):
24         bs.x = bs.x + self.f(bs).x
25         return bs
26
27
28 ######################################################################
29
30 # A BracketedSequence is a BxTx... tensor with a first and a nb time
31 # steps to compute.
32
33 # Modules able to process it expect that they will have to process a
34 # first bracket starting at t=0, followed by a succession of brackets
35 # that move forward in time, do not overlap, and cover the axis T with
36 # no holes.
37 #
38 # Although it is more general, for a classical prompt-conditioned
39 # auto-regressive process it will be a first bracket starting at 0 and
40 # of arbitrary length for the "prompt", followed by brackets of length
41 # 1 for the successive tokens.
42 #
43 # Modules able to process brackets may implement a cache that is
44 # resetted when the input bracket starts at t=0
45
46
47 class BracketedSequence:
48     def __init__(self, x, first=None, nb=None):
49         self.x = x
50         self.first = 0 if first is None else first
51         self.nb = x.size(1) if nb is None else nb
52
53     def slice(self):
54         return self.x[:, self.first : self.first + self.nb]
55
56
57 ######################################################################
58
59
60 class CacheWrapper(nn.Module):
61     def __init__(self, *f):
62         super().__init__()
63         self.f = f[0] if len(f) == 1 else nn.Sequential(*f)
64
65     def forward(self, bs):
66         if bs.first == 0:
67             y = self.f(bs.slice())
68             self.cache_y = y.new(*((y.size(0), bs.x.size(1)) + y.size()[2:]))
69             self.cache_y[:, bs.first : bs.first + bs.nb] = y
70         else:
71             self.cache_y[:, bs.first : bs.first + bs.nb] = self.f(bs.slice())
72
73         bs.x = self.cache_y
74
75         return bs
76
77
78 ##############################
79
80
81 class AddPositionalEncoding(nn.Module):
82     def __init__(self, len_max):
83         super().__init__()
84         self.len_max = len_max
85
86     # [Vaswani et al 2018] PE_{t,2i} = sin(t/(L^{2i/D})), PE_{t,2i+1} = cos(t/(L^{2i/D}))
87
88     def forward(self, bs):
89         if bs.first == 0:
90             t = torch.arange(bs.x.size(1), dtype=bs.x.dtype, device=bs.x.device)[
91                 :, None
92             ]
93             j = torch.arange(bs.x.size(2), dtype=bs.x.dtype, device=bs.x.device)[
94                 None, :
95             ]
96             k = j % 2
97             self.pe = torch.sin(
98                 t / (self.len_max ** ((j - k) / bs.x.size(2))) + math.pi / 2 * k
99             )
100             self.cache_y = bs.x.new(bs.x.size())
101
102         self.cache_y[:, bs.first : bs.first + bs.nb] = (
103             bs.slice() + self.pe[bs.first : bs.first + bs.nb]
104         )
105
106         bs.x = self.cache_y
107
108         return bs
109
110
111 ##############################
112
113
114 class QKVAttention(nn.Module):
115     def __init__(
116         self, dim_in, dim_qk, dim_v, nb_heads=1, causal=False, attention_dropout=0.0
117     ):
118         super().__init__()
119
120         def randw(*d):
121             return nn.Parameter(torch.randn(*d) / math.sqrt(d[-1]))
122
123         self.causal = causal
124         self.attention_dropout = attention_dropout
125
126         self.w_q = randw(nb_heads, dim_qk, dim_in)
127         self.w_k = randw(nb_heads, dim_qk, dim_in)
128         self.w_v = randw(nb_heads, dim_v, dim_in)
129         self.w_o = randw(dim_v * nb_heads, dim_in)
130
131     def forward(self, bs_q, x_kv=None):
132         x_q = bs_q.x
133         if x_kv is None:
134             x_kv = x_q
135
136         if bs_q.first == 0:
137             self.cache_k = x_q.new_zeros(
138                 x_q.size(0), self.w_k.size(0), x_kv.size(1), self.w_k.size(1)
139             )
140             self.cache_v = x_q.new_zeros(
141                 x_q.size(0), self.w_v.size(0), x_kv.size(1), self.w_v.size(1)
142             )
143             self.cache_y = x_q.new_zeros(x_q.size(0), x_q.size(1), self.w_o.size(1))
144
145         q = torch.einsum(
146             "ntc,hdc->nhtd", x_q[:, bs_q.first : bs_q.first + bs_q.nb], self.w_q
147         )
148         self.cache_k[:, :, bs_q.first : bs_q.first + bs_q.nb] = torch.einsum(
149             "ntc,hdc->nhtd", x_kv[:, bs_q.first : bs_q.first + bs_q.nb], self.w_k
150         )
151         self.cache_v[:, :, bs_q.first : bs_q.first + bs_q.nb] = torch.einsum(
152             "ntc,hdc->nhtd", x_kv[:, bs_q.first : bs_q.first + bs_q.nb], self.w_v
153         )
154
155         a = torch.einsum(
156             "nhtd,nhsd->nhts", q, self.cache_k[:, :, : bs_q.first + bs_q.nb]
157         ) / math.sqrt(self.w_q.size(1))
158
159         if self.causal:
160             if bs_q.first == 0:
161                 self.cache_attzero = (
162                     torch.arange(x_q.size(1), device=q.device)[None, None, :, None]
163                     < torch.arange(x_kv.size(1), device=q.device)[None, None, None, :]
164                 )
165             a = a.masked_fill(
166                 self.cache_attzero[
167                     :, :, bs_q.first : bs_q.first + bs_q.nb, : bs_q.first + bs_q.nb
168                 ],
169                 float("-inf"),
170             )
171
172         a = a.softmax(dim=3)
173         a = F.dropout(a, self.attention_dropout, self.training)
174
175         y = torch.einsum(
176             "nhts,nhsd->nthd", a, self.cache_v[:, :, : bs_q.first + bs_q.nb]
177         ).flatten(2)
178
179         self.cache_y[:, bs_q.first : bs_q.first + bs_q.nb] = y @ self.w_o
180
181         bs_q.x = self.cache_y
182
183         return bs_q
184
185
186 ##############################
187
188
189 class MyGPT(nn.Module):
190     def __init__(
191         self,
192         vocabulary_size,
193         dim_model,
194         dim_keys,
195         dim_hidden,
196         nb_heads,
197         nb_blocks,
198         causal=False,
199         dropout=0.0,
200         len_max=1e5,
201     ):
202
203         super().__init__()
204
205         assert dim_model % nb_heads == 0
206
207         self.embedding = nn.Sequential(
208             CacheWrapper(nn.Embedding(vocabulary_size, dim_model), nn.Dropout(dropout)),
209             AddPositionalEncoding(len_max),
210         )
211
212         trunk_blocks = []
213
214         for b in range(nb_blocks):
215             trunk_blocks += [
216                 WithResidual(
217                     CacheWrapper(nn.LayerNorm((dim_model,))),
218                     QKVAttention(
219                         dim_in=dim_model,
220                         dim_qk=dim_keys,
221                         dim_v=dim_model // nb_heads,
222                         nb_heads=nb_heads,
223                         causal=causal,
224                         attention_dropout=dropout,
225                     ),
226                 ),
227                 WithResidual(
228                     CacheWrapper(
229                         nn.LayerNorm((dim_model,)),
230                         nn.Linear(in_features=dim_model, out_features=dim_hidden),
231                         nn.ReLU(),
232                         nn.Linear(in_features=dim_hidden, out_features=dim_model),
233                         nn.Dropout(dropout),
234                     ),
235                 ),
236             ]
237
238         self.trunk = nn.Sequential(*trunk_blocks)
239
240         self.readout = CacheWrapper(
241             nn.Linear(in_features=dim_model, out_features=vocabulary_size)
242         )
243
244         with torch.no_grad():
245             for m in self.modules():
246                 if isinstance(m, nn.Embedding):
247                     m.weight.normal_(mean=0, std=2e-2)
248                 elif isinstance(m, nn.LayerNorm):
249                     m.bias.zero_()
250                     m.weight.fill_(1.0)
251
252     def forward(self, bs):
253         bs.x = F.pad(bs.x, (1, -1))
254         bs = self.embedding(bs)
255         bs = self.trunk(bs)
256         bs = self.readout(bs)
257         return bs
258
259
260 ######################################################################
261
262 if __name__ == "__main__":
263
264     print("Basic check.")
265
266     vocabulary_size = 10
267     x = torch.randint(vocabulary_size, (9, 7))
268
269     model = MyGPT(
270         vocabulary_size=vocabulary_size,
271         dim_model=18,
272         dim_keys=50,
273         dim_hidden=100,
274         nb_heads=2,
275         nb_blocks=1,
276         dropout=0.1,
277     )
278
279     model.eval()
280
281     y1 = model(BracketedSequence(x)).x
282
283     y2 = torch.randn_like(y1)
284     for s in range(x.size(1)):
285         z = model(BracketedSequence(x, s, 1))
286         y2[:, s] = z.x[:, s]
287
288     # print(y1.max(dim = 2).values)
289     # print(y2.max(dim = 2).values)
290     print(f"error={((y1 - y2).norm() / (y1.norm() + y2.norm())).item()}")
291
292 ######################################################################