Witllm/wit/train.py

70 lines
2.1 KiB
Python

import argparse
from functools import partial
from itertools import chain
from typing import Dict, Tuple
import pytorch_lightning as pl
import torch
from lit_module import LitModule
from tokenization_qwen import QWenTokenizer
from logger import TBLogger
from meaning_dataset import MeaningDataset, BatchGroupMeaningDataloader
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 = 16
val_batch_size = 32
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 = 1024 # 128 1024 2048 32
config.num_hidden_layers = 12 # 6 12 24 3
config.num_attention_heads = 16 # 8 8 16
lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
level_ratio = 6
start = vocab_size * level_ratio * level_ratio
end = start * level_ratio
size = end * level_ratio
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)
val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
# it = iter(train_dataloader)
# print("data samples:")
# for i in range(10):
# print(next(it)["input_ids"].numpy().tolist())
torch.set_float32_matmul_precision("medium")
lit_trainer = pl.Trainer(
accelerator="cuda",
devices=[0, 1],
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,
)