x-transformers 2.8.3__tar.gz → 2.9.0__tar.gz

{x_transformers-2.8.3 → x_transformers-2.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.8.3
+Version: 2.9.0
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -40,6 +40,7 @@ Requires-Dist: loguru
 Requires-Dist: packaging>=21.0
 Requires-Dist: torch>=2.0
 Provides-Extra: examples
+Requires-Dist: adam-atan2-pytorch>=0.2.2; extra == 'examples'
 Requires-Dist: lion-pytorch; extra == 'examples'
 Requires-Dist: tqdm; extra == 'examples'
 Provides-Extra: test

{x_transformers-2.8.3 → x_transformers-2.9.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.8.3"
+version = "2.9.0"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -35,6 +35,7 @@ Repository = "https://github.com/lucidrains/x-transformers"
 
 [project.optional-dependencies]
 examples = [
+    "adam-atan2-pytorch>=0.2.2",
     "lion-pytorch",
     "tqdm",
 ]

{x_transformers-2.8.3 → x_transformers-2.9.0}/tests/test_x_transformers.py

@@ -1362,7 +1362,7 @@ def test_vae():
     out = model.generate(seq[:, :512], 512, seq_for_latents = style)
 
 def test_muon_params():
-    from x_transformers import Attention, FeedForward
+    from x_transformers import Attention, FeedForward, Encoder
 
     attn = Attention(dim = 512, dim_out = 384)
     assert len(list(attn.muon_parameters())) == 2
@@ -1370,3 +1370,6 @@ def test_muon_params():
     ff = FeedForward(dim = 512)
 
     assert len(list(ff.muon_parameters())) == 2
+
+    enc = Encoder(dim = 512, depth = 2)
+    assert len(enc.muon_parameters()) > 0

x_transformers-2.9.0/train_with_muon.py

@@ -0,0 +1,132 @@
+# /// script
+# dependencies = [
+#   "x-transformers",
+#   "adam-atan2-pytorch>=0.2.2",
+# ]
+# ///
+
+from x_transformers import TransformerWrapper, Decoder
+from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
+
+import random
+import tqdm
+import gzip
+import numpy as np
+import torch
+import torch.optim as optim
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+
+from adam_atan2_pytorch import MuonAdamAtan2
+
+# constants
+
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 1e-4
+MUON_LEARNING_RATE = 1e-3
+VALIDATE_EVERY  = 100
+GENERATE_EVERY  = 500
+GENERATE_LENGTH = 1024
+SEQ_LEN = 1024
+
+# helpers
+
+def cycle(loader):
+    while True:
+        for data in loader:
+            yield data
+
+def decode_token(token):
+    return str(chr(max(32, token)))
+
+def decode_tokens(tokens):
+    return ''.join(list(map(decode_token, tokens)))
+
+# instantiate GPT-like decoder model
+
+model = TransformerWrapper(
+    num_tokens = 256,
+    max_seq_len = SEQ_LEN,
+    attn_layers = Decoder(
+        dim = 512,
+        depth = 6,
+        heads = 8,
+        rotary_pos_emb = True
+    )
+)
+
+ar_wrapper = AutoregressiveWrapper(model)
+model.cuda()
+
+# prepare enwik8 data
+
+with gzip.open('./data/enwik8.gz') as file:
+    data = np.frombuffer(file.read(int(95e6)), dtype=np.uint8).copy()
+    train_x, valid_x = np.split(data, [int(90e6)])
+    data_train, data_val = torch.from_numpy(train_x), torch.from_numpy(valid_x)
+
+class TextSamplerDataset(Dataset):
+    def __init__(self, data, seq_len):
+        super().__init__()
+        self.data = data
+        self.seq_len = seq_len
+
+    def __getitem__(self, index):
+        rand_start = torch.randint(0, self.data.size(0) - self.seq_len - 1, (1,))
+        full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()
+        return full_seq.cuda()
+
+    def __len__(self):
+        return self.data.size(0) // self.seq_len
+
+train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
+val_dataset   = TextSamplerDataset(data_val, SEQ_LEN)
+train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE, drop_last = True))
+val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE, drop_last = True))
+
+# optimizer
+
+optim = MuonAdamAtan2(
+    muon_params = model.muon_parameters(),
+    params = model.parameters(),
+    remove_muon_params_from_params = True,
+    lr = LEARNING_RATE,
+    muon_lr = MUON_LEARNING_RATE,
+)
+
+# training
+
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
+    model.train()
+
+    for __ in range(GRADIENT_ACCUMULATE_EVERY):
+        loss = ar_wrapper(next(train_loader))
+        (loss / GRADIENT_ACCUMULATE_EVERY).backward()
+
+    print(f'training loss: {loss.item()}')
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+    optim.step()
+    optim.zero_grad()
+
+    if i % VALIDATE_EVERY == 0:
+        model.eval()
+        with torch.no_grad():
+            loss = ar_wrapper(next(val_loader))
+            print(f'validation loss: {loss.item()}')
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:-1]
+        prime = decode_tokens(inp)
+        print(f'%s \n\n %s', (prime, '*' * 100))
+
+        sample = ar_wrapper.generate(
+            prompts = inp,
+            seq_len = GENERATE_LENGTH,
+            cache_kv = True
+        )
+
+        output_str = decode_tokens(sample)
+        print(output_str)

{x_transformers-2.8.3 → x_transformers-2.9.0}/x_transformers/x_transformers.py

@@ -2493,6 +2493,17 @@ class AttentionLayers(Module):
         for attn_layer, attn_inter in zip(attn_layers, attn_intermeds):
             attn_layer.qk_clip_(attn_inter, tau = tau)
 
+    def muon_parameters(self):
+        params = []
+
+        for m in self.modules():
+            if not isinstance(m, (Attention, FeedForward)):
+                continue
+
+            params.extend(list(m.muon_parameters()))
+
+        return params
+
     def forward(
         self,
         x,
@@ -3230,6 +3241,9 @@ class TransformerWrapper(Module):
     ):
         self.attn_layers.attn_qk_clip_(intermediates, tau = tau)
 
+    def muon_parameters(self):
+        return self.attn_layers.muon_parameters()
+
     def forward(
         self,
         x,