Refine train.py.
This commit is contained in:
parent
b0ca4dc35d
commit
33b351ff8a
35
wit/train.py
35
wit/train.py
|
@ -17,32 +17,49 @@ 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
|
||||
train_batch_size = 32
|
||||
val_batch_size = 32
|
||||
num_proc = 8
|
||||
max_epochs = 1000
|
||||
strategy = "auto"
|
||||
resume_from_ckpt_path = None
|
||||
seed = 42
|
||||
vocab_size = 256
|
||||
|
||||
vocab_size = 1024
|
||||
level_ratio = 4
|
||||
|
||||
hidden_size = 256 # 128 1024 2048 32
|
||||
num_attention_heads = 8 # 8 8 16
|
||||
num_hidden_layers = 1 # 6 12 24 3
|
||||
|
||||
name = "vocab_level_hidden_head_layer"
|
||||
version = (
|
||||
str(vocab_size)
|
||||
+ "_"
|
||||
+ str(level_ratio)
|
||||
+ "_"
|
||||
+ str(hidden_size)
|
||||
+ "_"
|
||||
+ str(num_attention_heads)
|
||||
+ "_"
|
||||
+ str(num_hidden_layers)
|
||||
)
|
||||
|
||||
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
|
||||
config.hidden_size = hidden_size
|
||||
config.num_hidden_layers = num_hidden_layers
|
||||
config.num_attention_heads = num_attention_heads
|
||||
|
||||
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
|
||||
start = vocab_size * level_ratio * level_ratio * level_ratio * level_ratio
|
||||
end = start * level_ratio
|
||||
size = end * level_ratio
|
||||
size = start + start
|
||||
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)
|
||||
|
@ -57,7 +74,7 @@ if __name__ == "__main__":
|
|||
accelerator="cuda",
|
||||
devices=[0, 1],
|
||||
precision=precision,
|
||||
logger=TBLogger("./", default_hp_metric=False),
|
||||
logger=TBLogger("./log/", name=name, version=version, default_hp_metric=False),
|
||||
strategy=strategy,
|
||||
max_epochs=max_epochs,
|
||||
)
|
||||
|
|
Loading…
Reference in New Issue