init

.gitignore

@@ -0,0 +1,3 @@
+wandb
+__pycache__
+.ipynb_checkpoints

app.py

@@ -0,0 +1,139 @@
+# pip install accelerate datasets transformers huggingface_hub wandb gated_state_spaces_pytorch
+import os
+
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
+
+import wandb
+from tqdm import tqdm
+from transformers import BloomForCausalLM, BloomTokenizerFast
+from gated_state_spaces_pytorch import GatedStateSpacesLM
+from gated_state_spaces_pytorch.autoregressive_wrapper import AutoregressiveWrapper
+
+# from c4x import C4X
+from pile_hf import ThePile, ThePileTokenized
+from accelerate import Accelerator
+
+
+def main():
+    accelerator = Accelerator(
+        log_with="wandb",
+        gradient_accumulation_steps=8192,
+    )
+    accelerator.init_trackers("gated-state-space")
+
+    emb_fn = "emb.pt"
+    model_name = "bigscience/bloomz-1b7"
+    if not os.path.isfile(emb_fn):
+        bloom = BloomForCausalLM.from_pretrained(model_name)
+        wte = bloom.transformer.word_embeddings.state_dict()
+        torch.save(wte, emb_fn)
+    else:
+        wte = torch.load(emb_fn)
+
+    f_emb = 2048
+    n_vocab = 250880
+    model = AutoregressiveWrapper(
+        GatedStateSpacesLM(
+            num_tokens=n_vocab,
+            dim=f_emb,
+            depth=24,
+        ),
+    )
+
+    model.net.token_emb.requires_grad_(False)
+    model.net.token_emb.load_state_dict(wte)
+
+    to_logits = nn.Linear(f_emb, n_vocab, bias=False)
+    to_logits.requires_grad_(False)
+    to_logits.load_state_dict(wte)
+
+    model.net.to_logits = nn.Sequential(
+        nn.LayerNorm(f_emb),
+        to_logits,
+    )
+    model.load_state_dict(torch.load("model3.pt"))
+    model = model.to(accelerator.device)
+
+    if accelerator.is_main_process:
+        wandb.watch(model)
+
+    optim = AdamW(model.parameters(), 1e-4)
+    sch = CosineAnnealingWarmRestarts(
+        optim,
+        T_0=1000,
+        T_mult=2,
+        eta_min=1e-7,
+    )
+
+    bs = 1
+    kk = 2048
+    tok: BloomTokenizerFast = BloomTokenizerFast.from_pretrained(model_name)
+    dsx = ThePileTokenized(
+        ThePile("train"),
+        tokenizer=tok,
+        max_length=kk,
+        repeat_factor=4 / 3,
+    )
+    dlx = DataLoader(
+        dsx,
+        batch_size=bs,
+        num_workers=12,
+    )
+
+    prog = tqdm(dlx, disable=not accelerator.is_main_process)
+
+    model = accelerator.prepare(model)
+    optim, dlx, sch = accelerator.prepare(optim, dlx, sch)
+
+    optim.zero_grad()
+    for i, batch in enumerate(prog):
+        batch = batch.to(accelerator.device)
+        with accelerator.accumulate(model):
+            with accelerator.autocast():
+                los = model(batch)
+            accelerator.backward(los)
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(model.parameters(), 1.0)
+            optim.step()
+            optim.zero_grad()
+            if not accelerator.optimizer_step_was_skipped:
+                sch.step()
+
+        if i % 1000 == 0:
+            unwrapped_model = accelerator.unwrap_model(model)
+            b, n = 1, 512
+            init = torch.tensor([[2]] * b).to(accelerator.device)
+            prd = unwrapped_model.generate(init, n)
+            prd = [tok.decode(p) for p in prd]
+            try:
+                accelerator.log(
+                    dict(
+                        text=wandb.Html(
+                            "<hr>".join(p.replace("\n", "<br>") for p in prd)
+                        )
+                    ),
+                    step=i,
+                )
+            except Exception as ex:
+                accelerator.print("Failed to log to W&B...", ex)
+            sd = unwrapped_model.state_dict()
+            # sd.pop('net.to_logits.weight')
+            accelerator.save(sd, "model4.pt")
+
+        if i % 10 == 0:
+            accelerator.log(
+                dict(
+                    loss=los.item(),
+                    lr=optim.param_groups[0]["lr"],
+                ),
+                step=i,
+            )
+            prog.set_postfix(loss=los.item())
+
+
+if __name__ == "__main__":
+    main()

c4x.py

@@ -0,0 +1,61 @@
+# 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 BloomTokenizerFast
+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 = BloomTokenizerFast.from_pretrained('bigscience/bloomz-1b7')
+        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)

model2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:176c772feff0cf8504a46f872f6a32ae4269632b3e805e9437438f29268b795b
+size 7609367025

model3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c89f900da2bae9f79193ba785df8be4118d99135ffe66848e60f1ee6627b4bac
+size 7609367025

model4.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c89f900da2bae9f79193ba785df8be4118d99135ffe66848e60f1ee6627b4bac
+size 7609367025

pile.py

@@ -0,0 +1,107 @@
+import json
+import time
+import random
+from typing import Literal
+
+import requests
+import zstandard as zstd
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+Subset = Literal["train", "val", "test"]
+URLs = {
+    "val": [
+        "https://the-eye.eu/public/AI/pile/val.jsonl.zst",
+    ],
+    "test": [
+        "https://the-eye.eu/public/AI/pile/test.jsonl.zst",
+    ],
+    "train": [
+        "https://the-eye.eu/public/AI/pile/train/00.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/01.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/02.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/03.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/04.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/05.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/06.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/07.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/08.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/09.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/10.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/11.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/12.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/13.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/14.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/15.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/16.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/17.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/18.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/19.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/20.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/21.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/22.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/23.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/24.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/25.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/26.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/27.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/28.jsonl.zst",
+        "https://the-eye.eu/public/AI/pile/train/29.jsonl.zst",
+    ],
+}
+
+
+def _read_line_from_stream(reader, initial_line="", buffer_size=4096):
+    line = initial_line
+    while True:
+        c = reader.read(buffer_size)
+        if not c:
+            raise StopIteration
+        line += c.decode("utf-8")
+        if "\n" in line:
+            break
+    return line.split("\n", 1)
+
+
+def _line_streamer(reader, buffer_size=4096):
+    rest = ""
+    while True:
+        try:
+            line, rest = _read_line_from_stream(
+                reader,
+                rest,
+                buffer_size,
+            )
+            yield line
+        except StopIteration:
+            break
+
+
+class ThePile(IterableDataset):
+    TEXT_BUFFER_SIZE = 4096
+
+    def __init__(self, subset: Subset):
+        self.subset = subset
+
+    def __iter__(self):
+        urls = URLs[self.subset].copy()
+        while True:
+            wi = get_worker_info()
+            seed = wi.id if wi is not None else None
+            rnd = random.Random(seed)
+            rnd.shuffle(urls)
+            for url in urls:
+                r = requests.get(url, stream=True)
+                with zstd.ZstdDecompressor().stream_reader(r.raw) as reader:
+                    for line in _line_streamer(reader, self.TEXT_BUFFER_SIZE):
+                        data = json.loads(line)
+                        yield data
+
+
+if __name__ == "__main__":
+    from tqdm import tqdm
+
+    dataset = ThePile("train")
+    for data in tqdm(dataset, smoothing=0.01):
+        pass
+    # Average: ~2000 samples/sec/worker

pile_hf.py

@@ -0,0 +1,50 @@
+import torch
+from torch.utils.data import IterableDataset
+
+from transformers import PreTrainedTokenizerBase
+
+from pile import ThePile
+
+
+class ThePileTokenized(IterableDataset):
+    def __init__(
+        self,
+        base_dataset: ThePile,
+        tokenizer: PreTrainedTokenizerBase,
+        max_length: int = 1024,
+        repeat_factor: float = 1.0,
+    ):
+        self.pile = base_dataset
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.repeat_factor = repeat_factor
+
+    def __iter__(self):
+        ds = iter(self.pile)
+        buffer = []
+        while True:
+            tokens = self.tokenizer.encode(next(ds)["text"])
+            buffer += [self.tokenizer.eos_token_id] + tokens
+            while len(buffer) > self.max_length:
+                yield torch.tensor(buffer[: self.max_length])
+                buffer = buffer[int(self.max_length / self.repeat_factor) :]
+
+
+if __name__ == "__main__":
+    from tqdm import tqdm
+    from torch.utils.data import DataLoader
+    from transformers import GPT2Tokenizer
+
+    dataset = ThePileTokenized(
+        ThePile("train"),
+        GPT2Tokenizer.from_pretrained("gpt2"),
+        max_length=2048,
+        repeat_factor=4 / 3,
+    )
+    dataloader = DataLoader(
+        dataset,
+        batch_size=1,
+    )
+    for batch in tqdm(dataloader, smoothing=0.01):
+        x = 0
+    # ~6 iters/s for 1 worker