206 lines
7.5 KiB
Python
206 lines
7.5 KiB
Python
from fastchat.train.llama_flash_attn_monkey_patch import (
|
|
replace_llama_attn_with_flash_attn,
|
|
)
|
|
|
|
replace_llama_attn_with_flash_attn()
|
|
|
|
import json
|
|
from torch.utils.data import Dataset
|
|
from accelerate import Accelerator
|
|
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer, AdamW
|
|
import torch
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
|
|
|
|
IGNORE_TOKEN_ID = -100
|
|
|
|
|
|
class MixData(Dataset):
|
|
def __init__(self, dataset, ratio, tokenizer):
|
|
super(Dataset, self).__init__()
|
|
self.dataset = dataset
|
|
self.data_size = [len(c) for c in self.dataset]
|
|
ratio = [r if isinstance(r, int) else s for r, s in zip(ratio, self.data_size)]
|
|
self.ratio = ratio
|
|
self.tokenizer = tokenizer
|
|
self.sample_size = [int(self.data_size[0] / self.ratio[0] * r) for r in self.ratio]
|
|
print(self.data_size, self.sample_size, [c1 / c2 for c1, c2 in zip(self.sample_size, self.data_size)])
|
|
|
|
@staticmethod
|
|
def rounder(number):
|
|
rand = np.random.rand()
|
|
if rand < number - int(number):
|
|
return int(number) + 1
|
|
else:
|
|
return int(number)
|
|
|
|
@staticmethod
|
|
def choice_index(number, sample_size):
|
|
for i in range(len(sample_size)):
|
|
if number < sum(sample_size[:i + 1]):
|
|
return i, number - sum(sample_size[:i])
|
|
|
|
def __getitem__(self, index):
|
|
corpus_id, index = self.choice_index(index, self.sample_size)
|
|
rand = np.random.rand()
|
|
index = self.rounder((index + rand) / self.sample_size[corpus_id] * self.data_size[corpus_id])
|
|
index = min(index, len(self.dataset[corpus_id]) - 1)
|
|
return self.dataset[corpus_id][index]
|
|
|
|
def __len__(self):
|
|
return sum(self.sample_size)
|
|
|
|
def set_ratio(self, ratio):
|
|
self.ratio = ratio
|
|
self.data_size = [len(c) for c in self.dataset]
|
|
self.sample_size = [int(self.data_size[0] / self.ratio[0] * r) for r in self.ratio]
|
|
print(self.data_size, self.sample_size, [c1 / c2 for c1, c2 in zip(self.sample_size, self.data_size)])
|
|
|
|
def collate_fn(self, data):
|
|
input_ids, labels = zip(*data)
|
|
input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
|
|
labels = pad_sequence(labels, batch_first=True, padding_value=-100)
|
|
attention_mask = input_ids.ne(self.tokenizer.pad_token_id)
|
|
features = {
|
|
'input_ids': input_ids.long(),
|
|
'labels': labels.long(),
|
|
'attention_mask': attention_mask.long(),
|
|
}
|
|
return features
|
|
|
|
|
|
def last_index(lst, value):
|
|
return next((len(lst) - i - 1 for i, x in enumerate(lst[::-1]) if x != value), -1)
|
|
|
|
|
|
def safe_ids(ids, max_value, pad_id):
|
|
return [i if i < max_value else pad_id for i in ids]
|
|
|
|
|
|
dummy_message = [{"role": "user", "content": "Who are you?"},
|
|
{"role": "assistant", "content": "I am vicuna, a language model trained by researchers from open-source community."},
|
|
{"role": "user", "content": "What can you do?"},
|
|
{"role": "assistant", "content": "I can chat with you."}]
|
|
|
|
|
|
def tokenize(messages, tokenizer):
|
|
roles = {"user": "USER", "assistant": "ASSISTANT"}
|
|
input_ids = []
|
|
labels = []
|
|
system = "A chat between a curious user and an artificial intelligence assistant. " \
|
|
"The assistant gives helpful, detailed, and polite answers to the user's questions."
|
|
system_ids = tokenizer.encode(system, add_special_tokens=False)
|
|
input_ids += system_ids
|
|
labels += [IGNORE_TOKEN_ID] * len(system_ids)
|
|
for i, turn in enumerate(messages):
|
|
role = roles.get(turn['role'], 'USER')
|
|
content = turn['content']
|
|
content = content.strip()
|
|
if role == 'ASSISTANT':
|
|
content += '</s>'
|
|
role_ids = tokenizer.encode(role + ":", add_special_tokens=False)
|
|
content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
|
|
max_length=tokenizer.model_max_length)
|
|
input_ids += role_ids + content_ids
|
|
if role == 'ASSISTANT':
|
|
labels += [IGNORE_TOKEN_ID] * len(role_ids) + content_ids
|
|
else:
|
|
labels += [IGNORE_TOKEN_ID] * (len(role_ids) + len(content_ids))
|
|
|
|
if tokenizer.add_bos_token:
|
|
input_ids = [tokenizer.bos_token_id] + input_ids
|
|
labels = [IGNORE_TOKEN_ID] + labels
|
|
|
|
input_ids = input_ids[:tokenizer.model_max_length]
|
|
labels = labels[:tokenizer.model_max_length]
|
|
|
|
trunc_id = last_index(labels, IGNORE_TOKEN_ID) + 1
|
|
input_ids = input_ids[:trunc_id]
|
|
labels = labels[:trunc_id]
|
|
if len(labels) == 0:
|
|
return tokenize(dummy_message, tokenizer)
|
|
input_ids = safe_ids(input_ids, tokenizer.vocab_size, tokenizer.pad_token_id)
|
|
labels = safe_ids(labels, tokenizer.vocab_size, IGNORE_TOKEN_ID)
|
|
return input_ids, labels
|
|
|
|
|
|
class VicunaData(Dataset):
|
|
def __init__(self, data, tokenizer):
|
|
self.data = data
|
|
self.tokenizer = tokenizer
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
def __getitem__(self, item):
|
|
item = self.data[item]
|
|
input_ids, labels = tokenize(item, self.tokenizer)
|
|
return torch.tensor(input_ids), torch.tensor(labels)
|
|
|
|
def collate_fn(self, data):
|
|
input_ids, labels = zip(*data)
|
|
input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
|
|
labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_TOKEN_ID)
|
|
attention_mask = input_ids.ne(self.tokenizer.pad_token_id)
|
|
features = {
|
|
'input_ids': input_ids.long(),
|
|
'labels': labels.long(),
|
|
'attention_mask': attention_mask.long(),
|
|
}
|
|
return features
|
|
|
|
|
|
def main():
|
|
accelerator = Accelerator(gradient_accumulation_steps=4)
|
|
batch_size = 4
|
|
|
|
save_path = 'out/baichuan-vicuna-7b'
|
|
model_name = 'fireballoon/baichuan-llama-7b'
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, padding_side="right", model_max_length=4096)
|
|
tokenizer.pad_token = tokenizer.unk_token
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(model_name)
|
|
model.config.use_cache = False
|
|
model.gradient_checkpointing_enable()
|
|
|
|
dataset = VicunaData(
|
|
json.load(open('data/new/share_gpt-90k.json')) +
|
|
json.load(open('data/new/cot-75k.json')) +
|
|
json.load(open('data/new/leet-9k.json')), tokenizer)
|
|
|
|
print(len(dataset))
|
|
|
|
data_loader = torch.utils.data.DataLoader(dataset, collate_fn=dataset.collate_fn,
|
|
batch_size=batch_size, num_workers=0, shuffle=True)
|
|
|
|
optimizer = AdamW(model.parameters(), 2e-5)
|
|
model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)
|
|
|
|
for epoch in range(10):
|
|
accelerator.print(f'Training {save_path} {epoch}')
|
|
accelerator.wait_for_everyone()
|
|
model.train()
|
|
tk0 = tqdm(data_loader, total=len(data_loader))
|
|
loss_report = []
|
|
for batch in tk0:
|
|
with accelerator.accumulate(model):
|
|
out = model(**batch)
|
|
loss = out.loss
|
|
|
|
accelerator.backward(loss)
|
|
accelerator.clip_grad_norm_(model.parameters(), 1.)
|
|
optimizer.step()
|
|
optimizer.zero_grad()
|
|
|
|
loss_report.append(accelerator.gather(loss).mean().item())
|
|
tk0.set_postfix(loss=sum(loss_report[-100:]) / len(loss_report[-100:]))
|
|
accelerator.wait_for_everyone()
|
|
model.save_checkpoint(f'{save_path}/{epoch}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|