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