Witllm/wit/train.py

86 lines
2.7 KiB
Python
Raw Normal View History

2024-02-25 20:20:32 +08:00
import argparse
from functools import partial
from itertools import chain
from typing import Dict, Tuple
import pytorch_lightning as pl
import torch
2024-03-04 21:41:46 +08:00
2024-02-25 20:20:32 +08:00
from lit_module import LitModule
from tokenization_qwen import QWenTokenizer
2024-03-05 22:09:28 +08:00
from logger import TBLogger
2024-02-25 20:20:32 +08:00
2024-03-18 11:43:41 +08:00
from meaning_dataset import MeaningDataset, BatchGroupMeaningDataloader
2024-03-14 11:40:26 +08:00
from wit.configuration import ModelConfig
2024-03-13 19:41:02 +08:00
2024-03-14 11:40:26 +08:00
pretrain_model_name = None # "qwen/Qwen-1_8B-Chat"
2024-02-25 20:20:32 +08:00
learning_rate = 0.0001
use_tril_attention_mask = None
precision = "32-true" # "precision:bf16-mixed,16-mixed,32-true"
2024-04-17 20:22:49 +08:00
train_batch_size = 1
2024-03-29 22:10:25 +08:00
val_batch_size = 1
2024-02-25 20:20:32 +08:00
num_proc = 8
max_epochs = 1000
2024-03-05 22:09:28 +08:00
strategy = "auto"
2024-02-25 20:20:32 +08:00
resume_from_ckpt_path = None
seed = 42
2024-04-07 00:25:21 +08:00
vocab_size = 256
level_ratio = 6
level = 4
2024-04-17 20:22:49 +08:00
dataset_level = 1.5
min_subitem = 2
2024-03-26 15:01:19 +08:00
2024-04-07 00:25:21 +08:00
hidden_size = 1024 # 128 1024 2048 32
2024-04-03 17:09:30 +08:00
num_attention_heads = 16 # 8 8 16
2024-04-17 20:22:49 +08:00
num_hidden_layers = 6 # 6 12 24 3
2024-03-26 15:01:19 +08:00
2024-04-17 20:22:49 +08:00
mask_level = [0]
mask_idx = [-1]
2024-04-10 14:43:16 +08:00
# name = "vocab_ratio_level_data_hidden_head_layer"
2024-04-12 20:04:04 +08:00
# name = "mask_level_idx"
2024-04-17 20:22:49 +08:00
name = "small"
2024-04-12 20:04:04 +08:00
2024-04-17 20:22:49 +08:00
ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{min_subitem}" + "_" + f"{dataset_level}"
2024-03-29 22:10:25 +08:00
ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}"
2024-04-10 14:43:16 +08:00
ver = ver + "_" + f"{mask_level}" + "_" + f"{mask_idx}"
2024-02-25 20:20:32 +08:00
if __name__ == "__main__":
2024-03-15 11:16:13 +08:00
torch.manual_seed(seed)
2024-02-25 20:20:32 +08:00
2024-03-14 11:40:26 +08:00
config = ModelConfig()
config.vocab_size = vocab_size
2024-03-26 15:01:19 +08:00
config.hidden_size = hidden_size
config.num_hidden_layers = num_hidden_layers
config.num_attention_heads = num_attention_heads
2024-03-14 11:40:26 +08:00
lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
2024-02-25 20:20:32 +08:00
tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
2024-03-26 18:15:55 +08:00
start = vocab_size * (level_ratio**level)
2024-04-17 20:22:49 +08:00
size = vocab_size * int((level_ratio**dataset_level))
raw_dataset = MeaningDataset(start, start + size, size, vocab_size, level_ratio, min_subitem)
2024-04-10 14:43:16 +08:00
raw_dataset.set_mask(mask_level, mask_idx)
2024-04-02 19:59:05 +08:00
train_dataset, val_dataset = raw_dataset.split(0.9)
2024-03-18 11:43:41 +08:00
train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
2024-04-03 17:09:30 +08:00
# for i in range(len(train_dataloader)):
# print(train_dataloader.print_mapping(i))
2024-02-26 00:31:47 +08:00
2024-02-25 20:20:32 +08:00
torch.set_float32_matmul_precision("medium")
2024-03-05 22:09:28 +08:00
lit_trainer = pl.Trainer(
2024-03-25 19:53:11 +08:00
accelerator="cuda",
2024-03-05 22:09:28 +08:00
precision=precision,
2024-03-29 22:10:25 +08:00
logger=TBLogger("./log/", name=name, version=ver, default_hp_metric=False),
2024-03-05 22:09:28 +08:00
strategy=strategy,
max_epochs=max_epochs,
)
2024-02-25 20:20:32 +08:00
lit_trainer.fit(
lit_module,
train_dataloaders=train_dataloader,
val_dataloaders=val_dataloader,
ckpt_path=resume_from_ckpt_path,
)