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, )