add training code

app.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torch.nn.utils import clip_grad_norm_
+
+import wandb
+from tqdm import tqdm
+from transformers import GPT2LMHeadModel
+from gated_state_spaces_pytorch import GatedStateSpacesLM
+from gated_state_spaces_pytorch.autoregressive_wrapper import AutoregressiveWrapper
+
+from c4x import C4X
+
+
+if __name__ == '__main__':
+    wandb.init(
+        project="gated-state-space", 
+        entity="naxalpha",
+    )
+    
+    gpt_2 = GPT2LMHeadModel.from_pretrained('gpt2-xl')
+    gpt_2.requires_grad_(False)
+    gpt_2 = gpt_2.cuda()
+
+    f_emb = 1600
+    model = AutoregressiveWrapper(
+        GatedStateSpacesLM(
+            num_tokens=50257,
+            dim=f_emb,
+            depth=24,
+        ),
+    )
+    wandb.watch(model)
+
+    emb = gpt_2.state_dict()['transformer.wte.weight']
+
+    model.net.token_emb.weight.requires_grad_(False)
+    model.net.token_emb.weight.copy_(emb)
+
+    model.net.to_logits.weight.requires_grad_(False)
+    model.net.to_logits.weight.copy_(emb)
+    
+    model.net.to_logits = nn.Sequential(
+        nn.LayerNorm(f_emb),
+        model.net.to_logits,
+    )
+
+    model = model.cuda()
+    optim = AdamW(model.parameters(), 2e-5)
+
+    bs = 8
+    kk = 128
+    dsx = C4X(kk+1)
+    dlx = DataLoader(
+        dsx,
+        batch_size=bs, 
+        num_workers=16,
+    )
+
+    k = 4
+    prog = tqdm(dlx)
+    optim.zero_grad()
+
+    for i, batch in enumerate(prog):
+        batch = batch.cuda()
+        if i % 2 == 0:  # distil
+            batch = batch[:, :-1]
+            with torch.no_grad():
+                logits = gpt_2(batch).logits
+                probs = logits.softmax(dim=-1)
+            out = model.net(batch)
+            los = F.cross_entropy(
+                out.flatten(0,1),
+                probs.flatten(0,1),
+            )
+        else:  # scratch
+            los = model(batch)
+
+        (los / k).backward()
+        if (i+1) % k == 0:
+            clip_grad_norm_(
+                model.parameters(),
+                max_norm=1.,
+            )
+            optim.step()
+            optim.zero_grad()
+
+        if i % 1000 == 0:
+            b, n = 4, 512
+            init = torch.tensor([[50256]]*b).cuda()
+            prd = model.generate(init, n)
+            prd = [dsx.decode(p) for p in prd]
+            try:
+                wandb.log(dict(
+                    text=wandb.Html(
+                        '<hr>'.join(
+                            p.replace('\n', '<br>') for p in prd
+                        )
+                )), step=i)
+            except Exception as ex:
+                print('Failed to log to W&B...', ex)
+            torch.save(model.state_dict(), 'model.pt')
+
+        wandb.log(dict(
+            loss=los.item(),
+        ), step=i)
+        prog.set_postfix(loss=los.item())

c4x.py

@@ -0,0 +1,62 @@
+# stream C4 dataset from Huggingface with GPT-2 Tokenizer for PyTorch Language Model Training
+import json
+import torch
+import random
+from datasets import load_dataset
+from transformers import GPT2Tokenizer
+from torch.utils.data import Dataset, get_worker_info
+
+
+def cycled(itr):
+    while True:
+        for itm in itr:
+            yield itm
+
+class C4X(Dataset):
+
+    def __init__(self, seq_len=512, split='train'):
+        self.seq = seq_len
+        self.ds = load_dataset(
+            'c4', 
+            name='en', 
+            split=split, 
+            streaming=True,
+        )
+        self.tok = GPT2Tokenizer.from_pretrained('gpt2')
+        self.init = False
+
+    def __len__(self):
+        return 1_000_000_000
+    
+    def _init(self):
+        if self.init:
+            return
+        wi = get_worker_info()
+        self.ds = cycled(
+            self.ds.shuffle(
+                seed=wi.seed,
+                buffer_size=10_000,
+            )
+        )
+        self.init = True
+
+    def _get_next(self):
+        self._init()
+        obj = next(self.ds)['text']
+        tkn = self.tok.encode(obj)
+        return tkn
+
+    def _get_full(self):
+        obj = []
+        while len(obj) < self.seq:
+            obj += self._get_next()
+            obj.append(self.tok.eos_token_id)
+        s = random.randint(0, len(obj)-self.seq)
+        return obj[s:s+self.seq]
+
+    def __getitem__(self, _):
+        return torch.tensor(self._get_full())
+
+    def decode(self, tkns):
+        return self.tok.decode(tkns)
+