506 lines
20 KiB
Python
506 lines
20 KiB
Python
# coding=utf-8
|
|
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
import bitsandbytes as bnb
|
|
import torch
|
|
from accelerate import Accelerator
|
|
from datasets import DatasetDict, builder, load_dataset
|
|
from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
|
|
from peft.tuners.lora import LoraLayer
|
|
from torch import nn
|
|
from tqdm import tqdm
|
|
from transformers import (
|
|
AutoModelForSequenceClassification,
|
|
AutoTokenizer,
|
|
BitsAndBytesConfig,
|
|
HfArgumentParser,
|
|
PreTrainedTokenizerBase,
|
|
TrainingArguments,
|
|
)
|
|
|
|
from trl import RewardTrainer
|
|
|
|
|
|
tqdm.pandas()
|
|
builder.has_sufficient_disk_space = lambda needed_bytes, directory=".": True
|
|
# torch.autograd.set_detect_anomaly(True)
|
|
|
|
### fix from https://github.com/huggingface/trl/issues/274
|
|
|
|
|
|
class GPTRewardTrainer(RewardTrainer):
|
|
def compute_loss(
|
|
self,
|
|
model,
|
|
inputs,
|
|
return_outputs=False,
|
|
):
|
|
rewards = model(
|
|
input_ids=inputs["input_ids"],
|
|
attention_mask=inputs["attention_mask"],
|
|
)[0]
|
|
bsz = rewards.size(0)
|
|
jidx = torch.arange(0, bsz, 2)
|
|
kidx = jidx + 1
|
|
rewards_chosen = rewards[jidx]
|
|
rewards_rejected = rewards[kidx]
|
|
loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
|
|
if return_outputs:
|
|
return loss, {
|
|
"rewards_chosen": rewards_chosen,
|
|
"rewards_rejected": rewards_rejected,
|
|
}
|
|
return loss
|
|
|
|
|
|
@dataclass
|
|
class GPTRewardDataCollatorWithPadding:
|
|
r"""
|
|
Reward DataCollator class that pads the inputs to the maximum length of the batch.
|
|
Args:
|
|
tokenizer (`PreTrainedTokenizerBase`):
|
|
The tokenizer used for encoding the data.
|
|
padding (`Union[bool, str, `PaddingStrategy`]`, `optional`, defaults to `True`):
|
|
padding_strategy to pass to the tokenizer.
|
|
max_length (`Optional[int]`, `optional`, defaults to `None`):
|
|
The maximum length of the sequence to be processed.
|
|
pad_to_multiple_of (`Optional[int]`, `optional`, defaults to `None`):
|
|
If set will pad the sequence to a multiple of the provided value.
|
|
return_tensors (`str`, `optional`, defaults to `"pt"`):
|
|
The tensor type to use.
|
|
"""
|
|
|
|
tokenizer: PreTrainedTokenizerBase
|
|
padding: Union[bool, str] = True
|
|
max_length: Optional[int] = None
|
|
pad_to_multiple_of: Optional[int] = None
|
|
return_tensors: str = "pt"
|
|
|
|
def __call__(self, features):
|
|
# features_chosen = []
|
|
# features_rejected = []
|
|
merged_features = []
|
|
for feature in features:
|
|
# check if the keys are named as expected
|
|
if (
|
|
"input_ids_chosen" not in feature
|
|
or "input_ids_rejected" not in feature
|
|
or "attention_mask_chosen" not in feature
|
|
or "attention_mask_rejected" not in feature
|
|
):
|
|
raise ValueError(
|
|
"The features should include `input_ids_chosen`, `attention_mask_chosen`, `input_ids_rejected` and `attention_mask_rejected`"
|
|
)
|
|
|
|
merged_features.append(
|
|
{
|
|
"input_ids": feature["input_ids_chosen"],
|
|
"attention_mask": feature["attention_mask_chosen"],
|
|
}
|
|
)
|
|
merged_features.append(
|
|
{
|
|
"input_ids": feature["input_ids_rejected"],
|
|
"attention_mask": feature["attention_mask_rejected"],
|
|
}
|
|
)
|
|
batch = self.tokenizer.pad(
|
|
merged_features,
|
|
padding=self.padding,
|
|
max_length=self.max_length,
|
|
pad_to_multiple_of=self.pad_to_multiple_of,
|
|
return_tensors=self.return_tensors,
|
|
)
|
|
batch = {
|
|
"input_ids": batch["input_ids"],
|
|
"attention_mask": batch["attention_mask"],
|
|
"return_loss": True,
|
|
}
|
|
return batch
|
|
|
|
|
|
# Define and parse arguments.
|
|
@dataclass
|
|
class ScriptArguments:
|
|
"""
|
|
The name of the Casual LM model we wish to fine with RewardTrainer
|
|
"""
|
|
|
|
model_name: Optional[str] = field(
|
|
default="/home/toolkit/huggingface/tldr_sft_pythia7b", metadata={"help": "the model name"}
|
|
)
|
|
dataset_name: Optional[str] = field(
|
|
default="mnoukhov/openai_summarize_comparisons_tldrprompt", metadata={"help": "the dataset name"}
|
|
)
|
|
dataset_text_field: Optional[str] = field(default="prompt", metadata={"help": "the text field of the dataset"})
|
|
log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
|
|
logging_steps: Optional[int] = field(default=100, metadata={"help": "the number of update steps between two logs"})
|
|
train_split: Optional[str] = field(
|
|
default="train", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
|
|
)
|
|
eval_split: Optional[str] = field(
|
|
default="test[:5000]", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
|
|
)
|
|
learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
|
|
weight_decay: Optional[float] = field(default=0.001)
|
|
num_warmup_steps: Optional[int] = field(default=100)
|
|
lr_scheduler_type: Optional[str] = field(default="cosine")
|
|
optimizer_type: Optional[str] = field(default="adamw_torch", metadata={"help": "the optimizer type"})
|
|
per_device_train_batch_size: Optional[int] = field(default=2, metadata={"help": "the per device train batch size"})
|
|
per_device_eval_batch_size: Optional[int] = field(default=1, metadata={"help": "the per device eval batch size"})
|
|
num_train_epochs: Optional[int] = field(default=1, metadata={"help": "the number of training epochs"})
|
|
seq_length: Optional[int] = field(default=560, metadata={"help": "Input sequence length"})
|
|
gradient_accumulation_steps: Optional[int] = field(
|
|
default=16, metadata={"help": "the number of gradient accumulation steps"}
|
|
)
|
|
bf16: Optional[bool] = field(
|
|
default=False,
|
|
metadata={
|
|
"help": "This essentially cuts the training time in half if you want to sacrifice a little precision and have a supported GPU."
|
|
},
|
|
)
|
|
fp16: Optional[bool] = field(
|
|
default=False,
|
|
metadata={
|
|
"help": "This essentially cuts the training time in half if you want to sacrifice a little precision and have a supported GPU."
|
|
},
|
|
)
|
|
fp16_model: Optional[bool] = field(
|
|
default=False,
|
|
metadata={},
|
|
)
|
|
load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 8 bits precision"})
|
|
load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
|
|
use_lora: Optional[bool] = field(
|
|
default=True,
|
|
)
|
|
lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"})
|
|
lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"})
|
|
lora_r: Optional[int] = field(default=8, metadata={"help": "the lora r parameter"})
|
|
lora_all_linear: Optional[bool] = field(default=False, metadata={"help": "lora adapter on all linear layers"})
|
|
trust_remote_code: Optional[bool] = field(default=True, metadata={"help": "Enable `trust_remote_code`"})
|
|
output_dir: Optional[str] = field(default="results", metadata={"help": "the output directory"})
|
|
gradient_checkpointing: Optional[bool] = field(
|
|
default=False,
|
|
metadata={"help": "Enables gradient checkpointing."},
|
|
)
|
|
mode: Optional[str] = field(default="train")
|
|
eval_steps: Optional[float] = field(default=None)
|
|
pretrained_adapter: Optional[str] = field(default=None)
|
|
padding: Optional[str] = field(
|
|
default="max_length", metadata={"help": "padding to use for preprocessing the dataset"}
|
|
)
|
|
save_strategy: Optional[str] = field(default="steps")
|
|
|
|
|
|
def find_all_linear_names(args, model):
|
|
cls = bnb.nn.Linear4bit if args.load_in_4bit else (bnb.nn.Linear8bitLt if args.load_in_8bit else torch.nn.Linear)
|
|
lora_module_names = set()
|
|
for name, module in model.named_modules():
|
|
if isinstance(module, cls):
|
|
names = name.split(".")
|
|
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
|
|
|
|
if "lm_head" in lora_module_names: # needed for 16-bit
|
|
lora_module_names.remove("lm_head")
|
|
|
|
if "score" in lora_module_names: # needed for 16-bit
|
|
lora_module_names.remove("score")
|
|
|
|
return list(lora_module_names)
|
|
|
|
|
|
def create_and_prepare_model(args):
|
|
if args.load_in_8bit and args.load_in_4bit:
|
|
raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
|
|
elif args.load_in_8bit or args.load_in_4bit:
|
|
quantization_config = BitsAndBytesConfig(load_in_8bit=args.load_in_8bit, load_in_4bit=args.load_in_4bit)
|
|
device_map = {"": Accelerator().local_process_index}
|
|
else:
|
|
device_map = None
|
|
quantization_config = None
|
|
|
|
if args.bf16:
|
|
torch_dtype = torch.bfloat16
|
|
elif args.fp16_model:
|
|
torch_dtype = torch.float16
|
|
else:
|
|
torch_dtype = torch.float32
|
|
|
|
model = AutoModelForSequenceClassification.from_pretrained(
|
|
args.model_name,
|
|
quantization_config=quantization_config,
|
|
device_map=device_map,
|
|
num_labels=1,
|
|
torch_dtype=torch_dtype,
|
|
)
|
|
|
|
model.config.torch_dtype = torch_dtype
|
|
model.config.use_cache = not args.gradient_checkpointing
|
|
|
|
# if script_args.ignore_bias_buffers:
|
|
# torch distributed hack
|
|
if quantization_config is not None:
|
|
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=args.gradient_checkpointing)
|
|
args.gradient_checkpointing = False
|
|
|
|
if args.use_lora:
|
|
# we add `score` to the list of modules to save to
|
|
# correctly save the score head.
|
|
if args.pretrained_adapter is not None:
|
|
model = PeftModel.from_pretrained(model, args.pretrained_adapter)
|
|
else:
|
|
if args.lora_all_linear:
|
|
target_modules = find_all_linear_names(args, model)
|
|
else:
|
|
target_modules = None
|
|
|
|
peft_config = LoraConfig(
|
|
r=args.lora_r,
|
|
lora_alpha=args.lora_alpha,
|
|
lora_dropout=args.lora_dropout,
|
|
bias="none",
|
|
task_type="SEQ_CLS",
|
|
target_modules=target_modules,
|
|
modules_to_save=["score"],
|
|
)
|
|
|
|
model = get_peft_model(model, peft_config)
|
|
|
|
modules_to_save = ["score"]
|
|
for key, _ in model.named_modules():
|
|
target_module_found = any(key.endswith(target_key) for target_key in modules_to_save)
|
|
if target_module_found:
|
|
model.get_submodule(key + ".original_module").requires_grad_(False)
|
|
|
|
if torch_dtype == torch.bfloat16:
|
|
for name, module in model.named_modules():
|
|
if isinstance(module, LoraLayer):
|
|
module = module.to(torch_dtype)
|
|
if "norm" in name:
|
|
module = module.to(torch.float32)
|
|
if "score" in name or "embed_tokens" in name:
|
|
if hasattr(module, "weight") and module.weight.dtype == torch.float32:
|
|
module = module.to(torch_dtype)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(script_args.model_name)
|
|
if getattr(tokenizer, "pad_token", None) is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
if getattr(model.config, "pad_token_id", None) is None:
|
|
model.config.pad_token_id = model.config.eos_token_id
|
|
|
|
return model, tokenizer
|
|
|
|
|
|
def prepare_dataset(args, dataset, tokenizer, num_proc=2):
|
|
# def summary_filter(example):
|
|
# return (example["chosen"] != example["rejected"]) and (
|
|
# len(example["chosen"].split()) >= 5 or len(example["rejected"].split()) >= 5
|
|
# )
|
|
#
|
|
# pre_filter = len(dataset)
|
|
# dataset = dataset.filter(summary_filter)
|
|
# print(f"filtered {pre_filter - len(dataset)} samples from {split}")
|
|
original_columns = dataset.column_names
|
|
|
|
def preprocess_function(examples):
|
|
new_examples = {
|
|
"input_ids_chosen": [],
|
|
"attention_mask_chosen": [],
|
|
"input_ids_rejected": [],
|
|
"attention_mask_rejected": [],
|
|
}
|
|
for prompt, chosen, rejected in zip(examples["prompt"], examples["chosen"], examples["rejected"]):
|
|
tokenized_chosen = tokenizer(
|
|
prompt + " " + chosen, padding=args.padding, truncation=True, max_length=script_args.seq_length
|
|
)
|
|
tokenized_rejected = tokenizer(
|
|
prompt + " " + rejected, padding=args.padding, truncation=True, max_length=script_args.seq_length
|
|
)
|
|
new_examples["input_ids_chosen"].append(tokenized_chosen["input_ids"])
|
|
new_examples["attention_mask_chosen"].append(tokenized_chosen["attention_mask"])
|
|
new_examples["input_ids_rejected"].append(tokenized_rejected["input_ids"])
|
|
new_examples["attention_mask_rejected"].append(tokenized_rejected["attention_mask"])
|
|
|
|
return new_examples
|
|
|
|
dataset = dataset.map(preprocess_function, batched=True, num_proc=num_proc, remove_columns=original_columns)
|
|
|
|
return dataset
|
|
|
|
|
|
parser = HfArgumentParser(ScriptArguments)
|
|
script_args = parser.parse_args_into_dataclasses()[0]
|
|
|
|
model, tokenizer = create_and_prepare_model(script_args)
|
|
if script_args.mode != "eval":
|
|
train_data = load_dataset(script_args.dataset_name, split=script_args.train_split)
|
|
train_dataset = prepare_dataset(script_args, train_data, tokenizer)
|
|
else:
|
|
train_dataset = None
|
|
|
|
if script_args.eval_split is not None and script_args.eval_split != "None":
|
|
eval_data = load_dataset(script_args.dataset_name, split=script_args.eval_split)
|
|
eval_dataset = prepare_dataset(script_args, eval_data, tokenizer)
|
|
else:
|
|
eval_dataset = None
|
|
|
|
# don't include gradient_checkpointing here, see trl#728
|
|
training_args = TrainingArguments(
|
|
output_dir=script_args.output_dir,
|
|
per_device_train_batch_size=script_args.per_device_train_batch_size,
|
|
per_device_eval_batch_size=script_args.per_device_eval_batch_size,
|
|
bf16=script_args.bf16,
|
|
fp16=script_args.fp16,
|
|
num_train_epochs=script_args.num_train_epochs,
|
|
gradient_accumulation_steps=script_args.gradient_accumulation_steps,
|
|
learning_rate=script_args.learning_rate,
|
|
report_to=script_args.log_with,
|
|
remove_unused_columns=False,
|
|
lr_scheduler_type=script_args.lr_scheduler_type,
|
|
weight_decay=script_args.weight_decay,
|
|
optim=script_args.optimizer_type,
|
|
warmup_steps=script_args.num_warmup_steps,
|
|
logging_steps=script_args.logging_steps,
|
|
evaluation_strategy=("steps" if script_args.eval_steps is not None else "epoch"),
|
|
eval_steps=script_args.eval_steps,
|
|
save_strategy="epoch",
|
|
gradient_checkpointing=script_args.gradient_checkpointing,
|
|
ddp_find_unused_parameters=False,
|
|
)
|
|
|
|
data_collator = GPTRewardDataCollatorWithPadding(tokenizer, max_length=script_args.seq_length, pad_to_multiple_of=8)
|
|
|
|
trainer = GPTRewardTrainer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
args=training_args,
|
|
train_dataset=train_dataset,
|
|
eval_dataset=eval_dataset,
|
|
max_length=script_args.seq_length,
|
|
data_collator=data_collator,
|
|
)
|
|
|
|
if script_args.mode == "train":
|
|
print("Training")
|
|
trainer.train()
|
|
trainer.evaluate()
|
|
|
|
print("Saving last checkpoint of the model")
|
|
trainer.save_model(script_args.output_dir)
|
|
|
|
output_dir = os.path.join(script_args.output_dir, "final_checkpoint")
|
|
trainer.model.save_pretrained(output_dir)
|
|
elif script_args.mode == "eval":
|
|
print("Evaluating")
|
|
# results = trainer.evaluate()
|
|
results = trainer.evaluate()
|
|
print(results)
|
|
elif script_args.mode == "relabel":
|
|
|
|
def relabel_with_preds(batch: Dict[str, List]):
|
|
relabel_batch = {
|
|
"prompt": [],
|
|
"chosen": [],
|
|
"rejected": [],
|
|
"pred_chosen": [],
|
|
"pred_rejected": [],
|
|
}
|
|
for prompt, chosen, rejected, pred_chosen, pred_rejected in zip(
|
|
batch["prompt"],
|
|
batch["chosen"],
|
|
batch["rejected"],
|
|
batch["pred_chosen"],
|
|
batch["pred_rejected"],
|
|
):
|
|
relabel_batch["prompt"].append(prompt)
|
|
if pred_chosen >= pred_rejected:
|
|
relabel_batch["chosen"].append(chosen)
|
|
relabel_batch["rejected"].append(rejected)
|
|
relabel_batch["pred_chosen"].append(pred_chosen)
|
|
relabel_batch["pred_rejected"].append(pred_rejected)
|
|
else:
|
|
relabel_batch["chosen"].append(rejected)
|
|
relabel_batch["rejected"].append(chosen)
|
|
relabel_batch["pred_chosen"].append(pred_rejected)
|
|
relabel_batch["pred_rejected"].append(pred_chosen)
|
|
|
|
return relabel_batch
|
|
|
|
relabel_dataset = DatasetDict()
|
|
for split, pred_dataset in [("train", train_dataset), ("test", eval_dataset)]:
|
|
if pred_dataset is None:
|
|
continue
|
|
trainer.accelerator.print(f"Prediction {split}")
|
|
preds, _, metrics = trainer.predict(pred_dataset)
|
|
trainer.accelerator.print(f"metrics {metrics}")
|
|
|
|
if trainer.accelerator.is_local_main_process:
|
|
print("Relabelling Dataset and Saving")
|
|
ds_split = script_args.train_split if split == "train" else script_args.eval_split
|
|
dataset = load_dataset(script_args.dataset_name, split=ds_split)
|
|
dataset = dataset.add_column("pred_chosen", preds[:, 0])
|
|
dataset = dataset.add_column("pred_rejected", preds[:, 1])
|
|
|
|
dataset = dataset.map(relabel_with_preds, batched=True)
|
|
|
|
dataset._info.description = f"{script_args.dataset_name} relabelled with {script_args.model_name}"
|
|
relabel_dataset[split] = dataset
|
|
|
|
if trainer.accelerator.is_local_main_process:
|
|
print("Saving")
|
|
relabel_dataset.save_to_disk(script_args.output_dir)
|
|
print("Pushing")
|
|
relabel_dataset.push_to_hub(os.path.basename(script_args.output_dir))
|
|
elif script_args.mode == "predict":
|
|
relabel_dataset = DatasetDict()
|
|
for split, pred_dataset in [("train", train_dataset), ("test", eval_dataset)]:
|
|
if pred_dataset is None:
|
|
continue
|
|
trainer.accelerator.print(f"Prediction {split}")
|
|
preds, _, metrics = trainer.predict(pred_dataset)
|
|
trainer.accelerator.print(f"metrics {metrics}")
|
|
|
|
if trainer.accelerator.is_local_main_process:
|
|
print("Relabelling Dataset and Saving")
|
|
ds_split = script_args.train_split if split == "train" else script_args.eval_split
|
|
dataset = load_dataset(script_args.dataset_name, split=ds_split)
|
|
model_basename = script_args.model_name.rsplit("/", 1)[-1]
|
|
dataset = dataset.add_column(f"pred_chosen_{model_basename}", preds[:, 0])
|
|
dataset = dataset.add_column(f"pred_rejected_{model_basename}", preds[:, 1])
|
|
|
|
dataset._info.description = f"{script_args.dataset_name} relabelled with {script_args.model_name}"
|
|
relabel_dataset[split] = dataset
|
|
|
|
if trainer.accelerator.is_local_main_process:
|
|
print("Saving")
|
|
relabel_dataset.save_to_disk(script_args.output_dir)
|
|
print("Pushing")
|
|
relabel_dataset.push_to_hub(os.path.basename(script_args.output_dir))
|
|
else:
|
|
raise Exception(f"incorrect mode {script_args.mode}")
|
|
# TODO this freezes for some reason
|
|
# for split, dataset in relabel_dataset.items():
|
|
# if trainer.accelerator.is_local_main_process:
|
|
# eval_dataset = prepare_dataset(script_args, dataset, tokenizer)
|
|
# trainer.accelerator.print(f"Re-evaluating relabel {split} dataset of size {len(dataset)}")
|
|
# trainer.accelerator.wait_for_everyone()
|
|
# results = trainer.evaluate(eval_dataset)
|
|
# trainer.accelerator.print(results)
|