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

Try model train.

This commit is contained in:
Colin 2024-03-05 22:09:28 +08:00
parent 11fc8f1d39
commit 9ef3e92b23
1 changed files with 16 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
wit/lit_train.py
View File

@ -17,6 +17,7 @@ from transformers import (
from modelscope import snapshot_download from modelscope import snapshot_download
from lit_module import LitModule from lit_module import LitModule
from tokenization_qwen import QWenTokenizer from tokenization_qwen import QWenTokenizer
from logger import TBLogger
model_name = "qwen/Qwen-1_8B-Chat" model_name = "qwen/Qwen-1_8B-Chat"
learning_rate = 0.0001 learning_rate = 0.0001
@ -27,20 +28,22 @@ train_batch_size = 256
val_batch_size = 16 val_batch_size = 16
num_proc = 8 num_proc = 8
max_epochs = 1000 max_epochs = 1000
strategy = "fsdp" strategy = "auto"
resume_from_ckpt_path = None resume_from_ckpt_path = None
seed = 42 seed = 42
vocab_size = 4096
class SpecialDataset(Dataset): class SpecialDataset(Dataset):
def __init__(self, start=1, end=4096, size=65536): def __init__(self, start=1, end=320, size=32768):
self.size = size self.size = size
self.features = [] self.features = []
a = torch.randint(start, end, [size]) a = torch.randint(start, end, [size])
b = torch.randint(start, end, [size]) b = torch.randint(start, end, [size])
c = torch.randint(start, end, [size]) c = torch.randint(start, end, [size])
d = torch.randint(start, end, [size]) d = torch.randint(start, end, [size])
self.data = torch.stack([a, b, c, d, ((a + b + c + d) / 4).long()]).permute(1, 0) # self.data = torch.stack([a, b, a + b, a + b]).permute(1, 0)
self.data = torch.stack([a, a + a, a + a, a + a]).permute(1, 0)
def __len__(self): def __len__(self):
return self.size return self.size
@ -50,7 +53,8 @@ class SpecialDataset(Dataset):
data = self.data[idx] data = self.data[idx]
output["input_ids"] = data output["input_ids"] = data
output["labels"] = data.clone() output["labels"] = data.clone()
output["labels"][:4] = 0 # output["labels"][:2] = 0
# output["labels"][:2] = vocab_size
output["token_type_ids"] = torch.zeros(data.shape) output["token_type_ids"] = torch.zeros(data.shape)
return output return output
@ -67,10 +71,7 @@ if __name__ == "__main__":
tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken") tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
raw_dataset = SpecialDataset() train_dataset, val_dataset = random_split(SpecialDataset(), [0.95, 0.05])
train_idx, val_idx = random_split(list(range(len(raw_dataset))), [0.95, 0.05])
train_dataset = Subset(raw_dataset, train_idx.indices)
val_dataset = Subset(raw_dataset, val_idx.indices)
train_dataloader = DataLoader( train_dataloader = DataLoader(
train_dataset, train_dataset,
@ -89,8 +90,13 @@ if __name__ == "__main__":
) )
torch.set_float32_matmul_precision("medium") torch.set_float32_matmul_precision("medium")
precision = precision lit_trainer = pl.Trainer(
lit_trainer = pl.Trainer(accelerator="gpu", precision=precision, strategy=strategy, max_epochs=max_epochs) accelerator="gpu",
precision=precision,
logger=TBLogger("./", default_hp_metric=False),
strategy=strategy,
max_epochs=max_epochs,
)
lit_trainer.fit( lit_trainer.fit(
lit_module, lit_module,
train_dataloaders=train_dataloader, train_dataloaders=train_dataloader,