# 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)