colin
/
Witllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refine train.py for train.

This commit is contained in:
Colin 2024-03-25 19:53:11 +08:00
parent 4c7fdbe817
commit d10e7a8396
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
wit/train.py
View File

@ -17,7 +17,7 @@ pretrain_model_name = None # "qwen/Qwen-1_8B-Chat"
learning_rate = 0.0001 learning_rate = 0.0001
use_tril_attention_mask = None use_tril_attention_mask = None
precision = "32-true" # "precision:bf16-mixed,16-mixed,32-true" precision = "32-true" # "precision:bf16-mixed,16-mixed,32-true"
train_batch_size = 32 train_batch_size = 16
val_batch_size = 32 val_batch_size = 32
num_proc = 8 num_proc = 8
max_epochs = 1000 max_epochs = 1000
@ -39,7 +39,7 @@ if __name__ == "__main__":
lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask) lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken") tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
level_ratio = 4 level_ratio = 6
start = vocab_size * level_ratio * level_ratio start = vocab_size * level_ratio * level_ratio
end = start * level_ratio end = start * level_ratio
size = end * level_ratio size = end * level_ratio
@ -47,15 +47,15 @@ if __name__ == "__main__":
train_dataset, val_dataset = raw_dataset.Split(0.95) train_dataset, val_dataset = raw_dataset.Split(0.95)
train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size) train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size) val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
it = iter(train_dataloader) # it = iter(train_dataloader)
print("data samples:") # print("data samples:")
for i in range(10): # for i in range(10):
print(next(it)["input_ids"].numpy().tolist()) # print(next(it)["input_ids"].numpy().tolist())
torch.set_float32_matmul_precision("medium") torch.set_float32_matmul_precision("medium")
lit_trainer = pl.Trainer( lit_trainer = pl.Trainer(
accelerator="gpu", accelerator="cuda",
# devices=[0], devices=[0, 1],
precision=precision, precision=precision,
logger=TBLogger("./", default_hp_metric=False), logger=TBLogger("./", default_hp_metric=False),
strategy=strategy, strategy=strategy,