Added tokenise method for streamed data, fixed issues with einsums

main.py

@@ -1,4 +1,5 @@
 import argparse
+import time
 import torch as t
 import torch.nn as nn
 import torch.functional as F
@@ -9,7 +10,8 @@ import wandb
 from typing import Tuple
 from torch.utils.data.dataloader import DataLoader
 from datasets import load_dataset
-from utils import OsSoluConfig
+from transformers import AutoTokenizer
+from utils import OsSoluConfig, tokenise
 from model import OsSoluModel
 
 WANDB_PROJECT_NAME = "os_solu"
@@ -32,7 +34,7 @@ def parse_arguments() -> dict:
     parser.add_argument("--dropout", type=float, default=0.1, help="Probability of dropout.")
     parser.add_argument("--learning_rate", type=float, default=1e-3, help="Learning rate for the optimiser.")
     parser.add_argument("--ln_eps", type=float, default=1e-3, help="Layer norm epsilon.")
-    parser.add_argument("--max_positional_embeddings", type=int, default=1024, help="Maximum number of positional embeddings.")
+    parser.add_argument("--max_positional_embeddings", type=int, default=1024, help="Maximum number of positional embeddings/sequence length.")
     parser.add_argument("--nonlinearity", type=str, default="solu", help=" Nonlinearity to use inside MLP block: must be relu or solu.")
     parser.add_argument("--num_blocks", type=int, default=1, help="Number of transformer blocks.")
     parser.add_argument("--num_embeddings", type=int, default=1024, help="Number of embeddings.")
@@ -40,7 +42,7 @@ def parse_arguments() -> dict:
     parser.add_argument("--num_heads", type=int, default=4, help="Number of attention heads in each attention layer.")
     parser.add_argument("--optimiser_type", type=str, default="adam", help="Optimiser type.")
     parser.add_argument("--self_attention_type", type=str, default="unidirectional", help="What type of attention to use: rotary or unidirectional.")
-    parser.add_argument("--vocab_size", type=int, default=65536, help="Vocabulary size of the input sequence.")
+    parser.add_argument("--vocab_size", type=int, default=50_278, help="Vocabulary size of the input sequence.")
     args = vars(parser.parse_args())
 
     # Parse string arguments.
@@ -67,7 +69,6 @@ def train(config: OsSoluConfig, model: OsSoluModel, train_dataloader: DataLoader
     Returns:
         OsSoluModel: The trained model.
     """
-    # TODO: training loop
     train_loss_fn = t.nn.CrossEntropyLoss()
     wandb.watch(model, criterion=train_loss_fn, log="all", log_freq=10, log_graph=True)
 
@@ -77,16 +78,17 @@ def train(config: OsSoluConfig, model: OsSoluModel, train_dataloader: DataLoader
 
     # Train loop.
     examples_seen = 0
+    train_data_iterator = iter(train_dataloader)
     for epoch in range(config.num_epochs):
-        for i, (data, target) in enumerate(tqdm(train_dataloader)):
-            print(data, target)
+        for i, batch in enumerate(tqdm(train_data_iterator
+    )):
+            data = batch["text"]
             data = data.to(DEVICE)
-            target = target.to(DEVICE)
 
             predictions = model(data)
             accuracy = (predictions.argmax(dim=-1) == target).sum() / len(data)
             optimiser.zero_grad()
-            loss = train_loss_fn(target, predictions)
+            # loss = train_loss_fn(data, predictions)
             loss.backward()
             optimiser.step()
 
@@ -109,9 +111,10 @@ def eval(model: OsSoluModel, test_dataloader: DataLoader) -> None:
     total_loss, num_correct = 0, 0
     model.eval()
     with t.inference_mode():
-        for i, (data, target) in enumerate(tqdm(test_dataloader)):
+        test_data_iterator = iter(test_dataloader)
+        for i, (data, target) in enumerate(tqdm(test_data_iterator)):
+            data = batch["text"]
             data = data.to(DEVICE)
-            target = target.to(DEVICE)
 
             predictions = model(data)
             num_correct += (predictions.argmax(dim=-1) == target).sum().item()
@@ -134,15 +137,31 @@ def setup() -> Tuple[OsSoluConfig, OsSoluModel]:
     args = parse_arguments()
     wandb.init(project=WANDB_PROJECT_NAME, config=args)
     config = OsSoluConfig(args)
-    model = OsSoluModel(config)
+    model = OsSoluModel(config).to(DEVICE)
 
+    start_data_time = time.time()
     # Load and prep data.
     ds = load_dataset("the_pile", streaming=True)
-    train_dataset = ds["train"].with_format("torch")
-    train_dataloader = DataLoader(train_dataset, batch_size=config.batch_size)
 
-    test_dataset = ds["test"].with_format("torch")
+    try:
+        ds = ds.remove_columns("meta")
+    except:
+        print("Dataset did not contain 'meta' column.")
+
+    train_dataset = ds["train"]
+    test_dataset = ds["test"]
+
+    # TODO: tokenise the data before sending it to the model.
+    tokeniser = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
+    tokeniser.add_special_tokens({"pad_token": "<PAD>"})
+
+    train_dataset = train_dataset.map(lambda x: tokenise(x, tokeniser), batched=True).with_format("torch")
+    test_dataset = test_dataset.map(tokenise, batched=True).with_format("torch")
+
+    train_dataloader = DataLoader(train_dataset, batch_size=config.batch_size)
     test_dataloader = DataLoader(test_dataset, batch_size=config.batch_size)
+    print(f"Data loaded in {time.time() - start_data_time:.1f}s.")
+
     return config, model, (train_dataloader, test_dataloader)
 
 if __name__=="__main__":

model.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 import torch.functional as F
 import torch.optim as optim
 import wandb
-import fancy_einsum as einsum
+from fancy_einsum import einsum
 from einops import rearrange, repeat, reduce
 from utils import OsSoluConfig
 
@@ -22,7 +22,7 @@ class OsSoluModel(nn.Module):
         self.final_ln = nn.LayerNorm(config.d_model, config.ln_eps)
 
     def forward(self, x: t.Tensor) -> t.Tensor:
-        positional_embeddings = self.embed_positions(t.arange(x.size(1)))
+        positional_embeddings = self.embed_positions(t.arange(x.size(1), device=x.device))
         token_embeddings = self.embed_tokens(x)
         embeddings = positional_embeddings + token_embeddings
         out = self.dropout(embeddings)
@@ -69,9 +69,9 @@ class UnidirectionalAttention(nn.Module):
         super().__init__()
         self.num_heads = config.num_heads
         self.d_model = config.d_model
-        self.project_q = nn.Linear(config.num_embeddings, config.d_model)
-        self.project_k = nn.Linear(config.num_embeddings, config.d_model)
-        self.project_v = nn.Linear(config.num_embeddings, config.d_model)
+        self.project_q = nn.Linear(config.d_model, config.d_model)
+        self.project_k = nn.Linear(config.d_model, config.d_model)
+        self.project_v = nn.Linear(config.d_model, config.d_model)
         self.project_out = nn.Linear(config.d_model, config.d_model)
         self.LARGE_NEGATIVE_VALUE = -1e5
 
@@ -84,7 +84,11 @@ class UnidirectionalAttention(nn.Module):
 
         Q = self.hidden_to_heads(Q)
         K = self.hidden_to_heads(K)
-        attention_pattern = einsum("batch num_heads seqlen_q head_size, batch num_heads seqlen_k head_size -> batch num_heads seqlen_q seqlen_k")
+        attention_pattern = einsum(
+            "batch num_heads seqlen_q head_size, "
+            "batch num_heads seqlen_k head_size ->"
+            "batch num_heads seqlen_q seqlen_k",
+            Q, K)
 
         return attention_pattern
 
@@ -95,18 +99,23 @@ class UnidirectionalAttention(nn.Module):
         
         # Masking attention. Since GPT is unidirectional, it should only attend to previous tokens.
         if seqlen > 1:
-            fst_range = t.arange(seqlen, device=self.device).unsqueeze(0).T
-            snd_range = t.arange(seqlen, device=self.device).unsqueeze(0)
+            fst_range = t.arange(seqlen, device=x.device).unsqueeze(0).T
+            snd_range = t.arange(seqlen, device=x.device).unsqueeze(0)
             bool_array = fst_range < snd_range
-            attention_score[..., bool_array] = self.LARGE_NEGATIVE_VALUE
+            attention_pattern[..., bool_array] = self.LARGE_NEGATIVE_VALUE
         
         
         attention_pattern = attention_pattern / t.sqrt(t.tensor(self.d_model // self.num_heads))
         attention_score = attention_pattern.softmax(dim=-1)
         
         V = self.hidden_to_heads(V)
-        out = einsum("batch num_heads seqlen_q seqlen_k, batch num_heads seqlen_k head_size -> batch num_heads seqlen_q head_size", attention_score, V)
-        out = rearrange("b nh s hs -> b s (nh hs)")
+        out = einsum(
+            "batch num_heads seqlen_q seqlen_k,"
+            "batch num_heads seqlen_k head_size ->"
+            "batch num_heads seqlen_q head_size", 
+            attention_score, V)
+
+        out = rearrange(out, "b nh s hs -> b s (nh hs)")
         out = self.project_out(out)
         
 

requirements.txt

@@ -9,6 +9,7 @@ notebook
 numpy-stl
 plotly
 torch
+transformers
 tqdm
 wandb
 zstandard

utils.py

@@ -1,3 +1,6 @@
+import numpy as np
+from einops import rearrange
+
 class OsSoluConfig:
     """A class to hold hyperparameters for the model itself and for the training process."""
     
@@ -32,4 +35,42 @@ class OsSoluConfig:
         self.num_heads = args["num_heads"]
         self.optimiser_type = args["optimiser_type"]
         self.self_attention_type = args["self_attention_type"]
-        self.vocab_size = args["vocab_size"]
+        self.vocab_size = args["vocab_size"]
+
+def tokenise(batch, tokeniser, num_gpus: int = 1, context_length: int = 1024):
+    """Tokenise a batch of text data. This implementation is idiosyncratic to the Pile dataset, but can be easily modified to work with e.g. C4.
+
+    Args:
+        batch (dict): The batch of text, as a dict with a 'text' field.
+        tokeniser (-): A huggingface-API tokeniser, of type returned by AutoTokenizer.from_pretrained (depends on model chosen).
+        num_gpus (int, optional): The number of GPUs available for data parallel training. Defaults to 1.
+        context_length (int, optional): The context length of the model that will be trained on this data. Defaults to 1024.
+
+    Returns:
+        dict: A single field dictionary, 'text', whose value is a tensor of shape (batch_size, sequence_length) containing tokenised sequences.
+    """
+    batch = batch["text"]
+    full_text = tokeniser.eos_token.join(batch)
+
+    # Divide entire batch among all GPUs available.
+    seq_len = len(full_text)//num_gpus
+    sequence_list = [full_text[i*seq_len:(i+1)*seq_len] for i in range(num_gpus)]
+
+    # Tokenise sequences, removing padding tokens.
+    all_tokens = tokeniser(sequence_list, return_tensors="pt", padding=True)["input_ids"].flatten()
+    all_tokens = all_tokens[all_tokens != tokeniser.pad_token_id]
+
+    # Reshape all_tokens to be (batch_size x sequence_length) where each sequence has
+    # a "beginning of sequence" token prepended to it.
+    num_tokens = len(all_tokens)
+    current_batch_size = num_tokens // (context_length-1)
+    all_tokens = all_tokens[:(context_length-1)*current_batch_size]
+    all_tokens = rearrange(all_tokens, "(batch_size seq_len) -> batch_size seq_len", batch_size=current_batch_size, seq_len=context_length-1)
+    prefix = np.full((current_batch_size, 1), tokeniser.bos_token_id, dtype=np.int64)
+
+    tokenised_text = np.concatenate([prefix, all_tokens], axis=1)
+    assert  tokenised_text.shape == (current_batch_size, context_length)
+    print(f"{current_batch_size=}, {context_length=}")
+    return {"text": tokenised_text}
+
+