code/dpo_costa.py

# 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, Literal, Optional

import bitsandbytes as bnb
import torch
from accelerate import Accelerator
from callbacks import GoldModelRewardCallback, PerplexityCallback, PerplexityGenCallback
from datasets import builder, concatenate_datasets, load_dataset
from peft import AutoPeftModelForCausalLM, LoraConfig, PeftConfig, get_peft_model, prepare_model_for_kbit_training
from scalar_rm_model import ScalarModel
from transformers import (
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    GenerationConfig,
    HfArgumentParser,
    TrainerCallback,
    TrainingArguments,
)
from transformers.trainer_utils import get_last_checkpoint

from trl import DPOTrainer


builder.has_sufficient_disk_space = lambda needed_bytes, directory=".": True


# Define and parse arguments.
@dataclass
class ScriptArguments:
    """
    The arguments for the DPO training script.
    """

    # data parameters
    dataset_name: Optional[str] = field(
        default="mnoukhov/openai_summarize_comparisons_tldrprompt_relabel1b", metadata={"help": "the dataset name"}
    )
    train_split: Optional[str] = field(default="train", metadata={"help": "the dataset split to train on"})
    eval_split: Optional[str] = field(
        default="test", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
    )
    beta: Optional[float] = field(default=0.1, metadata={"help": "the beta parameter for DPO loss"})

    pseudo_dataset_name: Optional[str] = field(default=None, metadata={"help": "the dataset name"})
    pseudo_dataset_split: Optional[str] = field(default="train", metadata={"help": "the dataset name"})
    prompt_field: Optional[str] = field(default="prompt")

    # model parameters
    model_name: Optional[str] = field(default="gpt2", metadata={"help": "the model name"})
    model_revision: Optional[str] = field(default=None, metadata={"help": "the model name"})
    tokenizer_name: Optional[str] = field(default=None, metadata={"help": "the model name"})
    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_model: 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."
        },
    )
    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_peft: Optional[bool] = field(default=True, metadata={"help": "Wether to use PEFT or not to train adapters"})
    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"})

    # training parameters
    optimizer_type: Optional[str] = field(default="adamw_torch", metadata={"help": "the optimizer type"})
    warmup_steps: Optional[int] = field(default=150)
    learning_rate: Optional[float] = field(default=1e-3, metadata={"help": "optimizer learning rate"})
    per_device_train_batch_size: Optional[int] = field(default=4, metadata={"help": "batch size per device"})
    per_device_eval_batch_size: Optional[int] = field(default=8, metadata={"help": "batch size per device"})
    gradient_accumulation_steps: Optional[int] = field(
        default=1, metadata={"help": "the number of gradient accumulation steps"}
    )
    max_length: Optional[int] = field(default=560, metadata={"help": "max length of each sample"})
    max_prompt_length: Optional[int] = field(default=512, metadata={"help": "max length of each sample's prompt"})
    max_target_length: Optional[int] = field(
        default=48, metadata={"help": "Only used for encoder decoder model. Max target of each sample's prompt"}
    )
    num_train_epochs: Optional[int] = field(default=1, metadata={"help": "the number of training epochs"})
    max_steps: Optional[int] = field(default=-1)
    gradient_checkpointing: Optional[bool] = field(
        default=False, metadata={"help": "whether to use gradient checkpointing"}
    )

    # instrumentation
    seed: Optional[int] = field(default=0)
    output_dir: Optional[str] = field(default="results", metadata={"help": "the output directory"})
    logging_steps: Optional[int] = field(default=100, metadata={"help": "the number of update steps between two logs"})
    log_n_samples_during_eval: Optional[int] = field(default=100)
    eval_steps: Optional[float] = field(default=None, metadata={"help": "the number of steps to eval at"})
    save_steps: Optional[float] = field(default=1000, metadata={"help": "the number of steps to save at"})
    save_strategy: Optional[str] = field(default="steps")
    report_to: Optional[str] = field(
        default="wandb",
        metadata={
            "help": 'The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`,'
            '`"comet_ml"`, `"mlflow"`, `"neptune"`, `"tensorboard"`,`"clearml"` and `"wandb"`. '
            'Use `"all"` to report to all integrations installed, `"none"` for no integrations.'
        },
    )
    # debug argument for distributed training
    ignore_bias_buffers: Optional[bool] = field(
        default=False,
        metadata={
            "help": "fix for DDP issues with LM bias/mask buffers - invalid scalar type,`inplace operation. See"
            "https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992"
        },
    )
    push_to_hub: Optional[bool] = field(default=False)
    push_to_hub_organization: Optional[str] = field(default=None)

    # gold model
    gold_eval: Literal["full", "gen", "ppl", "none"] = field(default="full")
    gold_model_name: str = field(default=None, metadata={"help": "the gold reward model name"})
    gold_model_revision: Optional[str] = field(default=None, metadata={"help": "the model name"})
    gold_in_8bit: Optional[bool] = field(default=False, metadata={"help": "gold the model in 8 bits precision"})
    gold_in_4bit: Optional[bool] = field(default=False, metadata={"help": "gold the model in 4 bits precision"})
    gold_bf16: Optional[bool] = field(
        default=False,
    )
    gold_fp16: Optional[bool] = field(
        default=False,
    )
    generate_greedy: Optional[bool] = field(default=True)
    gold_dataset_name: Optional[str] = field(
        default="CarperAI/openai_summarize_tldr", metadata={"help": "the dataset name"}
    )
    gold_eval_split: Optional[str] = field(default="valid")
    gold_prompt_field: Optional[str] = field(default="prompt")
    gold_target_field: Optional[str] = field(default="label")
    gold_load_and_unload: Optional[str] = field(default=False)
    mode: Optional[str] = field(default="train")
    eval_first_step: Optional[bool] = field(default=True)
    strip_prompt: Optional[bool] = field(default=True)


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:
        dtype = torch.bfloat16
    elif args.fp16_model:
        dtype = torch.float16
    else:
        dtype = torch.float32

    tokenizer_name = args.tokenizer_name

    if "adapter" in args.model_name:
        model_cls = AutoPeftModelForCausalLM
        config = PeftConfig.from_pretrained(args.model_name)
        if tokenizer_name is None:
            tokenizer_name = config.base_model_name_or_path
    else:
        model_cls = AutoModelForCausalLM
        if tokenizer_name is None:
            tokenizer_name = args.model_name

    model = model_cls.from_pretrained(
        args.model_name,
        revision=args.model_revision,
        quantization_config=quantization_config,
        device_map=device_map,
        torch_dtype=dtype,
    )

    model.config.torch_dtype = dtype
    model.config.use_cache = not script_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=script_args.gradient_checkpointing)

    # we add `score` to the list of modules to save to
    # correctly save the score head.
    # set target modules to be query_key_value for Pythia
    if args.lora_all_linear:
        modules = find_all_linear_names(args, model)
    else:
        modules = None

    if args.use_peft:
        modules_to_save = ["lm_head"]
        peft_config = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=modules,
            modules_to_save=modules_to_save,
        )

        model = get_peft_model(model, peft_config)

        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 args.bf16:
    #     for name, module in model.named_modules():
    #         if isinstance(module, LoraLayer):
    #             module = module.to(torch.bfloat16)
    #         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.bfloat16)

    # tokenizer_name = script_args.model_name if script_args.tokenizer_name is None else script_args.tokenizer_name
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

    if tokenizer_name.startswith("EleutherAI"):
        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
    elif 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 create_and_prepare_gold_model(args):
    if script_args.gold_in_8bit or script_args.gold_in_4bit:
        gold_quantization_config = BitsAndBytesConfig(
            load_in_8bit=script_args.gold_in_8bit, load_in_4bit=script_args.gold_in_4bit
        )
        gold_device_map = {"": Accelerator().local_process_index}
    else:
        gold_device_map = None
        gold_quantization_config = None

    if script_args.gold_bf16:
        torch_dtype = torch.bfloat16
    elif script_args.gold_fp16:
        torch_dtype = torch.float16
    else:
        torch_dtype = torch.float32

    if script_args.gold_model_name.startswith("vwxyzjn"):
        gold_model_cls = ScalarModel
    else:
        gold_model_cls = AutoModelForSequenceClassification

    gold_model = gold_model_cls.from_pretrained(
        script_args.gold_model_name,
        revision=script_args.gold_model_revision,
        quantization_config=gold_quantization_config,
        torch_dtype=torch_dtype,
        device_map=gold_device_map,
    )

    # if getattr(gold_model.config, "pad_token_id", None) is None:
    #     gold_model.config.pad_token_id = gold_model.config.eos_token_id

    return gold_model


def strip_prompt(examples):
    examples["prompt"] = [prompt.strip() for prompt in examples["prompt"]]

    return examples


def create_and_prepare_dataset(args):
    train_dataset = load_dataset(args.dataset_name, split=args.train_split)
    eval_dataset = load_dataset(args.dataset_name, split=args.eval_split)

    if args.prompt_field != "prompt":
        train_dataset = train_dataset.rename_column(args.prompt_field, "prompt")
        eval_dataset = eval_dataset.rename_column(args.prompt_field, "prompt")

    if args.pseudo_dataset_name is not None:
        all_train_datasets = [train_dataset]
        pseudo_dataset_names = args.pseudo_dataset_name.split(",")
        for ds_name in pseudo_dataset_names:
            dataset = load_dataset(ds_name, split=args.pseudo_dataset_split)
            if args.strip_prompt:
                dataset = dataset.map(strip_prompt, batched=True)
            all_train_datasets.append(dataset)

        train_dataset = concatenate_datasets(all_train_datasets)

    return train_dataset, eval_dataset


if __name__ == "__main__":
    parser = HfArgumentParser(ScriptArguments)
    script_args = parser.parse_args_into_dataclasses()[0]

    # 1. load a pretrained model
    model, tokenizer = create_and_prepare_model(script_args)

    if script_args.ignore_bias_buffers:
        # torch distributed hack
        model._ddp_params_and_buffers_to_ignore = [
            name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool
        ]

    train_dataset, eval_dataset = create_and_prepare_dataset(script_args)

    if script_args.push_to_hub:
        model_id = script_args.model_name.rsplit("/", 1)[-1] + "_" + os.getenv("WANDB_RUN_GROUP")
        hub_model_id = f"{script_args.push_to_hub_organization}/{model_id}"
        print(f"pushing model to {hub_model_id}")
    else:
        hub_model_id = None

    # 4. initialize training arguments:
    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,
        num_train_epochs=script_args.num_train_epochs,
        max_steps=script_args.max_steps,
        remove_unused_columns=False,
        gradient_accumulation_steps=script_args.gradient_accumulation_steps,
        learning_rate=script_args.learning_rate,
        evaluation_strategy="epoch" if script_args.eval_steps is None else "steps",
        save_strategy=script_args.save_strategy,
        logging_first_step=True,
        logging_steps=script_args.logging_steps,
        eval_steps=script_args.eval_steps,
        save_steps=script_args.save_steps,
        optim=script_args.optimizer_type,
        warmup_steps=script_args.warmup_steps,
        report_to=script_args.report_to,
        bf16=script_args.bf16,
        fp16=script_args.fp16,
        ddp_find_unused_parameters=(script_args.gradient_checkpointing),
        push_to_hub=script_args.push_to_hub,
        hub_model_id=hub_model_id,
    )

    # 5. initialize the DPO trainer
    dpo_trainer = DPOTrainer(
        model=model,
        args=training_args,
        beta=script_args.beta,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
        max_length=script_args.max_length,
        max_target_length=script_args.max_target_length,
        max_prompt_length=script_args.max_prompt_length,
    )

    # Gold Eval
    if script_args.gold_eval != "none":
        gold_eval_dataset = load_dataset(
            script_args.gold_dataset_name,
            split=script_args.gold_eval_split,
        )

        if script_args.strip_prompt:
            gold_eval_dataset = gold_eval_dataset.map(strip_prompt, batched=True)

        if script_args.generate_greedy:
            generation_config = GenerationConfig(
                max_new_tokens=script_args.max_target_length,
                do_sample=False,
                num_beams=1,
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.eos_token_id,
            )
        else:
            generation_config = GenerationConfig(
                max_new_tokens=script_args.max_target_length,
                min_length=-1,
                top_k=0.0,
                top_p=1.0,
                do_sample=True,
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.eos_token_id,
            )

        if script_args.gold_eval == "full":
            gold_model = create_and_prepare_gold_model(script_args)

            callback = GoldModelRewardCallback(
                training_args,
                gold_model,
                gold_eval_dataset,
                tokenizer,
                dpo_trainer.accelerator,
                script_args.max_length,
                script_args.max_prompt_length,
                script_args.gold_prompt_field,
                script_args.gold_target_field,
                script_args.gold_load_and_unload,
                script_args.log_n_samples_during_eval,
                generation_config,
            )
        else:
            if script_args.gold_eval == "gen":
                callback_cls = PerplexityGenCallback
            elif script_args.gold_eval == "ppl":
                callback_cls = PerplexityCallback
            else:
                raise NotImplementedError

            callback = callback_cls(
                args=training_args,
                dataset=gold_eval_dataset,
                tokenizer=tokenizer,
                accelerator=dpo_trainer.accelerator,
                max_length=script_args.max_length,
                max_prompt_length=script_args.max_prompt_length,
                prompt_field=script_args.gold_prompt_field,
                target_field=script_args.gold_target_field,
                log_n_samples_during_eval=script_args.log_n_samples_during_eval,
                generation_config=generation_config,
                hub_model_id=hub_model_id,
            )

        dpo_trainer.add_callback(callback)

    if script_args.eval_first_step:

        class EvaluateFirstStepCallback(TrainerCallback):
            def on_step_end(self, args, state, control, **kwargs):
                if state.global_step == 1:
                    control.should_evaluate = True

        dpo_trainer.add_callback(EvaluateFirstStepCallback())

    # 6. train
    if script_args.mode == "train":
        last_checkpoint = get_last_checkpoint(script_args.output_dir)
        dpo_trainer.train(resume_from_checkpoint=last_checkpoint)
    elif script_args.mode == "eval":
        print("evaluating")
        results = dpo_trainer.evaluate()
        print(results)
    elif script_args.mode == "relabel":

        def relabel_with_preds(batch: Dict[str, List]):
            relabel_batch = {
                "prompt": [],
                "chosen": [],
                "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)
                else:
                    relabel_batch["chosen"].append(rejected)
                    relabel_batch["rejected"].append(chosen)

            return relabel_batch

        dpo_trainer.accelerator.print(f"Prediction {script_args.eval_split}")
        preds, _, metrics = dpo_trainer.predict(eval_dataset)
        (
            chosen_rewards,
            rejected_rewards,
            policy_chosen_logps,
            policy_rejected_logps,
            reference_chosen_logps,
            reference_rejected_logps,
        ) = preds
        dpo_trainer.accelerator.print(f"metrics {metrics}")

        if dpo_trainer.accelerator.is_local_main_process:
            print("Relabelling Dataset and Saving")
            dataset = load_dataset(script_args.dataset_name, split=script_args.eval_split)
            dataset = dataset.add_column("pred_chosen", chosen_rewards)
            dataset = dataset.add_column("pred_rejected", rejected_rewards)

            relabel_dataset = dataset.map(
                relabel_with_preds,
                batched=True,
            )

            description = f"{script_args.dataset_name} relabelled with {script_args.model_name}"
            relabel_dataset._info.description = description

        if dpo_trainer.accelerator.is_local_main_process:
            # print("Saving")
            # relabel_dataset.save_to_disk(script_args.output_dir)
            print("Pushing")
            # repo_id = f"MilaRLHF/{os.path.basename(script_args.output_dir)}"
            relabel_dataset.push_to_hub(os.path.basename(script_args.output_dir), split=script_args.eval_split)
            # relabel_dataset_card = DatasetCard.load(repo_id)
            # relabel_dataset_card.text = description
            # relabel_dataset_card.push_to_hub(repo_id)
    elif script_args.mode == "predict":
        dpo_trainer.accelerator.print(f"Prediction {script_args.eval_split}")
        preds, _, metrics = dpo_trainer.predict(eval_dataset)
        (
            chosen_rewards,
            rejected_rewards,
            policy_chosen_logps,
            policy_rejected_logps,
            reference_chosen_logps,
            reference_rejected_logps,
        ) = preds
        dpo_trainer.accelerator.print(f"metrics {metrics}")

        if dpo_trainer.accelerator.is_local_main_process:
            print("Relabelling Dataset and Saving")
            dataset = load_dataset(script_args.dataset_name, split=script_args.eval_split)
            model_basename = script_args.model_name.rsplit("/", 1)[-1]
            dataset = dataset.add_column(f"pred_chosen_{model_basename}", chosen_rewards)
            dataset = dataset.add_column(f"pred_rejected_{model_basename}", rejected_rewards)

        if dpo_trainer.accelerator.is_local_main_process:
            # print("Saving")
            # relabel_dataset.save_to_disk(script_args.output_dir)
            print("Pushing")
            dataset.push_to_hub(os.path.basename(script_args.output_dir), split=script_args.eval_split)