100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
import argparse
|
|
from functools import partial
|
|
from itertools import chain
|
|
from typing import Dict, Tuple
|
|
|
|
import datasets
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split, Subset
|
|
|
|
from transformers import (
|
|
BatchEncoding,
|
|
DefaultDataCollator,
|
|
PreTrainedTokenizer,
|
|
set_seed,
|
|
)
|
|
from modelscope import snapshot_download
|
|
from lit_module import LitModule
|
|
from tokenization_qwen import QWenTokenizer
|
|
|
|
model_name = "qwen/Qwen-1_8B-Chat"
|
|
learning_rate = 0.0001
|
|
use_tril_attention_mask = None
|
|
precision = "32-true" # "precision:bf16-mixed,16-mixed,32-true"
|
|
tokenizer_name_or_path = None
|
|
train_batch_size = 256
|
|
val_batch_size = 16
|
|
num_proc = 8
|
|
max_epochs = 1000
|
|
strategy = "fsdp"
|
|
resume_from_ckpt_path = None
|
|
seed = 42
|
|
|
|
|
|
class SpecialDataset(Dataset):
|
|
def __init__(self, start=1, end=4096, size=65536):
|
|
self.size = size
|
|
self.features = []
|
|
a = torch.randint(start, end, [size])
|
|
b = torch.randint(start, end, [size])
|
|
c = 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)
|
|
|
|
def __len__(self):
|
|
return self.size
|
|
|
|
def __getitem__(self, idx):
|
|
output = {}
|
|
data = self.data[idx]
|
|
output["input_ids"] = data
|
|
output["labels"] = data.clone()
|
|
output["labels"][:4] = 0
|
|
output["token_type_ids"] = torch.zeros(data.shape)
|
|
return output
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if tokenizer_name_or_path is None:
|
|
tokenizer_name_or_path = model_name
|
|
|
|
set_seed(seed)
|
|
|
|
# lightning module
|
|
model_dir = snapshot_download(model_name)
|
|
lit_module = LitModule(model_dir, learning_rate, use_tril_attention_mask)
|
|
|
|
tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
|
|
|
|
raw_dataset = SpecialDataset()
|
|
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_dataset,
|
|
batch_size=train_batch_size,
|
|
num_workers=num_proc,
|
|
collate_fn=DefaultDataCollator(),
|
|
persistent_workers=True,
|
|
shuffle=True,
|
|
)
|
|
val_dataloader = DataLoader(
|
|
val_dataset,
|
|
batch_size=val_batch_size,
|
|
num_workers=num_proc,
|
|
collate_fn=DefaultDataCollator(),
|
|
persistent_workers=True,
|
|
)
|
|
|
|
torch.set_float32_matmul_precision("medium")
|
|
precision = precision
|
|
lit_trainer = pl.Trainer(accelerator="gpu", precision=precision, strategy=strategy, max_epochs=max_epochs)
|
|
lit_trainer.fit(
|
|
lit_module,
|
|
train_dataloaders=train_dataloader,
|
|
val_dataloaders=val_dataloader,
|
|
ckpt_path=resume_from_ckpt_path,
|
|
)
|