address #1

Author: lucidrains

nGPT_pytorch/nGPT.py

@@ -1,8 +1,10 @@
+from functools import partial
+
 import torch
 from torch import nn
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
-import torch.nn.utils.parametrize as parametrize
+from torch.nn.utils.parametrize import register_parametrization
 
 from einops import rearrange
 from einops.layers.torch import Rearrange
@@ -35,16 +37,33 @@ def __init__(
         self,
         dim,
         dim_out,
-        norm_dim_in = True
+        norm_dim_in = True,
+        parametrize = True
     ):
         super().__init__()
         self.linear = nn.Linear(dim, dim_out, bias = False)
 
-        parametrize.register_parametrization(
-            self.linear,
-            'weight',
-            L2Norm(dim = -1 if norm_dim_in else 0)
-        )
+        self.parametrize = parametrize
+        self.l2norm = L2Norm(dim = -1 if norm_dim_in else 0)
+
+        if parametrize:
+            register_parametrization(
+                self.linear,
+                'weight',
+                self.l2norm
+            )
+
+        self.norm_weights_()
+
+    @torch.no_grad()
+    def norm_weights_(self):
+        if self.parametrize:
+            normed = self.weight
+            original = self.linear.parametrizations.weight.original
+
+            original.copy_(normed)
+        else:
+            self.weight.copy_(self.l2norm(self.weight))
 
     @property
     def weight(self):
@@ -62,13 +81,16 @@ def __init__(
         *,
         dim_head = 64,
         heads = 8,
-        norm_qk = True
+        norm_qk = True,
+        manual_norm_weights = False
     ):
         super().__init__()
+        NormLinear_ = partial(NormLinear, parametrize = not manual_norm_weights)
+
         dim_inner = dim_head * heads
-        self.to_q = NormLinear(dim, dim_inner)
-        self.to_k = NormLinear(dim, dim_inner)
-        self.to_v = NormLinear(dim, dim_inner)
+        self.to_q = NormLinear_(dim, dim_inner)
+        self.to_k = NormLinear_(dim, dim_inner)
+        self.to_v = NormLinear_(dim, dim_inner)
 
         self.rotary_emb = RotaryEmbedding(dim_head)
         self.qk_scale = nn.Parameter(torch.ones(dim_head) * (dim_head ** 0.25))
@@ -77,7 +99,7 @@ def __init__(
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
 
-        self.to_out = NormLinear(dim_inner, dim, norm_dim_in = False)
+        self.to_out = NormLinear_(dim_inner, dim, norm_dim_in = False)
 
     def forward(
         self,
@@ -117,19 +139,22 @@ def __init__(
         self,
         dim,
         *,
-        expand_factor = 4
+        expand_factor = 4,
+        manual_norm_weights = False
     ):
         super().__init__()
+        NormLinear_ = partial(NormLinear, parametrize = not manual_norm_weights)
+
         self.dim = dim
         dim_inner = int(dim * expand_factor * 2 / 3)
 
-        self.to_hidden = NormLinear(dim, dim_inner)
-        self.to_gate = NormLinear(dim, dim_inner)
+        self.to_hidden = NormLinear_(dim, dim_inner)
+        self.to_gate = NormLinear_(dim, dim_inner)
 
         self.hidden_scale = nn.Parameter(torch.ones(dim_inner))
         self.gate_scale = nn.Parameter(torch.ones(dim_inner))
 
-        self.to_out = NormLinear(dim_inner, dim, norm_dim_in = False)
+        self.to_out = NormLinear_(dim_inner, dim, norm_dim_in = False)
 
     def forward(self, x):
         hidden, gate = self.to_hidden(x), self.to_gate(x)
@@ -154,30 +179,33 @@ def __init__(
         attn_norm_qk = True,  # they say the query/key normalization is optional
         ff_expand_factor = 4.,
         ce_ignore_index = -1,
-        residual_lerp_scale_init = None
+        residual_lerp_scale_init = None,
+        manual_norm_weights = False
     ):
         super().__init__()
+        NormLinear_ = partial(NormLinear, parametrize = not manual_norm_weights)
+
         self.dim = dim
 
         residual_lerp_scale_init = default(residual_lerp_scale_init, 1. / depth)
 
-        self.token_embed = NormLinear(dim, num_tokens)
+        self.token_embed = NormLinear_(dim, num_tokens)
 
         self.layers = ModuleList([])
         self.residual_lerp_scales = nn.ParameterList([])
 
         for _ in range(depth):
             self.layers.append(ModuleList([
-                Attention(dim, dim_head = dim_head, heads = heads, norm_qk = attn_norm_qk),
-                FeedForward(dim, expand_factor = ff_expand_factor),
+                Attention(dim, dim_head = dim_head, heads = heads, norm_qk = attn_norm_qk, manual_norm_weights = manual_norm_weights),
+                FeedForward(dim, expand_factor = ff_expand_factor, manual_norm_weights = manual_norm_weights),
             ]))
 
             self.residual_lerp_scales.append(nn.ParameterList([
                 nn.Parameter(torch.ones(dim) * residual_lerp_scale_init),
                 nn.Parameter(torch.ones(dim) * residual_lerp_scale_init),
             ]))
 
-        self.to_logits = NormLinear(dim, num_tokens)
+        self.to_logits = NormLinear_(dim, num_tokens)
 
         self.logit_scale = nn.Parameter(torch.ones(num_tokens))
 
@@ -189,10 +217,7 @@ def norm_weights_(self):
             if not isinstance(module, NormLinear):
                 continue
 
-            normed = module.weight
-            original = module.linear.parametrizations.weight.original
-
-            original.copy_(normed)
+            module.norm_weights_()
 
     def forward(
         self,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nGPT-pytorch"
-version = "0.0.8"
+version = "0.0.9"
 description = "nGPT"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

train.py

@@ -8,6 +8,7 @@
 from torch.optim import Adam
 from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
+import torch.nn.utils.parametrize as parametrize
 
 from nGPT_pytorch import nGPT
 
@@ -89,7 +90,8 @@ def base_decoding(
 model = nGPT(
     num_tokens = 256,
     dim = 512,
-    depth = 8
+    depth = 8,
+    manual_norm_weights = True
 ).to(device)
 
 # prepare enwik8 data