Update.
[mygptrnn.git] / mygpt.py
index ba93851..12b3631 100755 (executable)
--- a/mygpt.py
+++ b/mygpt.py
@@ -21,6 +21,8 @@ from torch.nn import functional as F
 
 import ffutils
 
+# from blanket import blanket
+
 # import memload
 
 ######################################################################
@@ -84,6 +86,18 @@ class CacheWrapper(nn.Module):
 ##############################
 
 
+class NaNChecker(nn.Module):
+    def __init__(self, name):
+        super().__init__()
+        self.name = name
+
+    def forward(self, bs):
+        x = bs.x if type(bs) is BracketedSequence else bs
+        assert not x.isnan().any(), f"${self.name} detected NaN"
+        assert not x.isinf().any(), f"${self.name} detected Inf"
+        return bs
+
+
 class WithResidual(nn.Module):
     def __init__(self, *f):
         super().__init__()
@@ -126,7 +140,6 @@ class AddPositionalEncoding(nn.Module):
 
 import pscan
 
-
 # X is /.../xTxD   A is /.../xT   Y_init is /.../xD
 
 
@@ -147,6 +160,18 @@ def pscan_dim(A, X, Y_init, dim=-2):
     return Y
 
 
+def pscan_rgrad(grad_Y, A, X, Y_init, dim=-2, eps=1e-2):
+    with torch.no_grad():
+        s_A, s_X = 0, 0
+        for t in range(X.size(dim) - 1, 0, -1):
+            delta = (grad_Y[t] - s_A) / A[t].grad
+            s_A += A[t].grad * delta
+            A[t].grad = delta
+            delta = (grad_Y[t] - s_X) / X[t].grad
+            s_X += X[t].grad * delta
+            X[t].grad = delta
+
+
 def pscan_shape(A, X, Y_init):
     s = X.size()
     A = A.reshape(-1, s[-2])
@@ -190,6 +215,8 @@ class DumbRec(nn.Module):
         nb_lines,
         attention_dropout=0.0,
         len_max=1e5,
+        logger=print,
+        args=None,
     ):
         super().__init__()
 
@@ -203,19 +230,9 @@ class DumbRec(nn.Module):
 
         self.w_qw = randw(nb_heads, dim_qk, dim_model)
         self.w_qr = randw(nb_heads, dim_qk, dim_model)
-        # self.w_k = randw(nb_heads, dim_qk, dim_model)
         self.w_v = randw(nb_heads, dim_v, dim_model)
         self.w_o = randw(dim_v * nb_heads, dim_model)
 
-    def reset_inner_loss(self):
-        self.acc_attention = 0
-        self.acc_nb = 0
-
-    def get_inner_loss(self):
-        warnings.warn("l2 regularization", RuntimeWarning)
-        return (self.acc_attention / self.acc_nb).pow(2).sum()
-        # return torch.tensor([0], device=self.w_qw.device)
-
     def forward(self, bs):
         x_q, t0, t1 = bs.x, bs.first, bs.first + bs.nb
 
@@ -223,61 +240,33 @@ class DumbRec(nn.Module):
             self.rec_v = x_q.new_zeros(
                 x_q.size(0), self.nb_lines, x_q.size(1), self.w_v.size(1)
             )
-            # self.rec_k = x_q.new_zeros(
-            # x_q.size(0), self.nb_lines, x_q.size(1), self.w_k.size(1)
-            # )
             self.cache_y = x_q.new_zeros(x_q.size(0), x_q.size(1), self.w_o.size(1))
 
-        ######################################################################
-        # Prepare the keys
-
-        k_star = self.k_star[:, None, :].expand(-1, t1 - t0, -1)
-
-        warnings.warn("rotating key barrel", RuntimeWarning)
-        k_star = self.k_star[:, None, :].expand(-1, x_q.size(1), -1)
-        t_barrel = torch.arange(t0, t1, device=k_star.device)
-        t_barrel = t_barrel[None, :].expand(k_star.size(0), t1 - t0)
-        l_barrel = (
-            torch.arange(k_star.size(0), device=k_star.device)[:, None] + t_barrel
-        ) % k_star.size(0)
-        k_star = k_star[l_barrel, t_barrel]
-
         ######################################################################
         # Compute the recurrent state
 
         qw = torch.einsum("ntc,hdc->nhtd", x_q[:, t0:t1], self.w_qw)
 
         v = torch.einsum("ntc,hdc->nhtd", x_q[:, t0:t1], self.w_v)
-        # k = torch.einsum("ntc,hdc->nhtd", x_q[:, t0:t1], self.w_k)
 
-        aw = torch.einsum(
-            "nhtd,ltd->nhlt",
-            qw,
-            k_star,
-        ) / math.sqrt(self.w_qw.size(1))
+        aw = torch.einsum("nhtd,ld->nhlt", qw, self.k_star) / math.sqrt(
+            self.w_qw.size(1)
+        )
 
         aw = aw.softmax(dim=2)  # nhlt
 
-        if self.train:
-            self.acc_attention += aw.sum(dim=(0, 1, 3))
-            self.acc_nb += aw.size(0) * aw.size(1) * aw.size(3)
-
         aw = F.dropout(aw, self.attention_dropout, self.training)
 
         A = 1 - aw.sum(dim=1)  # nlt
 
         V = torch.einsum("nhlt,nhtd->nltd", aw, v).contiguous()
-        # K = torch.einsum("nhlt,nhtd->nltd", aw, k).contiguous()
 
         if t0 == 0:
             V0 = None
-            # K0 = None
         else:
             V0 = self.rec_v[:, :, t0 - 1]
-            # K0 = self.rec_k[:, :, t0 - 1]
 
         self.rec_v[:, :, t0:t1] = pscan_shape(A, V, V0)
-        # self.rec_k[:, :, t0:t1] = pscan_shape(A, K, K0)
 
         ######################################################################
         # compute the readout
@@ -287,7 +276,6 @@ class DumbRec(nn.Module):
         ar = torch.einsum(
             "nhtd,ld->nhlt",
             qr,
-            # self.rec_k[:, :, t0:t1],
             self.k_star,
         ) / math.sqrt(self.w_qr.size(1))
 
@@ -319,6 +307,8 @@ class KVRec(nn.Module):
         nb_lines,
         attention_dropout=0.0,
         len_max=1e5,
+        logger=print,
+        args=None,
     ):
         super().__init__()
 
@@ -341,9 +331,9 @@ class KVRec(nn.Module):
         self.acc_nb = 0
 
     def get_inner_loss(self):
-        warnings.warn("l2 regularization", RuntimeWarning)
-        return (self.acc_attention / self.acc_nb).pow(2).sum()
-        return torch.tensor([0], device=self.w_qw.device)
+        warnings.warn("l2 regularization", RuntimeWarning)
+        return (self.acc_attention / self.acc_nb).pow(2).sum()
+        return torch.tensor([0], device=self.w_qw.device)
         # warnings.warn("side regularization", RuntimeWarning)
         # return (
         # (0.5 / self.nb_lines - self.acc_attention / self.acc_nb).clamp(min=0).sum()
@@ -367,12 +357,12 @@ class KVRec(nn.Module):
 
         k_star = self.k_star[:, None, :].expand(-1, t1 - t0, -1)
 
-        warnings.warn("rotating key barrel", RuntimeWarning)
+        warnings.warn("rotating key barrel", RuntimeWarning)
         k_star = self.k_star[:, None, :].expand(-1, x_q.size(1), -1)
         t_barrel = torch.arange(t0, t1, device=k_star.device)
         t_barrel = t_barrel[None, :].expand(k_star.size(0), t1 - t0)
         l_barrel = (
-            torch.arange(k_star.size(0), device=k_star.device)[:, None] + t_barrel
+            torch.arange(k_star.size(0), device=k_star.device)[:, None]  # + t_barrel
         ) % k_star.size(0)
         k_star = k_star[l_barrel, t_barrel]
 
@@ -471,28 +461,24 @@ class Caterpillar(nn.Module):
         caterpillar_height,
         attention_dropout=0.0,
         len_max=1e5,
+        logger=print,
+        args=None,
     ):
         super().__init__()
 
         warnings.warn("Caterpillar", RuntimeWarning)
 
-        def randw(*d, amplitude=None):
-            if amplitude is None:
-                amplitude = 1 / math.sqrt(d[-1])
-            return nn.Parameter(amplitude * torch.randn(*d))
+        def randw(*d, factor=1):
+            return nn.Parameter(torch.randn(*d) * factor / math.sqrt(d[-1]))
 
         self.caterpillar_length = caterpillar_length
         self.caterpillar_height = caterpillar_height
         self.attention_dropout = attention_dropout
 
-        self.proba_gate_dropout = 0.25
+        ######################################################################
 
-        self.w_G = randw(nb_heads, caterpillar_height, dim_model, amplitude=1e-5)
-        self.b_G = nn.Parameter(
-            torch.full(
-                (nb_heads, caterpillar_height), -math.log(caterpillar_height - 1)
-            )
-        )
+        self.w_G = randw(nb_heads, caterpillar_height, dim_model)
+        self.b_G = nn.Parameter(torch.full((nb_heads, caterpillar_height), 0.0))
 
         self.w_K = randw(nb_heads, dim_qk, dim_model)
         self.w_V = randw(nb_heads, dim_v, dim_model)
@@ -500,20 +486,24 @@ class Caterpillar(nn.Module):
         self.w_O = randw(dim_v * nb_heads, dim_model)
 
         self.init_K_rec = randw(
-            caterpillar_height, caterpillar_length, dim_qk, amplitude=1e-5
+            caterpillar_height,
+            caterpillar_length,
+            dim_qk,
         )
         self.init_V_rec = randw(
-            caterpillar_height, caterpillar_length, dim_v, amplitude=1e-5
+            caterpillar_height,
+            caterpillar_length,
+            dim_v,
         )
 
-    def reset_inner_loss(self):
-        self.acc_attention = 0
-        self.acc_nb = 0
+    def reset_inner_loss(self):
+    # self.acc_attention = 0
+    # self.acc_nb = 0
 
-    def get_inner_loss(self):
-        # warnings.warn("l2 regularization", RuntimeWarning)
-        # return (self.acc_attention / self.acc_nb).pow(2).sum()
-        return torch.tensor([0], device=self.w_Q.device)
+    def get_inner_loss(self):
+    # warnings.warn("l2 regularization", RuntimeWarning)
+    # return (self.acc_attention / self.acc_nb).pow(2).sum()
+    # return torch.tensor([0], device=self.w_Q.device)
 
     def forward(self, bs):
         # Dimensions to make the source a bit clearer, that's needed
@@ -526,22 +516,22 @@ class Caterpillar(nn.Module):
         DV = self.w_V.size(1)
         DK = self.w_K.size(1)
         DM = self.w_O.size(1)
-        CH = self.caterpillar_height
-        CL = self.caterpillar_length
+        R = self.caterpillar_height
+        L = self.caterpillar_length
 
         assert (
-            t0 >= CL and (t1 - t0) % CL == 0
+            t0 >= L and (t1 - t0) % L == 0
         ), f"bs.first should be greater than caterpillar_length, and bs.nb should be a multiple of caterpillar_length"
 
         # We cache values to deal efficiently with auto-regression
 
         if bs.init_cache:
-            self.rec_V = X.new_zeros(N, CH, T, DV)
-            self.rec_K = X.new_zeros(N, CH, T, DK)
+            self.rec_V = X.new_zeros(N, R, T, DV)
+            self.rec_K = X.new_zeros(N, R, T, DK)
             # We start the recurrent sequences with optimizable
             # initial values. No idea if it helps.
-            self.rec_V[:, :, t0 - CL : t0] = self.init_V_rec[None, :, :, :]
-            self.rec_K[:, :, t0 - CL : t0] = self.init_K_rec[None, :, :, :]
+            self.rec_V[:, :, t0 - L : t0, :] = self.init_V_rec[None, :, :, :]
+            self.rec_K[:, :, t0 - L : t0, :] = self.init_K_rec[None, :, :, :]
 
             self.cache_Y = X.new_zeros(N, T, DM)
 
@@ -552,8 +542,8 @@ class Caterpillar(nn.Module):
         # Compute the recurrent state
 
         # This is the Gating sequence that modulates the storing of
-        # the new key and value in the CH pairs of the current
-        # stack. There are CH independent gating values, which means
+        # the new key and value in the R pairs of the current
+        # stack. There are R independent gating values, which means
         # that the current K/V may be stored in multiple pairs of the
         # recurrent state, or not at all.
 
@@ -561,102 +551,62 @@ class Caterpillar(nn.Module):
             torch.einsum("ntc,hrc->nhrt", X, self.w_G) + self.b_G[None, :, :, None]
         ).sigmoid()
 
-        ######################################################################
-        # The "flashbacks"
-
-        if self.training and self.proba_gate_dropout > 0.0:
-            # This is a better implementation of "flashbacks".
-
-            # G is NxHxExT where e is the caterpillar's row.
-
-            warnings.warn("gate dropout", RuntimeWarning)
-            epsilon = 0.5
-
-            dropout_head = (
-                (
-                    torch.rand(G.size(), device=G.device)
-                    .flatten(2, 3)
-                    .sort(dim=2)
-                    .indices
-                    == 0
-                )
-                .unflatten(2, (CH, t1 - t0))
-                .float()
-            )
-
-            dropout_tail = dropout_head.cumsum(dim=3) - dropout_head
-
-            dropout_active = (
-                torch.rand(N, 1, 1, 1, device=G.device) < self.proba_gate_dropout
-            ).long()
-
-            dropout_head *= dropout_active
-            dropout_tail *= dropout_active
-
-            G = (
-                G
-                # + dropout_head * (1 - epsilon - G.detach())
-                - dropout_tail * G.detach()
-            )
-
-        ######################################################################
-
-        # We prepare the arguments for the parallel scan
-
         # Clip the gating to avoid values greater than 1 when several
         # heads hit the same row
 
         G = G / G.sum(1, keepdim=True).clamp(min=1)
 
-        A = 1 - G.sum(1)
+        ######################################################################
+
+        A = 1 - G.sum(dim=1)
+
         gated_V = torch.einsum("nhrt,nhtd->nrtd", G, V)
         gated_K = torch.einsum("nhrt,nhtd->nrtd", G, K)
 
         # We start from cached values, which matters in inference
 
-        init_rec_V = self.rec_V[:, :, t0 - CL : t0]
-        init_rec_K = self.rec_K[:, :, t0 - CL : t0]
-
-        #################################################################
-        # Associative scan
+        init_rec_V = self.rec_V[:, :, t0 - L : t0]
+        init_rec_K = self.rec_K[:, :, t0 - L : t0]
 
         # Here there is a trick: Since the stack at position t is
-        # computed by updating that at position t-CL, the parallel
-        # scan operates with a period of CL. To do so we split the
-        # sequence indexing in two axes, the second of size CL, and
+        # computed by updating that at position t-L, the parallel
+        # scan operates with a period of L. To do so we split the
+        # sequence indexing in two axes, the second of size L, and
         # run the parallel scan using the first as the sequence index.
 
-        A = A.unflatten(2, (-1, CL))
-        gated_V = gated_V.unflatten(2, (-1, CL))
-        gated_K = gated_K.unflatten(2, (-1, CL))
+        A = A.unflatten(2, (-1, L))
+        gated_V = gated_V.unflatten(2, (-1, L))
+        gated_K = gated_K.unflatten(2, (-1, L))
 
-        next_V = pscan_dim(A, gated_V, init_rec_V, dim=2)
-        next_K = pscan_dim(A, gated_K, init_rec_K, dim=2)
+        next_V = pscan_dim(A, gated_V, init_rec_V, dim=2).flatten(2, 3)
+        next_K = pscan_dim(A, gated_K, init_rec_K, dim=2).flatten(2, 3)
 
-        self.rec_V[:, :, t0:t1] = next_V.flatten(2, 3)
-        self.rec_K[:, :, t0:t1] = next_K.flatten(2, 3)
+        self.rec_V[:, :, t0:t1] = next_V
+        self.rec_K[:, :, t0:t1] = next_K
 
         ######################################################################
         # compute the readout
 
         Q = torch.einsum("ntc,hdc->nhtd", X, self.w_Q)
 
-        # We build tensors NxHxTxFxL where N is the sample index, H
-        # the head, T the time, F the row in the caterpillar, and L
+        # Q = blanket(Q)
+
+        # We build tensors NxHxTxRxL where N is the sample index, H
+        # the head, T the time, R the row in the caterpillar, and L
         # the column in the caterpillar
 
         windowed_V = moving_window(
-            self.rec_V[:, :, t0 - CL + 1 : t1], dim=2, win_dim=3, win_size=CL
+            self.rec_V[:, :, t0 - L + 1 : t1], dim=2, win_dim=3, win_size=L
         )
 
         windowed_K = moving_window(
-            self.rec_K[:, :, t0 - CL + 1 : t1], dim=2, win_dim=3, win_size=CL
+            self.rec_K[:, :, t0 - L + 1 : t1], dim=2, win_dim=3, win_size=L
         )
 
-        # We have an attention score for each of the CHxCL values
+        # We have an attention score for each of the RxL values
 
         ar = torch.einsum(
-            "nhtd,nftld->nhtfl",
+            "nhtd,nrtld->nhtrl",
             Q,
             windowed_K,
         ) / math.sqrt(DK)
@@ -676,8 +626,6 @@ class Caterpillar(nn.Module):
             windowed_V,
         ).flatten(2)
 
-        # Compute the final output
-
         self.cache_Y[:, t0:t1] = Y @ self.w_O
 
         return BracketedSequence(self.cache_Y, t0, t1 - t0, bs.init_cache)
@@ -694,7 +642,10 @@ class QKVAttention(nn.Module):
         dim_v,
         nb_heads=1,
         causal=False,
+        horizon=None,
         attention_dropout=0.0,
+        logger=print,
+        args=None,
     ):
         super().__init__()
 
@@ -702,6 +653,7 @@ class QKVAttention(nn.Module):
             return nn.Parameter(torch.randn(*d) / math.sqrt(d[-1]))
 
         self.causal = causal
+        self.horizon = horizon
         self.attention_dropout = attention_dropout
         self.record_attention = False
 
@@ -745,6 +697,17 @@ class QKVAttention(nn.Module):
                     torch.arange(x_q.size(1), device=q.device)[None, None, :, None]
                     < torch.arange(x_q.size(1), device=q.device)[None, None, None, :]
                 )
+
+                if self.horizon is not None:
+                    self.cache_attzero = torch.logical_or(
+                        self.cache_attzero,
+                        torch.arange(x_q.size(1), device=q.device)[None, None, :, None]
+                        >= torch.arange(x_q.size(1), device=q.device)[
+                            None, None, None, :
+                        ]
+                        + self.horizon,
+                    )
+
             a = a.masked_fill(
                 self.cache_attzero[
                     :, :, bs.first : bs.first + bs.nb, : bs.first + bs.nb
@@ -785,18 +748,23 @@ class MyGPT(nn.Module):
         causal=False,
         dropout=0.0,
         len_max=1e5,
-        attention_layer="kvrec",
+        attention_layer="caterpillar",
+        logger=print,
+        args=None,
     ):
         super().__init__()
 
+        self.vocabulary_size = vocabulary_size
+
         assert attention_layer in {
             "mha",
             "dumbrec",
             "kvrec",
             "caterpillar",
+            "attcat",
         }, f"Unknown attention operator {attention_layer}."
 
-        if attention_layer == "caterpillar":
+        if attention_layer == "caterpillar" or attention_layer == "attcat":
             assert nb_lines % caterpillar_height == 0
             self.caterpillar_length = nb_lines // caterpillar_height
             self.caterpillar_height = caterpillar_height
@@ -815,51 +783,99 @@ class MyGPT(nn.Module):
 
         def attlayer():
             if attention_layer == "mha":
-                return QKVAttention(
-                    dim_model=dim_model,
-                    dim_qk=dim_keys,
-                    dim_v=dim_model // nb_heads,
-                    nb_heads=nb_heads,
-                    causal=causal,
-                    attention_dropout=dropout,
+                return WithResidual(
+                    CacheWrapper(nn.LayerNorm((dim_model,))),
+                    QKVAttention(
+                        dim_model=dim_model,
+                        dim_qk=dim_keys,
+                        dim_v=dim_model // nb_heads,
+                        nb_heads=nb_heads,
+                        causal=causal,
+                        attention_dropout=dropout,
+                        logger=logger,
+                        args=args,
+                    ),
                 )
             elif attention_layer == "dumbrec":
-                return DumbRec(
-                    dim_model=dim_model,
-                    dim_qk=dim_keys,
-                    dim_v=dim_model // nb_heads,
-                    nb_heads=nb_heads,
-                    nb_lines=nb_lines,
-                    attention_dropout=dropout,
+                return WithResidual(
+                    CacheWrapper(nn.LayerNorm((dim_model,))),
+                    DumbRec(
+                        dim_model=dim_model,
+                        dim_qk=dim_keys,
+                        dim_v=dim_model // nb_heads,
+                        nb_heads=nb_heads,
+                        nb_lines=nb_lines,
+                        attention_dropout=dropout,
+                        logger=logger,
+                        args=args,
+                    ),
                 )
             elif attention_layer == "kvrec":
-                return KVRec(
-                    dim_model=dim_model,
-                    dim_qk=dim_keys,
-                    dim_v=dim_model // nb_heads,
-                    nb_heads=nb_heads,
-                    nb_lines=nb_lines,
-                    attention_dropout=dropout,
+                return WithResidual(
+                    CacheWrapper(nn.LayerNorm((dim_model,))),
+                    KVRec(
+                        dim_model=dim_model,
+                        dim_qk=dim_keys,
+                        dim_v=dim_model // nb_heads,
+                        nb_heads=nb_heads,
+                        nb_lines=nb_lines,
+                        attention_dropout=dropout,
+                        logger=logger,
+                        args=args,
+                    ),
                 )
             elif attention_layer == "caterpillar":
-                return Caterpillar(
-                    dim_model=dim_model,
-                    dim_qk=dim_keys,
-                    dim_v=dim_model // nb_heads,
-                    nb_heads=nb_heads,
-                    caterpillar_length=self.caterpillar_length,
-                    caterpillar_height=self.caterpillar_height,
-                    attention_dropout=dropout,
+                return WithResidual(
+                    CacheWrapper(nn.LayerNorm((dim_model,))),
+                    Caterpillar(
+                        dim_model=dim_model,
+                        dim_qk=dim_keys,
+                        dim_v=dim_model // nb_heads,
+                        nb_heads=nb_heads,
+                        caterpillar_length=self.caterpillar_length,
+                        caterpillar_height=self.caterpillar_height,
+                        attention_dropout=dropout,
+                        logger=logger,
+                        args=args,
+                    ),
+                )
+            elif attention_layer == "attcat":
+                return nn.Sequential(
+                    WithResidual(
+                        CacheWrapper(nn.LayerNorm((dim_model,))),
+                        QKVAttention(
+                            dim_model=dim_model,
+                            dim_qk=dim_keys,
+                            dim_v=dim_model // nb_heads,
+                            nb_heads=nb_heads,
+                            causal=causal,
+                            horizon=self.caterpillar_length,
+                            attention_dropout=dropout,
+                            logger=logger,
+                            args=args,
+                        ),
+                    ),
+                    WithResidual(
+                        CacheWrapper(nn.LayerNorm((dim_model,))),
+                        Caterpillar(
+                            dim_model=dim_model,
+                            dim_qk=dim_keys,
+                            dim_v=dim_model // nb_heads,
+                            nb_heads=nb_heads,
+                            caterpillar_length=self.caterpillar_length,
+                            caterpillar_height=self.caterpillar_height,
+                            attention_dropout=dropout,
+                            logger=logger,
+                            args=args,
+                        ),
+                    ),
                 )
             else:
                 raise ValueError(f"Unknown attention type {attention_layer}.")
 
         for b in range(nb_blocks):
             trunk_blocks += [
-                WithResidual(
-                    CacheWrapper(nn.LayerNorm((dim_model,))),
-                    attlayer(),
-                ),
+                attlayer(),
                 WithResidual(
                     CacheWrapper(
                         nn.LayerNorm((dim_model,)),
@@ -980,7 +996,115 @@ class MyGPT(nn.Module):
 ######################################################################
 
 if __name__ == "__main__":
-    print("Basic check.")
+    import argparse
+
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import matplotlib.collections as mc
+
+    args = argparse.Namespace(
+        gate_dropout_proba=0.0, gate_dropout_sync=True, gate_dropout_replace=False
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    dim_model, dim_keys, nb_heads = 512, 64, 1
+    dropout = 0.1
+
+    caterpillar = Caterpillar(
+        dim_model=dim_model,
+        dim_qk=dim_keys,
+        dim_v=dim_model // nb_heads,
+        nb_heads=nb_heads,
+        caterpillar_length=16,
+        caterpillar_height=32,
+        attention_dropout=dropout,
+        args=args,
+    ).to(device)
+
+    qkv = QKVAttention(
+        dim_model=dim_model,
+        dim_qk=dim_keys,
+        dim_v=dim_model // nb_heads,
+        nb_heads=nb_heads,
+        causal=True,
+        attention_dropout=dropout,
+        args=args,
+    ).to(device)
+
+    linear = CacheWrapper(nn.Linear(512, 512)).to(device)
+
+    x = torch.randn(1, 256, dim_model)
+
+    x = x.to(device)
+    x.requires_grad_()
+
+    ######################################################################
+
+    fig = plt.figure()
+    fig.set_figheight(6)
+    fig.set_figwidth(8)
+
+    ax = fig.add_subplot(1, 1, 1)
+
+    # ax.set_xlim(-1.5, 1.5)
+    # ax.set_ylim(-1.5, 1.5)
+    # ax.set(aspect=1)
+    # ax.spines.right.set_visible(False)
+    # ax.spines.top.set_visible(False)
+
+    # dt = 0.01
+    # t = np.arange(dt, 20.0, dt)
+    # ax.semilogx(t, np.exp(-t / 5.0))
+    # ax.grid()
+    ax.set_yscale("log")
+
+    ######################################################################
+
+    for label, model, thickness in [
+        ("nn.Linear", linear, 0.2),
+        ("mygpy.QKVAttention", qkv, 1),
+        ("mygpt.Caterpillar", caterpillar, 2),
+    ]:
+        y = model(BracketedSequence(x, 32, x.size(1) - 32, init_cache=True)).x
+
+        for n, p in [("input", x)] + list(model.named_parameters()):
+            print(f"Processing {model}.{n}")
+            data = []
+            for t in range(y.size(1)):
+                sg = 0
+                for d in torch.randperm(y.size(2))[:8]:
+                    sg += torch.autograd.grad(y[0, t, d], p, retain_graph=True)[0]
+                assert not sg.isinf().any()
+                assert not sg.isnan().any()
+                data.append([t, sg.sum().item()])
+
+            data = torch.tensor(data)
+            # cx, cy = data[:, 0], data[:, 1]
+            cy = data[:, 1].sort().values
+            cx = torch.linspace(0, 1, cy.size(0))
+            ax.plot(
+                cx, cy, label=label + "." + n, linewidth=thickness
+            )  # , color='gray', label='Input')
+
+    # ax.legend(frameon=False, loc="top right")
+
+    # Put a legend to the right of the current axis
+    box = ax.get_position()
+    ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
+    ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+
+    filename = "plot.pdf"
+    print(f"saving {filename}")
+    fig.savefig(filename, bbox_inches="tight")
+
+    # if args.window and hasattr(plt.get_current_fig_manager(), 'window'):
+    # plt.get_current_fig_manager().window.setGeometry(2, 2, 1024, 768)
+    # plt.show()
+
+    exit(0)
+
+    ######################################################################
 
     m = Caterpillar(
         dim_model=4,
@@ -1002,8 +1126,6 @@ if __name__ == "__main__":
     print((y1 - torch.cat([y3a, y3b], dim=1)).abs().max())
     exit(0)
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
     vocabulary_size = 128
     x = torch.randint(vocabulary_size, (6, 1024))