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