mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-05 17:30:48 +08:00
Yannic Kilcher
@@ -0,0 +1,38 @@
|
||||
# Train using supervised examples
|
||||
|
||||
Requirements
|
||||
|
||||
```
|
||||
wandb
|
||||
evaluate
|
||||
datasets
|
||||
transformers
|
||||
torch
|
||||
```
|
||||
|
||||
Start training reward model
|
||||
|
||||
```bash
|
||||
python trainer.py --configs defaults galactica-125
|
||||
```
|
||||
|
||||
## Dataset
|
||||
|
||||
For now we only support webgpt and summary dataset from OpenAI. Once
|
||||
open-asisstant dataset are available it will be added here.
|
||||
|
||||
## Model
|
||||
|
||||
TBD
|
||||
|
||||
## Results
|
||||
|
||||
Experimental results in wandb
|
||||
[here](https://wandb.ai/sanagnos/supervised-finetuning?workspace=user-sanagnos).
|
||||
|
||||
## TODOS
|
||||
|
||||
- decide on a model
|
||||
- add special token to declare prompt and reply. Do nto freeze the weights for
|
||||
these
|
||||
- Merge utils etc with reward model
|
||||
@@ -0,0 +1,37 @@
|
||||
defaults:
|
||||
learning_rate: 1e-5
|
||||
gradient_checkpointing: false
|
||||
gradient_accumulation_steps: 32
|
||||
per_device_train_batch_size: 2
|
||||
per_device_eval_batch_size: 2
|
||||
weight_decay: 0.00
|
||||
warmup_steps: 600
|
||||
eval_steps: 200
|
||||
save_steps: 500
|
||||
max_length: 512
|
||||
num_train_epochs: 3
|
||||
logging_steps: 10
|
||||
max_grad_norm: 2.0
|
||||
save_total_limit: 4
|
||||
eval_accumulation_steps:
|
||||
freeze_layer:
|
||||
datasets:
|
||||
- webgpt
|
||||
cache_dir: ~/.cache
|
||||
loss_fn: CrossEntropyLoss
|
||||
eval_size:
|
||||
log_dir: "base"
|
||||
|
||||
galactica-125:
|
||||
learning_rate: 5e-5
|
||||
model_name: facebook/galactica-125m
|
||||
weight_decay: 0.01
|
||||
warmup_steps: 600
|
||||
gradient_checkpointing: false
|
||||
gradient_accumulation_steps: 2
|
||||
per_device_train_batch_size: 4
|
||||
per_device_eval_batch_size: 4
|
||||
|
||||
debug:
|
||||
eval_steps: 20
|
||||
eval_size: 100
|
||||
@@ -0,0 +1,67 @@
|
||||
from datasets import load_dataset
|
||||
from sklearn.model_selection import train_test_split
|
||||
from torch.utils.data import Dataset, Subset
|
||||
|
||||
|
||||
class SquadV2Dataset(Dataset):
|
||||
def __init__(self, cache_dir, split):
|
||||
self.dataset = load_dataset("squad_v2", cache_dir=cache_dir, split=split)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
data = self.dataset[idx]
|
||||
# dummy return first answer
|
||||
return "".join([data["title"], ". ", data["context"], " " + data["question"]]), data["answers"]["text"][0]
|
||||
|
||||
|
||||
class WebGPT(Dataset):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
dataset = load_dataset("openai/webgpt_comparisons")
|
||||
questions = {}
|
||||
# using prompt as our index will allows us
|
||||
# to add additional generated prompt later
|
||||
self.index2question = {}
|
||||
for row in dataset["train"]:
|
||||
question = row["question"]["full_text"]
|
||||
if question not in self.index2question:
|
||||
self.index2question[len(self.index2question)] = question
|
||||
|
||||
# only keep the best answer
|
||||
questions[question] = row["answer_0" if row["score_0"] > row["score_1"] else "answer_1"]
|
||||
|
||||
self.questions = questions
|
||||
|
||||
def __len__(self):
|
||||
return len(self.index2question)
|
||||
|
||||
def __getitem__(self, index):
|
||||
question = self.index2question[index]
|
||||
answer = self.questions[question]
|
||||
return [question, answer]
|
||||
|
||||
|
||||
def train_val_dataset(dataset, val_split=0.2):
|
||||
train_idx, val_idx = train_test_split(
|
||||
list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
|
||||
)
|
||||
return Subset(dataset, train_idx), Subset(dataset, val_idx)
|
||||
|
||||
|
||||
def get_one_dataset(conf, dataset_name):
|
||||
dataset_name = dataset_name.lower()
|
||||
|
||||
if dataset_name == "squadv2":
|
||||
raise ValueError("SquadV2 is not diverse enough for generation .. ")
|
||||
train = SquadV2Dataset(conf.cache_dir, "train")
|
||||
eval = SquadV2Dataset(conf.cache_dir, "validation")
|
||||
elif dataset_name == "webgpt":
|
||||
dataset = WebGPT()
|
||||
train, eval = train_val_dataset(dataset, val_split=0.2)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset {dataset_name}")
|
||||
|
||||
return train, eval
|
||||
@@ -0,0 +1,85 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
|
||||
|
||||
|
||||
@dataclass
|
||||
class DialogueDataCollator:
|
||||
"""
|
||||
Expects a list of texts corresponding to a sequence of [question, answer, question, answer, ...] pairs.
|
||||
"""
|
||||
|
||||
tokenizer: PreTrainedTokenizerBase
|
||||
padding: Union[bool, str, PaddingStrategy] = True
|
||||
max_length: Optional[int] = None
|
||||
pad_to_multiple_of: Optional[int] = None
|
||||
|
||||
def __call__(self, features):
|
||||
# TODO add special tokens for question and answer here
|
||||
# additional_special_tokens = ['<question>', '<answer>']
|
||||
prompt_tokens = ["Question: ", "Answer: "]
|
||||
|
||||
flatten_messages = []
|
||||
label_masks = []
|
||||
|
||||
for messages in features:
|
||||
assert len(messages) % 2 == 0, "Number of messages must be even"
|
||||
messages = [
|
||||
(prompt_tokens[0] if i % 2 == 0 else "") + x + ((" " + prompt_tokens[1]) if i % 2 == 0 else "")
|
||||
for i, x in enumerate(messages)
|
||||
]
|
||||
|
||||
# Add a way for the model to terminate generation, reinitialize prompter
|
||||
messages.append(prompt_tokens[0])
|
||||
|
||||
flatten_messages.append(
|
||||
self.tokenizer(
|
||||
"".join(messages),
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
return_offsets_mapping=True,
|
||||
)
|
||||
)
|
||||
|
||||
message_change_indices = np.cumsum([len(x) for x in messages[:-1]])
|
||||
# for each token an integer indicating the index of the message it belongs to. Just to create the label mask.
|
||||
# TEXT: Question: Hello, how are you? Answer: I am fine. Question: What is your name? Answer: My name is John.
|
||||
# MESSAGE_INDICES: 0 0 0 0 0 0 1 1 1 2 2 2 2 2 2 3 3 3 3
|
||||
|
||||
# If no result in next, we are predicting the last termination token(s)
|
||||
message_indices = list(
|
||||
map(
|
||||
lambda x: next((i for i, val in enumerate(message_change_indices) if val >= x), -2),
|
||||
list(map(lambda x: x[1], flatten_messages[-1]["offset_mapping"])),
|
||||
)
|
||||
)
|
||||
label_mask = np.roll(list(map(lambda x: x % 2 == 1, message_indices)), -1, -1)
|
||||
try:
|
||||
label_mask[[i for i in range(len(message_indices)) if message_indices[i] == -2][0] - 1] = True
|
||||
except IndexError:
|
||||
# an aftermath of padding
|
||||
pass
|
||||
|
||||
label_masks.append(label_mask)
|
||||
flatten_messages[-1].pop("offset_mapping")
|
||||
|
||||
batch = self.tokenizer.pad(
|
||||
flatten_messages,
|
||||
padding=self.padding,
|
||||
max_length=self.max_length,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors="pt",
|
||||
)
|
||||
dim = batch["input_ids"].shape[-1]
|
||||
|
||||
batch["label_masks"] = torch.stack([F.pad(torch.tensor(x), (0, dim - len(x))) for x in label_masks])
|
||||
|
||||
for k in list(batch.keys()):
|
||||
if k not in ["input_ids", "attention_mask", "label_masks"]:
|
||||
batch.pop(k)
|
||||
|
||||
return batch
|
||||
@@ -0,0 +1,15 @@
|
||||
from torch import nn
|
||||
|
||||
|
||||
class CrossEntropyLoss(nn.CrossEntropyLoss):
|
||||
def __init__(self, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean"):
|
||||
super(CrossEntropyLoss, self).__init__(weight, size_average, ignore_index, reduce, reduction)
|
||||
|
||||
def forward(self, input, target, mask=None):
|
||||
if mask is not None:
|
||||
mask = mask.view(-1)
|
||||
input = input.view(-1, input.size(-1))
|
||||
target = target.view(-1)
|
||||
input = input[mask]
|
||||
target = target[mask]
|
||||
return super(CrossEntropyLoss, self).forward(input, target)
|
||||
@@ -0,0 +1,200 @@
|
||||
import argparse
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from distutils.util import strtobool
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import (
|
||||
DataCollator,
|
||||
EvalPrediction,
|
||||
PreTrainedModel,
|
||||
PreTrainedTokenizerBase,
|
||||
Trainer,
|
||||
TrainerCallback,
|
||||
TrainingArguments,
|
||||
get_cosine_schedule_with_warmup,
|
||||
)
|
||||
from utils import get_dataset, get_loss, get_model, get_tokenizer, read_yamls
|
||||
|
||||
os.environ["WANDB_PROJECT"] = "supervised-finetuning"
|
||||
|
||||
|
||||
@dataclass
|
||||
class CustomTrainingArguments(TrainingArguments):
|
||||
loss_function: str = "CrossEntropyLoss"
|
||||
|
||||
|
||||
def compute_metrics(eval_pred):
|
||||
pred_ids = eval_pred.predictions
|
||||
labels = eval_pred.label_ids
|
||||
|
||||
return {"accuracy": (pred_ids[labels > 0] == labels[labels > 0]).mean()}
|
||||
|
||||
|
||||
def preprocess_logits_for_metrics(logits, labels):
|
||||
pred_ids = torch.argmax(logits, dim=-1)
|
||||
return pred_ids
|
||||
|
||||
|
||||
class SFTTrainer(Trainer):
|
||||
def __init__(
|
||||
self,
|
||||
model: Union[PreTrainedModel, nn.Module] = None,
|
||||
args: TrainingArguments = None,
|
||||
data_collator: Optional[DataCollator] = None,
|
||||
train_dataset: Optional[Dataset] = None,
|
||||
eval_dataset: Optional[Dataset] = None,
|
||||
tokenizer: Optional[PreTrainedTokenizerBase] = None,
|
||||
model_init: Callable[[], PreTrainedModel] = None,
|
||||
compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
|
||||
callbacks: Optional[List[TrainerCallback]] = None,
|
||||
optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
|
||||
preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
|
||||
):
|
||||
super().__init__(
|
||||
model,
|
||||
args,
|
||||
data_collator,
|
||||
train_dataset,
|
||||
eval_dataset,
|
||||
tokenizer,
|
||||
model_init,
|
||||
compute_metrics,
|
||||
callbacks,
|
||||
optimizers,
|
||||
preprocess_logits_for_metrics,
|
||||
)
|
||||
self.loss_fct = get_loss(args.loss_function)
|
||||
|
||||
def fetch_scheduler(self):
|
||||
return get_cosine_schedule_with_warmup(
|
||||
self.optimizer,
|
||||
num_warmup_steps=self.args.warmup_steps,
|
||||
num_training_steps=self.num_train_steps,
|
||||
num_cycles=1,
|
||||
last_epoch=-1,
|
||||
)
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False):
|
||||
labels_mask = inputs.pop("label_masks")
|
||||
|
||||
outputs = model(**inputs)
|
||||
|
||||
loss = self.loss_fct(outputs.get("logits"), torch.roll(inputs["input_ids"], -1, -1), mask=labels_mask)
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
|
||||
def _compute_loss(self, model, inputs):
|
||||
|
||||
labels_mask = inputs.pop("label_masks")
|
||||
|
||||
inputs = self._prepare_inputs(inputs)
|
||||
|
||||
outputs = model(**inputs)
|
||||
|
||||
logits = outputs.get("logits")
|
||||
|
||||
targets = torch.roll(inputs["input_ids"], -1, -1)
|
||||
loss = self.loss_fct(outputs.get("logits"), targets, mask=labels_mask)
|
||||
|
||||
return loss, logits, targets, labels_mask
|
||||
|
||||
def prediction_step(
|
||||
self,
|
||||
model: nn.Module,
|
||||
inputs: Dict[str, Union[torch.Tensor, Any]],
|
||||
prediction_loss_only: bool,
|
||||
ignore_keys: Optional[List[str]] = None,
|
||||
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
|
||||
|
||||
with torch.no_grad():
|
||||
loss, logits, labels, labels_mask = self._compute_loss(model, inputs)
|
||||
labels[~labels_mask] = -1
|
||||
|
||||
loss = loss.mean().detach()
|
||||
|
||||
if self.args.prediction_loss_only:
|
||||
return (loss, None, None)
|
||||
|
||||
return (loss, logits, labels)
|
||||
|
||||
|
||||
def _strtobool(x):
|
||||
return bool(strtobool(x))
|
||||
|
||||
|
||||
def argument_parsing(notebook=False, notebook_args=None):
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--configs", nargs="+", required=True)
|
||||
|
||||
if notebook:
|
||||
args, remaining = parser.parse_known_args(notebook_args)
|
||||
else:
|
||||
args, remaining = parser.parse_known_args()
|
||||
|
||||
# Config from YAML
|
||||
conf = {}
|
||||
configs = read_yamls("./configs")
|
||||
for name in args.configs:
|
||||
if "," in name:
|
||||
for n in name.split(","):
|
||||
conf.update(configs[n])
|
||||
else:
|
||||
conf.update(configs[name])
|
||||
|
||||
# Override config from command-line
|
||||
parser = argparse.ArgumentParser()
|
||||
for key, value in conf.items():
|
||||
type_ = type(value) if value is not None else str
|
||||
if type_ == bool:
|
||||
type_ = _strtobool
|
||||
parser.add_argument(f"--{key}", type=type_, default=value)
|
||||
|
||||
return parser.parse_args(remaining)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
training_conf = argument_parsing()
|
||||
|
||||
model = get_model(training_conf)
|
||||
tokenizer = get_tokenizer(training_conf)
|
||||
|
||||
train, evals, collate_fn = get_dataset(training_conf, tokenizer)
|
||||
|
||||
args = CustomTrainingArguments(
|
||||
output_dir=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned",
|
||||
num_train_epochs=training_conf.num_train_epochs,
|
||||
warmup_steps=training_conf.warmup_steps,
|
||||
loss_function=training_conf.loss_fn,
|
||||
learning_rate=float(training_conf.learning_rate),
|
||||
fp16=True,
|
||||
gradient_checkpointing=training_conf.gradient_checkpointing,
|
||||
gradient_accumulation_steps=training_conf.gradient_accumulation_steps,
|
||||
per_device_train_batch_size=training_conf.per_device_train_batch_size,
|
||||
per_device_eval_batch_size=training_conf.per_device_eval_batch_size,
|
||||
weight_decay=training_conf.weight_decay,
|
||||
max_grad_norm=training_conf.max_grad_norm,
|
||||
logging_steps=training_conf.logging_steps,
|
||||
save_total_limit=training_conf.save_total_limit,
|
||||
evaluation_strategy="steps",
|
||||
eval_steps=training_conf.eval_steps,
|
||||
save_steps=training_conf.save_steps,
|
||||
eval_accumulation_steps=training_conf.eval_accumulation_steps,
|
||||
report_to="wandb",
|
||||
)
|
||||
|
||||
assert len(evals) > 0
|
||||
trainer = SFTTrainer(
|
||||
model,
|
||||
args,
|
||||
train_dataset=train,
|
||||
eval_dataset=evals,
|
||||
data_collator=collate_fn,
|
||||
tokenizer=tokenizer,
|
||||
compute_metrics=compute_metrics,
|
||||
preprocess_logits_for_metrics=preprocess_logits_for_metrics,
|
||||
)
|
||||
trainer.train()
|
||||
@@ -0,0 +1,111 @@
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from custom_datasets import get_one_dataset
|
||||
from custom_datasets.dialogue_collator import DialogueDataCollator
|
||||
from losses import CrossEntropyLoss
|
||||
from sklearn.model_selection import train_test_split
|
||||
from torch.utils.data import ConcatDataset, Subset
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
SUPPORTED_MODELS = ["galactica"]
|
||||
|
||||
|
||||
def get_tokenizer(conf):
|
||||
tokenizer = AutoTokenizer.from_pretrained(conf.model_name, cache_dir=conf.cache_dir)
|
||||
|
||||
if "galactica" in conf.model_name:
|
||||
tokenizer.add_special_tokens({"pad_token": "<pad>", "eos_token": "</s>"})
|
||||
|
||||
return tokenizer
|
||||
|
||||
|
||||
def get_model(conf):
|
||||
if not any([x in conf.model_name for x in SUPPORTED_MODELS]):
|
||||
raise ValueError(
|
||||
f"Model {conf.model_name} not supported. Supported models: {SUPPORTED_MODELS}. "
|
||||
"To include more make sure the masking is dne correctly... (decoder only supported for now)"
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(conf.model_name, cache_dir=conf.cache_dir)
|
||||
|
||||
if conf.freeze_layer:
|
||||
model = freeze_top_n_layers(model, conf.freeze_layer)
|
||||
|
||||
model_parameters = filter(lambda p: p.requires_grad, model.parameters())
|
||||
params = sum([p.numel() for p in model_parameters])
|
||||
print("Number of trainable parameters: {}M".format(int(params / 1e6)))
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def get_dataset(conf, tokenizer):
|
||||
train_datasets, evals = [], {}
|
||||
|
||||
for dataset_name in conf.datasets:
|
||||
train, val = get_one_dataset(conf, dataset_name)
|
||||
train_datasets.append(train)
|
||||
evals[dataset_name] = Subset(val, list(range(min(len(val), conf.eval_size)))) if conf.eval_size else val
|
||||
|
||||
train = ConcatDataset(train_datasets)
|
||||
|
||||
collate_fn = DialogueDataCollator(tokenizer, max_length=conf.max_length)
|
||||
|
||||
return train, evals, collate_fn
|
||||
|
||||
|
||||
def get_loss(loss):
|
||||
if loss == "CrossEntropyLoss":
|
||||
return CrossEntropyLoss()
|
||||
else:
|
||||
raise ValueError(f"Loss {loss} not supported")
|
||||
|
||||
|
||||
def read_yamls(dir):
|
||||
conf = {}
|
||||
no_conf = True
|
||||
|
||||
for config_file in Path(dir).glob("**/*.yaml"):
|
||||
no_conf = False
|
||||
with config_file.open("r") as f:
|
||||
conf.update(yaml.safe_load(f))
|
||||
|
||||
if no_conf:
|
||||
print(f"WARNING: No yaml files found in {dir}")
|
||||
|
||||
return conf
|
||||
|
||||
|
||||
def train_val_dataset(dataset, val_split=0.2):
|
||||
train_idx, val_idx = train_test_split(
|
||||
list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
|
||||
)
|
||||
return Subset(dataset, train_idx), Subset(dataset, val_idx)
|
||||
|
||||
|
||||
def freeze_top_n_layers(model, target_layers):
|
||||
# its possible we can simply detect which module is a ModuleList
|
||||
# and simply freeze the module without doing string parsing
|
||||
for name, param in model.named_parameters():
|
||||
if "embed" in name:
|
||||
param.requires_grad = False
|
||||
elif ".layer" in name or ".h." in name:
|
||||
tokens = name.split(".")
|
||||
layer_ = None
|
||||
for token in tokens:
|
||||
if token.isdigit():
|
||||
layer_ = int(token)
|
||||
break
|
||||
|
||||
if layer_ is not None and layer_ < target_layers:
|
||||
# print('freeze ', layer_, name)
|
||||
param.requires_grad = False
|
||||
return model
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
|
||||
model = AutoModelForSequenceClassification.from_pretrained("bigscience/bloomz-560m")
|
||||
freeze_top_n_layers(model, 10)
|
||||
print(model.state_dict().keys())
|
||||
Reference in New Issue
Block a user