diff --git a/wit/meaning_dataset.py b/wit/meaning_dataset.py index 1e89819..40d8b87 100644 --- a/wit/meaning_dataset.py +++ b/wit/meaning_dataset.py @@ -31,6 +31,7 @@ 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) + print("Load end") else: print("Disk cache miss, build new one.") @@ -123,6 +124,11 @@ class MeaningDataset(Dataset): self.data.append(sq) self.length.append(len(sq)) + unique, counts = np.unique(self.length, return_counts=True) + print("MeaningDataset size: " + str(len(self.length))) + print("MeaningDataset max sequence length: " + str(max(unique))) + print("MeaningDataset most popular sequence length: " + str(unique[np.argmax(counts)])) + def __len__(self): return len(self.data) @@ -205,7 +211,8 @@ if __name__ == "__main__": md = MeaningDataset(4096, 8100, size=1024) train, val = md.Split(0.95) - dl = BatchGroupMeaningDataloader(train, 2) + dl = BatchGroupMeaningDataloader(train, 32) + length = len(dl) it = iter(dl) ne1 = next(it) ne2 = next(it) diff --git a/wit/train.py b/wit/train.py index 54eb5f6..79d958b 100644 --- a/wit/train.py +++ b/wit/train.py @@ -18,15 +18,16 @@ learning_rate = 0.0001 use_tril_attention_mask = None precision = "32-true" # "precision:bf16-mixed,16-mixed,32-true" train_batch_size = 32 -val_batch_size = 32 +val_batch_size = 4 num_proc = 8 max_epochs = 1000 strategy = "auto" resume_from_ckpt_path = None seed = 42 -vocab_size = 1024 +vocab_size = 2048 level_ratio = 4 +level = 4 hidden_size = 256 # 128 1024 2048 32 num_attention_heads = 8 # 8 8 16 @@ -57,9 +58,9 @@ if __name__ == "__main__": lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask) tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken") - start = vocab_size * level_ratio * level_ratio * level_ratio * level_ratio + start = vocab_size * (level_ratio**level) end = start * level_ratio - size = start + start + size = vocab_size * (level_ratio ** (level / 2)) raw_dataset = MeaningDataset(start, end, size, vocab_size, level_ratio) train_dataset, val_dataset = raw_dataset.Split(0.95) train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)