Add Batch dataloader support.

wit/meaning_dataset.py

@@ -7,6 +7,7 @@ from itertools import chain
 from typing import Dict, Tuple
 from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split
 import numpy as np
+from torch.utils.data import BatchSampler
 
 
 class MeaningMap:  # 16777216 1048576 8192
@@ -26,22 +27,10 @@ class MeaningMap:  # 16777216 1048576 8192
             self.ms_data = np.load(file_data)
             self.ms_start = np.load(file_start)
             self.ms_len = np.load(file_len)
-            return None
+        else:
-
             print("Disk cache miss, build new one.")
 
             mm = np.empty((size, max_subitem), dtype=np.int32)
-        # total_level = int(math.log(size / vocab_size, max_subitem))
-
-        # start = [0]
-        # end = [vocab_size]
-        # shift = vocab_size
-        # for i in range(total_level):
-        #     shift = end[-1]
-        #     start.append(end[-1])
-        #     end.append(shift * self.max_subitem)
-        # start.append(end[-1])
-        # end.append(size)
 
             index = np.arange(0, size)
             mm = np.random.random((size, max_subitem))
@@ -97,32 +86,44 @@ class MeaningMap:  # 16777216 1048576 8192
         len = self.ms_len[meaning]
         return self.ms_data[start : start + len]
 
+    def MaxLength(self):
+        return max(self.ms_len)
+
 
 class MeaningDataset(Dataset):
 
-    def __init__(self, start=131072, end=1048576, size=32768, vocab_size=4096, max_subitem=10, seed=42):
+    def __init__(
-        self.seed = seed
+        self,
+        start=131072,
+        end=1048576,
+        size=32768,
+        vocab_size=4096,
+        max_subitem=10,
+        min_seq_len=2,
+        seed=42,
+        data=None,
+        length=None,
+    ):
+        if data != None and length != None:
+            self.data = data
+            self.length = length
+            return
         np.random.seed(seed)
-        self.size = size
+        mm = MeaningMap(size=end, vocab_size=vocab_size, max_subitem=max_subitem)  # 1048576
-        self.mm = MeaningMap(size=end, vocab_size=vocab_size, max_subitem=max_subitem)  # 1048576
         self.data = []
+        self.length = []
         meanings = np.random.randint(start, end, size=(size))
         for m in meanings:
-            sq = self.mm.GetSequence(m)
+            sq = mm.GetSequence(m)
-            if len(sq) > 1:
+            if len(sq) >= min_seq_len:
                 self.data.append(sq)
-        left = size - len(self.data)
+                self.length.append(len(sq))
-        while True:
-            if left <= 0:
-                break
-            index = np.random.randint(start, end)
-            sq = self.mm.GetSequence(index)
-            if len(sq) > 1:
-                self.data.append(sq)
-                left = left - 1
 
     def __len__(self):
-        return self.size
+        return len(self.data)
+
+    def len(self):
+        return len(self.data)
 
     def __getitem__(self, idx):
         output = {}
@@ -132,11 +133,93 @@ class MeaningDataset(Dataset):
         output["token_type_ids"] = torch.zeros(data.shape)
         return output
 
+    def GetBatch(self, index_list):
+        data = []
+        for i in index_list:
+            data.append(self.data[i])
+        output = {}
+        data = torch.tensor(data).long()
+        output["input_ids"] = data
+        output["labels"] = data.clone()
+        output["token_type_ids"] = torch.zeros(data.shape)
+        return output
+
+    def Split(self, ratio):
+        l = len(self.data)
+        middle = int(l * ratio)
+        d_shuffle = self.data.copy()
+        l_shuffle = self.length.copy()
+        md1 = MeaningDataset(data=d_shuffle[:middle], length=l_shuffle[:middle])
+        md2 = MeaningDataset(data=d_shuffle[middle:], length=l_shuffle[middle:])
+        return md1, md2
+
+
+class BatchGroupMeaningDataloader(Dataset):
+
+    def __init__(self, dataset: MeaningDataset, batch_size, shuffle=True, drop_last=True):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+
+        length = dataset.length
+        unique, counts = np.unique(length, return_counts=True)
+        gl = {}
+        for u in unique:
+            gl[u] = np.where(length == u)[0]
+
+        lens = list(gl.keys())
+        gs = {}
+        if shuffle:
+            for k in gl.keys():
+                sl = gl[k].copy()
+                np.random.shuffle(sl)
+                gs[k] = sl
+        else:
+            for k in gl.keys():
+                sl = gl[k].copy()
+                gs[k] = sl
+
+        index = np.zeros((0, batch_size), dtype=np.int64)
+        for l in lens:
+            batch = len(gs[l]) // batch_size
+            new = gs[l][0 : batch * batch_size].reshape(batch, batch_size)
+            index = np.concatenate((index, new), axis=0)
+        if shuffle:
+            index_shuffle = np.arange(0, index.shape[0])
+            np.random.shuffle(index_shuffle)
+            index = index[index_shuffle]
+        self.index = index
+
+    def __len__(self):
+        return len(self.index)
+
+    def __getitem__(self, idx):
+        # print("get idx" + str(idx))
+        return self.dataset.GetBatch(self.index[idx])
+
 
 if __name__ == "__main__":
 
-    md = MeaningDataset(4096, 4100, size=32768)
+    md = MeaningDataset(4096, 8100, size=1024)
-    it = iter(md)
+    train, val = md.Split(0.95)
+
+    dl = BatchGroupMeaningDataloader(train, 2)
+    it = iter(dl)
+    ne1 = next(it)
+    ne2 = next(it)
+    ne3 = next(it)
+
+    dl = DataLoader(
+        train,
+        num_workers=1,
+        persistent_workers=True,
+        shuffle=False,
+    )
+    it = iter(dl)
+    ne1 = next(it)
+    ne2 = next(it)
+    ne3 = next(it)
+
     for i in range(10):
         daf = next(it)["input_ids"].numpy().tolist()
 

wit/train.py

@@ -3,25 +3,22 @@ from functools import partial
 from itertools import chain
 from typing import Dict, Tuple
 
-import datasets
 import pytorch_lightning as pl
 import torch
-from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split, Subset
 
 from lit_module import LitModule
 from tokenization_qwen import QWenTokenizer
 from logger import TBLogger
 
-from special_dataset import SpecialDataset
+from meaning_dataset import MeaningDataset, BatchGroupMeaningDataloader
-from meaning_dataset import MeaningDataset
 from wit.configuration import ModelConfig
 
 pretrain_model_name = None  # "qwen/Qwen-1_8B-Chat"
 learning_rate = 0.0001
 use_tril_attention_mask = None
 precision = "32-true"  # "precision:bf16-mixed,16-mixed,32-true"
-train_batch_size = 1
+train_batch_size = 32
-val_batch_size = 1
+val_batch_size = 32
 num_proc = 8
 max_epochs = 1000
 strategy = "auto"
@@ -42,38 +39,19 @@ if __name__ == "__main__":
     lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
     tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
 
-    # raw_dataset = SpecialDataset()
+    level_ratio = 4
-
+    start = vocab_size * level_ratio * level_ratio
-    level_scale = 4
+    end = start * level_ratio
-    start = vocab_size * level_scale * level_scale
+    size = end * level_ratio
-    raw_dataset = MeaningDataset(
+    raw_dataset = MeaningDataset(start, end, size, vocab_size, level_ratio)
-        start=start,
+    train_dataset, val_dataset = raw_dataset.Split(0.95)
-        end=start * level_scale,
+    train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
-        size=start * level_scale * level_scale,
+    val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
-        max_subitem=level_scale,
+    it = iter(train_dataloader)
-        vocab_size=vocab_size,
-    )
-
-    train_dataset, val_dataset = random_split(raw_dataset, [0.95, 0.05])
-    it = iter(train_dataset)
     print("data samples:")
     for i in range(10):
         print(next(it)["input_ids"].numpy().tolist())
 
-    train_dataloader = DataLoader(
-        train_dataset,
-        batch_size=train_batch_size,
-        num_workers=num_proc,
-        persistent_workers=True,
-        shuffle=True,
-    )
-    val_dataloader = DataLoader(
-        val_dataset,
-        batch_size=val_batch_size,
-        num_workers=num_proc,
-        persistent_workers=True,
-    )
-
     torch.set_float32_matmul_precision("medium")
     lit_trainer = pl.Trainer(
         accelerator="gpu",

wit/train_special.py

@@ -0,0 +1,79 @@
+import argparse
+from functools import partial
+from itertools import chain
+from typing import Dict, Tuple
+
+import datasets
+import pytorch_lightning as pl
+import torch
+from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split, Subset
+
+from lit_module import LitModule
+from tokenization_qwen import QWenTokenizer
+from logger import TBLogger
+
+from special_dataset import SpecialDataset
+from meaning_dataset import MeaningDataset
+from wit.configuration import ModelConfig
+
+pretrain_model_name = None  # "qwen/Qwen-1_8B-Chat"
+learning_rate = 0.0001
+use_tril_attention_mask = None
+precision = "32-true"  # "precision:bf16-mixed,16-mixed,32-true"
+train_batch_size = 128
+val_batch_size = 128
+num_proc = 8
+max_epochs = 1000
+strategy = "auto"
+resume_from_ckpt_path = None
+seed = 42
+vocab_size = 256
+
+
+if __name__ == "__main__":
+    torch.manual_seed(seed)
+
+    config = ModelConfig()
+    config.vocab_size = vocab_size
+    config.hidden_size = 128  # 128 1024 2048  32
+    config.num_hidden_layers = 3  # 6 12 24  3
+    config.num_attention_heads = 8  # 8 8 16
+
+    lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
+    tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
+
+    raw_dataset = SpecialDataset()
+    train_dataset, val_dataset = random_split(raw_dataset, [0.95, 0.05])
+    it = iter(train_dataset)
+    print("data samples:")
+    for i in range(10):
+        print(next(it)["input_ids"].numpy().tolist())
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_size=train_batch_size,
+        num_workers=num_proc,
+        persistent_workers=True,
+        shuffle=True,
+    )
+    val_dataloader = DataLoader(
+        val_dataset,
+        batch_size=val_batch_size,
+        num_workers=num_proc,
+        persistent_workers=True,
+    )
+
+    torch.set_float32_matmul_precision("medium")
+    lit_trainer = pl.Trainer(
+        accelerator="gpu",
+        precision=precision,
+        logger=TBLogger("./", default_hp_metric=False),
+        strategy=strategy,
+        max_epochs=max_epochs,
+    )
+    lit_trainer.fit(
+        lit_module,
+        train_dataloaders=train_dataloader,
+        val_dataloaders=val_dataloader,
+        ckpt_path=resume_from_ckpt_path,
+    )