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 # A BracketedSequence is a BxTx... tensor with a first and a nb time
  18 # steps to compute.
  19
  20 # Modules able to process it expect that they will have to process a
  21 # first bracket starting at t=0, followed by a succession of brackets
  22 # that move forward in time, do not overlap, and cover the axis T with
  23 # no holes.
  24 #
  25 # Although it is more general, for a classical prompt-conditioned
  26 # auto-regressive process it will be a first bracket starting at 0 and
  27 # of arbitrary length for the "prompt", followed by brackets of length
  28 # 1 for the successive tokens.
  29 #
  30 # Modules able to process brackets may implement a cache that is
  31 # resetted when the input bracket starts at t=0
  32
  33
  34 class BracketedSequence:
  35     def __init__(self, x, first=None, nb=None):
  36         self.x = x
  37         self.first = 0 if first is None else first
  38         self.nb = x.size(1) if nb is None else nb
  39
  40     def slice(self):
  41         return self.x[:, self.first : self.first + self.nb]
  42
  43
  44 ######################################################################
  45
  46
  47 class WithResidual(nn.Module):
  48     def __init__(self, *f):
  49         super().__init__()
  50         self.f = f[0] if len(f) == 1 else nn.Sequential(*f)
  51
  52     def forward(self, bs):
  53         bs.x = bs.x + self.f(bs).x
  54         return bs
  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, order):  # NxTxD, T
  89         if bs.first == 0:
  90             t = (
  91                 torch.arange(bs.x.size(1) + 1, dtype=bs.x.dtype, device=bs.x.device)[
  92                     :, None
  93                 ]
  94                 - 1
  95             )
  96             j = torch.arange(bs.x.size(2) // 2, dtype=bs.x.dtype, device=bs.x.device)[
  97                 None, :
  98             ]
  99             k = j % 2
 100             pe = (
 101                 torch.sin(
 102                     t / (self.len_max ** ((j - k) / bs.x.size(2))) + math.pi / 2 * k
 103                 )
 104                 .unsqueeze(0)
 105                 .expand(bs.x.size(0), -1, -1)
 106             )
 107
 108             order_output = order + 1
 109             order_input = F.pad(order + 1, (1, -1))
 110
 111             pe_input = pe.gather(
 112                 1, order_input.unsqueeze(-1).expand(-1, -1, pe.size(-1))
 113             )
 114             pe_output = pe.gather(
 115                 1, order_output.unsqueeze(-1).expand(-1, -1, pe.size(-1))
 116             )
 117
 118             self.pe = torch.cat((pe_input, pe_output), 2)
 119             self.cache_y = bs.x.new(bs.x.size())
 120
 121         self.cache_y[:, bs.first : bs.first + bs.nb] = (
 122             bs.slice() + self.pe[:, bs.first : bs.first + bs.nb]
 123         )
 124
 125         bs.x = self.cache_y
 126
 127         return bs
 128
 129
 130 ##############################
 131
 132
 133 class QKVAttention(nn.Module):
 134     def __init__(
 135         self,
 136         dim_in,
 137         dim_qk,
 138         dim_v,
 139         nb_heads=1,
 140         causal=False,
 141         attention_dropout=0.0,
 142         amm_generator=None,
 143     ):
 144         super().__init__()
 145
 146         def randw(*d):
 147             return nn.Parameter(torch.randn(*d) / math.sqrt(d[-1]))
 148
 149         if amm_generator is None:
 150             self.amm_generator = (
 151                 lambda d: torch.arange(d)[:, None] < torch.arange(d)[None, :]
 152             )
 153         else:
 154             self.amm_generator = amm_generator
 155
 156         self.causal = causal
 157         self.attention_dropout = attention_dropout
 158
 159         self.w_q = randw(nb_heads, dim_qk, dim_in)
 160         self.w_k = randw(nb_heads, dim_qk, dim_in)
 161         self.w_v = randw(nb_heads, dim_v, dim_in)
 162         self.w_o = randw(dim_v * nb_heads, dim_in)
 163
 164     def forward(self, bs_q):
 165         x_q = bs_q.x
 166
 167         if bs_q.first == 0:
 168             self.cache_k = x_q.new_zeros(
 169                 x_q.size(0), self.w_k.size(0), x_q.size(1), self.w_k.size(1)
 170             )
 171             self.cache_v = x_q.new_zeros(
 172                 x_q.size(0), self.w_v.size(0), x_q.size(1), self.w_v.size(1)
 173             )
 174             self.cache_y = x_q.new_zeros(x_q.size(0), x_q.size(1), self.w_o.size(1))
 175
 176         q = torch.einsum(
 177             "ntc,hdc->nhtd", x_q[:, bs_q.first : bs_q.first + bs_q.nb], self.w_q
 178         )
 179         self.cache_k[:, :, bs_q.first : bs_q.first + bs_q.nb] = torch.einsum(
 180             "ntc,hdc->nhtd", x_q[:, bs_q.first : bs_q.first + bs_q.nb], self.w_k
 181         )
 182         self.cache_v[:, :, bs_q.first : bs_q.first + bs_q.nb] = torch.einsum(
 183             "ntc,hdc->nhtd", x_q[:, bs_q.first : bs_q.first + bs_q.nb], self.w_v
 184         )
 185
 186         a = torch.einsum(
 187             "nhtd,nhsd->nhts", q, self.cache_k[:, :, : bs_q.first + bs_q.nb]
 188         ) / math.sqrt(self.w_q.size(1))
 189
 190         if self.causal:
 191             if bs_q.first == 0:
 192                 self.cache_attzero = self.amm_generator(x_q.size(1)).to(q.device)[
 193                     None, None, :, :
 194                 ]
 195             a = a.masked_fill(
 196                 self.cache_attzero[
 197                     :, :, bs_q.first : bs_q.first + bs_q.nb, : bs_q.first + bs_q.nb
 198                 ],
 199                 float("-inf"),
 200             )
 201
 202         a = a.softmax(dim=3)
 203         a = F.dropout(a, self.attention_dropout, self.training)
 204
 205         y = torch.einsum(
 206             "nhts,nhsd->nthd", a, self.cache_v[:, :, : bs_q.first + bs_q.nb]
 207         ).flatten(2)
 208
 209         self.cache_y[:, bs_q.first : bs_q.first + bs_q.nb] = y @ self.w_o
 210
 211         bs_q.x = self.cache_y
 212
 213         return bs_q
 214
 215
 216 ##############################
 217
 218
 219 class MyGPT(nn.Module):
 220     def __init__(
 221         self,
 222         vocabulary_size,
 223         dim_model,
 224         dim_keys,
 225         dim_hidden,
 226         nb_heads,
 227         nb_blocks,
 228         causal=False,
 229         dropout=0.0,
 230         len_max=1e5,
 231         amm_generator=None,
 232     ):
 233         super().__init__()
 234
 235         assert dim_model % nb_heads == 0
 236
 237         self.embedding = CacheWrapper(
 238             nn.Embedding(vocabulary_size, dim_model), nn.Dropout(dropout)
 239         )
 240         self.pe = AddPositionalEncoding(len_max)
 241
 242         trunk_blocks = []
 243
 244         for b in range(nb_blocks):
 245             trunk_blocks += [
 246                 WithResidual(
 247                     CacheWrapper(nn.LayerNorm((dim_model,))),
 248                     QKVAttention(
 249                         dim_in=dim_model,
 250                         dim_qk=dim_keys,
 251                         dim_v=dim_model // nb_heads,
 252                         nb_heads=nb_heads,
 253                         causal=causal,
 254                         attention_dropout=dropout,
 255                         amm_generator=amm_generator,
 256                     ),
 257                 ),
 258                 WithResidual(
 259                     CacheWrapper(
 260                         nn.LayerNorm((dim_model,)),
 261                         nn.Linear(in_features=dim_model, out_features=dim_hidden),
 262                         nn.ReLU(),
 263                         nn.Linear(in_features=dim_hidden, out_features=dim_model),
 264                         nn.Dropout(dropout),
 265                     ),
 266                 ),
 267             ]
 268
 269         self.trunk = nn.Sequential(*trunk_blocks)
 270
 271         self.readout = CacheWrapper(
 272             nn.Linear(in_features=dim_model, out_features=vocabulary_size)
 273         )
 274
 275         with torch.no_grad():
 276             for m in self.modules():
 277                 if isinstance(m, nn.Embedding):
 278                     m.weight.normal_(mean=0, std=2e-2)
 279                 elif isinstance(m, nn.LayerNorm):
 280                     m.bias.zero_()
 281                     m.weight.fill_(1.0)
 282
 283     def forward(self, bs, mode="standard", order=None):
 284         bs = BracketedSequence(F.pad(bs.x, (1, -1)), bs.first, bs.nb)
 285         if order is None:
 286             order = torch.arange(bs.x.size(1), device=bs.x.device)[None, :].expand_as(
 287                 bs.x
 288             )
 289         bs = self.embedding(bs)
 290         bs = self.pe(bs, order)
 291
 292         if mode == "standard":
 293             bs = self.trunk(bs)
 294             bs = self.readout(bs)
 295         elif mode == "head":
 296             bs = self.trunk(bs)
 297         elif mode == "deep":
 298             r = []
 299             for l in self.trunk:
 300                 bs = l(bs)
 301                 r += [bs.slice()]
 302             bs = BracketedSequence(torch.cat(r, -1))
 303         else:
 304             raise ValueError(f"{mode=}")
 305         return bs
 306
 307
 308 ######################################################################
 309
 310 if __name__ == "__main__":
 311     print("Basic check.")
 312
 313     vocabulary_size = 10
 314     x = torch.randint(vocabulary_size, (9, 7))
 315
 316     model = MyGPT(
 317         vocabulary_size=vocabulary_size,
 318         dim_model=18,
 319         dim_keys=50,
 320         dim_hidden=100,
 321         nb_heads=2,
 322         nb_blocks=1,
 323         dropout=0.1,
 324     )
 325
 326     model.eval()
 327
 328     y1 = model(BracketedSequence(x)).x
 329
 330     y2 = torch.randn_like(y1)
 331     for s in range(x.size(1)):
 332         z = model(BracketedSequence(x, s, 1))
 333         y2[:, s] = z.x[:, s]
 334
 335     # print(y1.max(dim = 2).values)
 336     # print(y2.max(dim = 2).values)
 337     print(f"error={((y1 - y2).norm() / (y1.norm() + y2.norm())).item()}")
 338
 339 ######################################################################