From 0b19fd576abad3b10eedc4ec75a9acd098ed6527 Mon Sep 17 00:00:00 2001 From: Colin Date: Wed, 19 Feb 2025 19:39:59 +0800 Subject: [PATCH] Refine train code. --- wit/configuration.py | 27 +++++++++++++++++++++++---- wit/dataset/dataset.py | 18 +++++++++--------- wit/model/modeling_wit.py | 1 + wit/train.py | 23 +++++++++++++---------- 4 files changed, 46 insertions(+), 23 deletions(-) diff --git a/wit/configuration.py b/wit/configuration.py index dc7ba37..d0f63b3 100644 --- a/wit/configuration.py +++ b/wit/configuration.py @@ -1,5 +1,3 @@ - - class ModelConfig: def __init__(self): self.vocab_size = 4096 @@ -47,14 +45,16 @@ class MeaningDatasetConfig: self.mask_level = None self.mask_idx = None + class DatasetConfig: def __init__(self): self.name = "meaning" self.meaning = MeaningDatasetConfig() + class TrainConfig: def __init__(self): - self.name = "bigger" # current train process name + self.name = "bigger" # current train process name self.pretrain_model_name = None # "qwen/Qwen-1_8B-Chat" self.learning_rate = 0.0001 self.use_tril_attention_mask = None @@ -69,4 +69,23 @@ class TrainConfig: self.dataloader_works = 2 self.model_config = ModelConfig() - self.dataset = DatasetConfig() \ No newline at end of file + self.dataset = DatasetConfig() + + +def class_to_dict(obj): + if isinstance(obj, (int, float, str, bool, type(None))): + return obj + elif isinstance(obj, dict): + return {k: class_to_dict(v) for k, v in obj.items()} + elif isinstance(obj, list): + return {str(index): value for index, value in enumerate(obj)} + elif hasattr(obj, "__dict__"): + return {k: class_to_dict(v) for k, v in obj.__dict__.items()} + else: + return obj + + +# train_config = TrainConfig() +# train_config_dict = class_to_dict(train_config) +# import pprint +# pprint.pprint(train_config_dict) diff --git a/wit/dataset/dataset.py b/wit/dataset/dataset.py index 7431610..842567a 100644 --- a/wit/dataset/dataset.py +++ b/wit/dataset/dataset.py @@ -29,27 +29,27 @@ def InitDataset(config): return train_dataloader, val_dataloader if config.dataset.name == "meaning": - conf = config.dataset.meaning + c = config.dataset.meaning vocab = config.model_config.vocab_size - start = vocab * (conf.level_ratio**conf.level) - size = vocab * int((conf.level_ratio**conf.dataset_level)) + start = vocab * (c.level_ratio**c.level) + size = vocab * int((c.level_ratio**c.dataset_level)) path = "./data/" - trainfile = path + f"MeaningDataset_train_v{size}_s{start}_s{size}_lr{conf.level_ratio}_ms{conf.min_subitem}.pt" - valfile = path + f"MeaningDataset_val_v{size}_s{start}_s{size}_lr{conf.level_ratio}_ms{conf.min_subitem}.pt" + trainfile = path + f"MeaningDataset_train_v{size}_s{start}_s{size}_lr{c.level_ratio}_ms{c.min_subitem}.pt" + valfile = path + f"MeaningDataset_val_v{size}_s{start}_s{size}_lr{c.level_ratio}_ms{c.min_subitem}.pt" if not os.path.exists(path): os.mkdir(path) if os.path.exists(trainfile) and os.path.exists(valfile): print(f"INFO: Load dataset from {trainfile}") train_dataset = torch.load(trainfile, weights_only=False) - train_dataset.set_mask(conf.mask_level, conf.mask_idx) + train_dataset.set_mask(c.mask_level, c.mask_idx) print(f"INFO: Load dataset from {valfile}") val_dataset = torch.load(valfile, weights_only=False) - val_dataset.set_mask(conf.mask_level, conf.mask_idx) + val_dataset.set_mask(c.mask_level, c.mask_idx) print(f"INFO: Load dataset end") else: - raw_dataset = MeaningDataset(start, start + size, vocab, None, conf.level_ratio, conf.min_subitem) - raw_dataset.set_mask(conf.mask_level, conf.mask_idx) + raw_dataset = MeaningDataset(start, start + size, vocab, None, c.level_ratio, c.min_subitem) + raw_dataset.set_mask(c.mask_level, c.mask_idx) train_dataset, val_dataset = raw_dataset.split(0.9) torch.save(train_dataset, trainfile) torch.save(val_dataset, valfile) diff --git a/wit/model/modeling_wit.py b/wit/model/modeling_wit.py index 6281710..94ec99f 100644 --- a/wit/model/modeling_wit.py +++ b/wit/model/modeling_wit.py @@ -193,6 +193,7 @@ class QWenLMHeadModel(nn.Module): class QwenRunner: def __init__(self, qwen): self.qwen = qwen + # torch.backends.cuda.enable_flash_sdp(True) @torch.no_grad() def Chat( diff --git a/wit/train.py b/wit/train.py index bac1e3a..db1e249 100644 --- a/wit/train.py +++ b/wit/train.py @@ -17,41 +17,44 @@ if __name__ == "__main__": conf.pretrain_model_name = None # "qwen/Qwen-1_8B-Chat" conf.learning_rate = 0.0001 conf.use_tril_attention_mask = None - conf.precision = "32-true" # "precision:bf16-mixed,16-mixed,32-true" - conf.train_batch_size = 8 + conf.precision = "bf16-mixed" # "precision:bf16-mixed,16-mixed,32-true" + conf.train_batch_size = 16 conf.val_batch_size = 4 conf.num_proc = 8 conf.max_epochs = 1000 conf.strategy = "auto" conf.resume_from_ckpt_path = None conf.seed = 42 - conf.dataloader_works = 2 + conf.dataloader_works = 4 + + conf.dataset.meaning.mask_level = [0, 1, 2] + conf.dataset.meaning.mask_idx = [0, -1, -1] - conf.mask_level = None # [0, 1, 2] - conf.mask_idx = None # [0, 0, -1] - config.vocab_size = 256 config.hidden_size = 128 # 128 1024 2048 32 - config.num_hidden_layers = 6 # 6 12 24 3 + config.num_hidden_layers = 3 # 6 12 24 3 config.num_attention_heads = 16 # 8 8 16 torch.manual_seed(conf.seed) - lit_module = LitModule(conf.pretrain_model_name, conf.learning_rate, config, conf.use_tril_attention_mask) - tokenizer = QWenTokenizer("./model/wit_b64.tiktoken", "./model/wit_char.tiktoken") + lit_module = LitModule(conf) train_dataloader, val_dataloader = ds.InitDataset(conf) # for i in range(len(train_dataloader)): # print(train_dataloader.print_mapping(i)) + logger = TBLogger("./log/", name=conf.name) + logger.log_hyperparams(configuration.class_to_dict(conf)) + torch.set_float32_matmul_precision("medium") lit_trainer = pl.Trainer( accelerator="cuda", precision=conf.precision, # logger=MLFLogger("./log/", run_name=conf.name), - logger=TBLogger("./log/", name=conf.name), + logger=logger, strategy=conf.strategy, max_epochs=conf.max_epochs, ) + lit_trainer.fit( lit_module, train_dataloaders=train_dataloader,