ready palm to be fine-tuneable

Author: lucidrains

data/README.md

@@ -0,0 +1,3 @@
+# Data source
+
+The enwik8 data was downloaded from the Hutter prize page: http://prize.hutter1.net/

data/enwik8.gz

@@ -0,0 +1 @@
+from palm_rlhf_pytorch.palm_rlhf_pytorch import PaLM

palm_rlhf_pytorch/__init__.py

@@ -4,6 +4,8 @@
 
 from einops import rearrange
 
+from palm_rlhf_pytorch.utils import eval_decorator, top_p, top_k
+
 # normalization
 # they use layernorm without bias, something that pytorch does not offer
 
@@ -27,10 +29,11 @@ def __init__(self, fn):
 
     def forward(self, x):
         y = self.fn(x)
+
         if not y.requires_grad and not x.requires_grad:
             return x.add_(y)
-        return y + x
 
+        return y + x
 
 # rotary positional embedding
 # https://arxiv.org/abs/2104.09864
@@ -57,9 +60,6 @@ def apply_rotary_pos_emb(pos, t):
     return (t * pos.cos()) + (rotate_half(t) * pos.sin())
 
 
-def l2norm(t):
-    return F.normalize(t, dim = -1)
-
 # classic Noam Shazeer paper, except here they use SwiGLU instead of the more popular GEGLU for gating the feedforward
 # https://arxiv.org/abs/2002.05202
 
@@ -133,7 +133,7 @@ def forward(self, x):
 
         # attention queries, keys, values, and feedforward inner
 
-        q, kv, v, ff = self.fused_attn_ff_proj(x).split(self.fused_dims, dim=-1)
+        q, k, v, ff = self.fused_attn_ff_proj(x).split(self.fused_dims, dim=-1)
 
 
         # split heads
@@ -178,19 +178,75 @@ def forward(self, x):
 # transformer
 
 
-def PaLM(*, dim, num_tokens, depth, dim_head=64, heads=8, ff_mult=4):
-    net = nn.Sequential(
-        nn.Embedding(num_tokens, dim),
-        *[
-            Residual(ParallelTransformerBlock(dim=dim, dim_head=dim_head, heads=heads, ff_mult=ff_mult))
-            for _ in range(depth)
-        ],
-        LayerNorm(dim),
-        nn.Linear(dim, num_tokens, bias=False)
-    )
+class PaLM(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        num_tokens,
+        depth,
+        dim_head=64,
+        heads=8,
+        ff_mult=4
+    ):
+        super().__init__()
+
+        self.token_emb = nn.Embedding(num_tokens, dim)
+        self.layers = nn.ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(Residual(ParallelTransformerBlock(dim=dim, dim_head=dim_head, heads=heads, ff_mult=ff_mult)))
+
+        self.to_logits = nn.Sequential(
+            LayerNorm(dim),
+            nn.Linear(dim, num_tokens, bias=False)
+        )        
+        
+        self.to_logits[-1].weight = self.token_emb.weight
+
+        nn.init.normal_(self.token_emb.weight, std=0.02)
+
+    @eval_decorator
+    @torch.no_grad()
+    def generate(
+        self,
+        prime,
+        seq_len,
+        temperature = 1.,
+        filter_logits_fn = top_k,
+        filter_thres = 0.9,
+        **kwargs
+    ):
+        n, out = prime.shape[-1], prime.clone()
+
+        for _ in range(seq_len):
+            logits = self.forward(out, **kwargs)
+
+            filtered_logits = filter_logits_fn(logits[:, -1], thres = filter_thres)
+            probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+            sample = torch.multinomial(probs, 1)
+            out = torch.cat((out, sample), dim=-1)
+
+        return out[:, n:]
+
+    def forward(
+        self,
+        x,
+        return_loss = False
+    ):
+        if return_loss:
+            x, labels = x[:, :-1], x[:, 1:]
+
+        x = self.token_emb(x)
+
+        for layer in self.layers:
+            x = layer(x) + x
+
+        logits = self.to_logits(x)
 
-    # they used embedding weight tied projection out to logits, not common, but works
-    net[-1].weight = net[0].weight
+        if not return_loss:
+            return logits
 
-    nn.init.normal_(net[0].weight, std=0.02)
-    return net
+        logits = rearrange(logits, 'b n c -> b c n')
+        return F.cross_entropy(logits, labels)

palm_rlhf_pytorch/palm_rlhf_pytorch.py

@@ -0,0 +1,35 @@
+import math
+import torch
+from torch import einsum, nn
+import torch.nn.functional as F
+
+# decorators
+
+def eval_decorator(fn):
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+    return inner
+
+# sampling helpers
+
+def top_p(logits, thres = 0.9):
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+    sorted_indices_to_remove = cum_probs > (1 - thres)
+    sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
+    sorted_indices_to_remove[:, 0] = 0
+
+    sorted_logits[sorted_indices_to_remove] = float('-inf')
+    return sorted_logits.scatter(1, sorted_indices, sorted_logits)
+
+def top_k(logits, thres = 0.9):
+    k = math.ceil((1 - thres) * logits.shape[-1])
+    val, ind = torch.topk(logits, k)
+    probs = torch.full_like(logits, float('-inf'))
+    probs.scatter_(1, ind, val)
+    return probs

palm_rlhf_pytorch/utils.py

@@ -0,0 +1,106 @@
+import gzip
+import random
+import tqdm
+import numpy as np
+
+import torch
+from torch.optim import Adam
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+
+from palm_rlhf_pytorch import PaLM
+
+# constants
+
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 2e-4
+VALIDATE_EVERY = 100
+PRIME_LENGTH = 128
+GENERATE_EVERY = 500
+GENERATE_LENGTH = 512
+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 = PaLM(
+    num_tokens=256,
+    dim=512,
+    depth=8
+).cuda()
+
+# prepare enwik8 data
+
+with gzip.open("./data/enwik8.gz") as file:
+    X = np.fromstring(file.read(int(95e6)), dtype=np.uint8)
+    trX, vaX = np.split(X, [int(90e6)])
+    data_train, data_val = torch.from_numpy(trX), torch.from_numpy(vaX)
+
+
+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,))
+        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))
+val_loader = cycle(DataLoader(val_dataset, batch_size=BATCH_SIZE))
+
+# optimizer
+
+optim = Adam(model.parameters(), lr=LEARNING_RATE)
+
+# training
+
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10.0, desc="training"):
+    model.train()
+
+    for __ in range(GRADIENT_ACCUMULATE_EVERY):
+        loss = model(next(train_loader), return_loss = True)
+        loss.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 = model(next(val_loader), return_loss = True)
+            print(f"validation loss: {loss.item()}")
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:PRIME_LENGTH]
+        prime = decode_tokens(inp)
+        print(f"%s \n\n %s", (prime, "*" * 100))
+
+        sample = model.generate(inp[None, ...], GENERATE_LENGTH)
+        output_str = decode_tokens(sample[0])
+        print(output_str)