Refine train code.

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()
+        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)

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)

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(

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,