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
  49 ######################################################################
  50
  51
  52 class CacheWrapper(nn.Module):
  53     def __init__(self, *f):
  54         super().__init__()
  55         self.f = f[0] if len(f) == 1 else nn.Sequential(*f)
  56
  57     def forward(self, bs):
  58         if bs.first == 0:
  59             y = self.f(bs.slice())
  60             self.cache_y = y.new(*((y.size(0), bs.x.size(1)) + y.size()[2:]))
  61             self.cache_y[:, bs.first : bs.first + bs.nb] = y
  62         else:
  63             self.cache_y[:, bs.first : bs.first + bs.nb] = self.f(bs.slice())
  64
  65         bs.x = self.cache_y
  66
  67         return bs
  68
  69
  70 ##############################
  71
  72
  73 class WithResidual(nn.Module):
  74     def __init__(self, *f):
  75         super().__init__()
  76         self.f = f[0] if len(f) == 1 else nn.Sequential(*f)
  77
  78     def forward(self, bs):
  79         bs.x = bs.x + self.f(bs).x
  80         return bs
  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         bs.x = self.cache_y
 112
 113         return bs
 114
 115
 116 ##############################
 117
 118
 119 class QKVAttention(nn.Module):
 120     def __init__(
 121         self, dim_in, dim_qk, dim_v, nb_heads=1, causal=False, attention_dropout=0.0
 122     ):
 123         super().__init__()
 124
 125         def randw(*d):
 126             return nn.Parameter(torch.randn(*d) / math.sqrt(d[-1]))
 127
 128         self.causal = causal
 129         self.attention_dropout = attention_dropout
 130
 131         self.w_q = randw(nb_heads, dim_qk, dim_in)
 132         self.w_k = randw(nb_heads, dim_qk, dim_in)
 133         self.w_v = randw(nb_heads, dim_v, dim_in)
 134         self.w_o = randw(dim_v * nb_heads, dim_in)
 135
 136     def forward(self, bs_q):
 137         x_q = bs_q.x
 138
 139         if bs_q.first == 0:
 140             self.cache_k = x_q.new_zeros(
 141                 x_q.size(0), self.w_k.size(0), x_q.size(1), self.w_k.size(1)
 142             )
 143             self.cache_v = x_q.new_zeros(
 144                 x_q.size(0), self.w_v.size(0), x_q.size(1), self.w_v.size(1)
 145             )
 146             self.cache_y = x_q.new_zeros(x_q.size(0), x_q.size(1), self.w_o.size(1))
 147
 148         q = torch.einsum(
 149             "ntc,hdc->nhtd", x_q[:, bs_q.first : bs_q.first + bs_q.nb], self.w_q
 150         )
 151         self.cache_k[:, :, bs_q.first : bs_q.first + bs_q.nb] = torch.einsum(
 152             "ntc,hdc->nhtd", x_q[:, bs_q.first : bs_q.first + bs_q.nb], self.w_k
 153         )
 154         self.cache_v[:, :, 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_v
 156         )
 157
 158         a = torch.einsum(
 159             "nhtd,nhsd->nhts", q, self.cache_k[:, :, : bs_q.first + bs_q.nb]
 160         ) / math.sqrt(self.w_q.size(1))
 161
 162         if self.causal:
 163             if bs_q.first == 0:
 164                 self.cache_attzero = (
 165                     torch.arange(x_q.size(1), device=q.device)[None, None, :, None]
 166                     < torch.arange(x_q.size(1), device=q.device)[None, None, None, :]
 167                 )
 168             a = a.masked_fill(
 169                 self.cache_attzero[
 170                     :, :, bs_q.first : bs_q.first + bs_q.nb, : bs_q.first + bs_q.nb
 171                 ],
 172                 float("-inf"),
 173             )
 174
 175         a = a.softmax(dim=3)
 176         a = F.dropout(a, self.attention_dropout, self.training)
 177
 178         y = torch.einsum(
 179             "nhts,nhsd->nthd", a, self.cache_v[:, :, : bs_q.first + bs_q.nb]
 180         ).flatten(2)
 181
 182         self.cache_y[:, bs_q.first : bs_q.first + bs_q.nb] = y @ self.w_o
 183
 184         bs_q.x = self.cache_y
 185
 186         return bs_q
 187
 188
 189 ##############################
 190
 191
 192 class MyGPT(nn.Module):
 193     def __init__(
 194         self,
 195         vocabulary_size,
 196         dim_model,
 197         dim_keys,
 198         dim_hidden,
 199         nb_heads,
 200         nb_blocks,
 201         causal=False,
 202         dropout=0.0,
 203         len_max=1e5,
 204     ):
 205         super().__init__()
 206
 207         assert dim_model % nb_heads == 0
 208
 209         self.embedding = nn.Sequential(
 210             CacheWrapper(nn.Embedding(vocabulary_size, dim_model), nn.Dropout(dropout)),
 211             AddPositionalEncoding(len_max),
 212         )
 213
 214         trunk_blocks = []
 215
 216         for b in range(nb_blocks):
 217             trunk_blocks += [
 218                 WithResidual(
 219                     CacheWrapper(nn.LayerNorm((dim_model,))),
 220                     QKVAttention(
 221                         dim_in=dim_model,
 222                         dim_qk=dim_keys,
 223                         dim_v=dim_model // nb_heads,
 224                         nb_heads=nb_heads,
 225                         causal=causal,
 226                         attention_dropout=dropout,
 227                     ),
 228                 ),
 229                 WithResidual(
 230                     CacheWrapper(
 231                         nn.LayerNorm((dim_model,)),
 232                         nn.Linear(in_features=dim_model, out_features=dim_hidden),
 233                         nn.ReLU(),
 234                         nn.Linear(in_features=dim_hidden, out_features=dim_model),
 235                         nn.Dropout(dropout),
 236                     ),
 237                 ),
 238             ]
 239
 240         self.trunk = nn.Sequential(*trunk_blocks)
 241
 242         self.readout = CacheWrapper(
 243             nn.Linear(in_features=dim_model, out_features=vocabulary_size)
 244         )
 245
 246         with torch.no_grad():
 247             for m in self.modules():
 248                 if isinstance(m, nn.Embedding):
 249                     m.weight.normal_(mean=0, std=2e-2)
 250                 elif isinstance(m, nn.LayerNorm):
 251                     m.bias.zero_()
 252                     m.weight.fill_(1.0)
 253
 254     def forward(self, bs):
 255         bs.x = F.pad(bs.x, (1, -1))
 256         bs = self.embedding(bs)
 257         bs = self.trunk(bs)
 258         bs = self.readout(bs)
 259         return bs
 260
 261
 262 ######################################################################
 263
 264 if __name__ == "__main__":
 265     print("Basic check.")
 266
 267     vocabulary_size = 10
 268     x = torch.randint(vocabulary_size, (9, 7))
 269
 270     model = MyGPT(
 271         vocabulary_size=vocabulary_size,
 272         dim_model=18,
 273         dim_keys=50,
 274         dim_hidden=100,
 275         nb_heads=2,
 276         nb_blocks=1,
 277         dropout=0.1,
 278     )
 279
 280     model.eval()
 281
 282     y1 = model(BracketedSequence(x)).x
 283
 284     y2 = torch.randn_like(y1)
 285     for s in range(x.size(1)):
 286         z = model(BracketedSequence(x, s, 1))
 287         y2[:, s] = z.x[:, s]
 288
 289     # print(y1.max(dim = 2).values)
 290     # print(y2.max(dim = 2).values)
 291     print(f"error={((y1 - y2).norm() / (y1.norm() + y2.norm())).item()}")
 292
 293 ######################################################################