Self-Alignment with Principle-Following Reward Models
# Load the model.
# Note: It can take a while to download LLaMA and add the adapter modules.

import torch
import fire

import datetime
import os
from threading import Event, Thread
from uuid import uuid4

import gradio as gr
import requests

from peft import PeftModel
from transformers import (
    BitsAndBytesConfig,
    AutoModelForCausalLM,
    LlamaTokenizerFast,
    StoppingCriteria,
    StoppingCriteriaList,
    TextIteratorStreamer,
)

torch.backends.cuda.matmul.allow_tf32 = True


def main(
    model_name: str = "/path/to/llama-2-70b-hf",
    adapters_name: str = "/path/to/adapter_model",
    tokenizer_name: str = "TheBloke/dromedary-65b-lora-HF",  # a random loadable llama tokenizer
    max_queue_size: int = 16,
    sharable_link: bool = True,
):
    print(f"Starting to load the model {model_name} into memory")

    m = AutoModelForCausalLM.from_pretrained(
        model_name,
        load_in_4bit=True,
        torch_dtype=torch.bfloat16,
        device_map={"": 0},
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
        ),
    )
    m = PeftModel.from_pretrained(m, adapters_name, is_trainable=False)

    start_message = (
        "# Dromedary"
        "\n\n## System Overview"
        "\n\nConsider an AI assistant whose codename is Dromedary, developed by the Self-Align team. "
        "Dromedary is trained on data from before Sept-2022, and it endeavors to be a helpful, ethical and reliable assistant."
        "\n\n## User Conversation\n\n"
    )

    tok = LlamaTokenizerFast.from_pretrained(tokenizer_name)

    print(f"Successfully loaded the model {model_name} into memory")

    # Setup the gradio Demo.

    stop_tokens = ["### User"]

    class StopOnTokens(StoppingCriteria):
        def __call__(
            self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
        ) -> bool:
            decoded_tokens = tok.batch_decode(input_ids, skip_special_tokens=True)

            for stop_token in stop_tokens:
                if decoded_tokens[0].strip().endswith(stop_token):
                    return True
            return False

    def convert_history_to_text(history):
        text = start_message + "".join(
            [
                "".join(
                    [
                        f"### User\n{item[0]}\n\n",
                        f"### Dromedary\n{item[1]}\n\n",
                    ]
                )
                for item in history[:-1]
            ]
        )
        text += "".join(
            [
                "".join(
                    [
                        f"### User\n{history[-1][0]}\n\n",
                        f"### Dromedary\n{history[-1][1]}",
                    ]
                )
            ]
        )
        return text

    def log_conversation(conversation_id, history, messages, generate_kwargs):
        logging_url = os.getenv("LOGGING_URL", None)
        if logging_url is None:
            return

        timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S")

        data = {
            "conversation_id": conversation_id,
            "timestamp": timestamp,
            "history": history,
            "messages": messages,
            "generate_kwargs": generate_kwargs,
        }

        try:
            requests.post(logging_url, json=data)
        except requests.exceptions.RequestException as e:
            print(f"Error logging conversation: {e}")

    def user(message, history):
        # Append the user's message to the conversation history
        return "", history + [[message, ""]]

    def bot(
        history, temperature, top_p, max_new_tokens, repetition_penalty, conversation_id
    ):
        # Initialize a StopOnTokens object
        stop = StopOnTokens()

        # Construct the input message string for the model by concatenating the current system message and conversation history
        messages = convert_history_to_text(history)

        print(f"history:")
        print(messages)

        # Tokenize the messages string
        input_ids = tok(messages, return_tensors="pt").input_ids
        input_ids = input_ids.to(m.device)
        streamer = TextIteratorStreamer(
            tok, timeout=10.0, skip_prompt=True, skip_special_tokens=True
        )
        generate_kwargs = dict(
            input_ids=input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=temperature > 0.0,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            streamer=streamer,
            stopping_criteria=StoppingCriteriaList([stop]),
        )

        stream_complete = Event()

        def generate_and_signal_complete():
            m.generate(**generate_kwargs)
            stream_complete.set()

        def log_after_stream_complete():
            stream_complete.wait()
            log_conversation(
                conversation_id,
                history,
                messages,
                {
                    "max_new_tokens": max_new_tokens,
                    "top_p": top_p,
                    "temperature": temperature,
                    "repetition_penalty": repetition_penalty,
                },
            )

        t1 = Thread(target=generate_and_signal_complete)
        t1.start()

        t2 = Thread(target=log_after_stream_complete)
        t2.start()

        # Initialize an empty string to store the generated text
        partial_text = ""
        for new_text in streamer:
            partial_text += new_text
            history[-1][1] = partial_text.split("### User")[0]
            yield history

        print(f"output:")
        print(history[-1][1])

    def get_uuid():
        return str(uuid4())

    with gr.Blocks(
        theme=gr.themes.Soft(),
        css=".disclaimer {font-variant-caps: all-small-caps;}",
    ) as demo:
        conversation_id = gr.State(get_uuid)
        gr.Markdown(
            """Dromedary-2 Demo
    """
        )
        chatbot = gr.Chatbot().style(height=500)
        with gr.Row():
            with gr.Column():
                msg = gr.Textbox(
                    label="Chat Message Box",
                    placeholder="Chat Message Box",
                    show_label=False,
                ).style(container=False)
            with gr.Column():
                with gr.Row():
                    submit = gr.Button("Submit")
                    # stop = gr.Button("Stop")
                    clear = gr.Button("Clear")
        with gr.Row():
            with gr.Accordion("Advanced Options:", open=False):
                with gr.Row():
                    with gr.Column():
                        with gr.Row():
                            temperature = gr.Slider(
                                label="Temperature",
                                value=0.7,
                                minimum=0.0,
                                maximum=1.0,
                                step=0.1,
                                interactive=True,
                                info="Higher values produce more diverse outputs",
                            )
                    with gr.Column():
                        with gr.Row():
                            top_p = gr.Slider(
                                label="Top-p (nucleus sampling)",
                                value=0.9,
                                minimum=0.0,
                                maximum=1,
                                step=0.01,
                                interactive=True,
                                info=(
                                    "Sample from the smallest possible set of tokens whose cumulative probability "
                                    "exceeds top_p. Set to 1 to disable and sample from all tokens."
                                ),
                            )
                    with gr.Column():
                        with gr.Row():
                            max_new_tokens = gr.Slider(
                                label="Response Length",
                                value=768,
                                minimum=64,
                                maximum=1024,
                                step=64,
                                interactive=True,
                                info="Sample from a shortlist of top-k tokens — 0 to disable and sample from all tokens.",
                            )
                    with gr.Column():
                        with gr.Row():
                            repetition_penalty = gr.Slider(
                                label="Repetition Penalty",
                                value=1.0,
                                minimum=1.0,
                                maximum=2.0,
                                step=0.1,
                                interactive=True,
                                info="Penalize repetition — 1.0 to disable.",
                            )
        with gr.Row():
            gr.Markdown(
                "Disclaimer: The model can produce factually incorrect output, and should not be relied on to produce "
                "factually accurate information. The model was trained on various public datasets; while great efforts "
                "have been taken to clean the pretraining data, it is possible that this model could generate lewd, "
                "biased, or otherwise offensive outputs.",
                elem_classes=["disclaimer"],
            )
        with gr.Row():
            gr.Markdown(
                "[Privacy policy](https://gist.github.com/samhavens/c29c68cdcd420a9aa0202d0839876dac)",
                elem_classes=["disclaimer"],
            )

        submit_event = msg.submit(
            fn=user,
            inputs=[msg, chatbot],
            outputs=[msg, chatbot],
            queue=False,
        ).then(
            fn=bot,
            inputs=[
                chatbot,
                temperature,
                top_p,
                max_new_tokens,
                repetition_penalty,
                conversation_id,
            ],
            outputs=chatbot,
            queue=True,
        )
        submit_click_event = submit.click(
            fn=user,
            inputs=[msg, chatbot],
            outputs=[msg, chatbot],
            queue=False,
        ).then(
            fn=bot,
            inputs=[
                chatbot,
                temperature,
                top_p,
                max_new_tokens,
                repetition_penalty,
                conversation_id,
            ],
            outputs=chatbot,
            queue=True,
        )
        # stop.click(
        #     fn=None,
        #     inputs=None,
        #     outputs=None,
        #     cancels=[submit_event, submit_click_event],
        #     queue=False,
        # )
        clear.click(lambda: None, None, chatbot, queue=False)

    demo.queue(max_size=max_queue_size, concurrency_count=1)

    # Launch your Dromedary-2 Demo!
    demo.launch(share=sharable_link)


if __name__ == "__main__":
    fire.Fire(main)

inference/demo.py
IBM/SALMON
# Copyright 2023 The Alpaca Team
#
# 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 argparse
import glob
import os
import random
from typing import (
    Callable,
    Dict,
    Optional,
    Sequence,
    Union,
    Mapping,
    Any,
)

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

Numeric = Union[int, float]


def zip_(*args: Sequence):
    """Assert sequences of same length before zipping."""
    if len(args) == 0:
        return []
    assert alleq(args, lambda x, y: len(x) == len(y))
    return zip(*args)


def mean(*seqs: Sequence[Numeric]) -> Union[Numeric, Sequence[Numeric]]:
    singleton = len(seqs) == 1
    means = [float(np.mean(seq)) for seq in seqs]
    return means[0] if singleton else means


def alleq(l: Sequence, f: Optional[Callable] = lambda x, y: x == y):
    """Check all arguments in a sequence are equal according to a given criterion.

    Args:
        f: A bi-variate boolean function.
        l: A list/tuple.

    Returns:
        True if everything is equal; otherwise False.
    """
    return all(f(l[0], li) for li in l[1:])


def flatten_dict(nested, sep=".", postprocess_fn=lambda *args: args):
    def rec(nest, prefix, into):
        for k, v in nest.items():
            if sep in k:
                raise ValueError(f"separator '{sep}' not allowed to be in key '{k}'")
            if isinstance(v, dict):  # collections.Mapping fails in py3.10.
                rec(v, prefix + k + sep, into)
            else:
                v = postprocess_fn(v)
                into[prefix + k] = v

    flat = {}
    rec(nested, "", flat)
    return flat


def unpack_dict(
    d: Dict, keys: Sequence[str], return_type: type = tuple
) -> Union[Sequence, Dict]:
    if return_type in (tuple, list):
        return return_type(d[key] for key in keys)
    elif return_type == dict:
        return {key: d[key] for key in keys}
    else:
        raise ValueError(f"Unknown return_type: {return_type}")


def merge_dict(dicts: Sequence[dict], merge_fn: Callable = lambda *args: args) -> dict:
    """Merge a sequence of dicts (with the same set of keys) into a single dict."""
    if len(dicts) == 0:
        return dict()
    return {key: merge_fn([dict_[key] for dict_ in dicts]) for key in dicts[0].keys()}


def prepare_inputs(
    data: Union[torch.Tensor, Any], device: Union[str, int, torch.device]
) -> Union[torch.Tensor, Any]:
    if isinstance(data, Mapping):
        return type(data)(
            {k: prepare_inputs(v, device) for k, v in data.items()}
        )  # noqa
    elif isinstance(data, (tuple, list)):
        return type(data)(prepare_inputs(v, device) for v in data)
    elif isinstance(data, torch.Tensor):
        return data.to(device)  # This can break with deepspeed.
    return data


def compute_logprobs(
    logits: torch.Tensor, labels: torch.Tensor, ignore_index: int
) -> torch.Tensor:
    """Compute per-token logprobs, zeroing out places with ignore_index (padding)."""
    return -F.cross_entropy(
        logits.permute(0, 2, 1), labels, reduction="none", ignore_index=ignore_index
    )


def pad(
    inputs: torch.Tensor,
    target_size: Union[torch.Size, Sequence[int]],
    value=0.0,
    left=True,
):
    current_size = inputs.size()
    diffs = tuple(ti - ci for ti, ci in zip_(target_size, current_size))
    pad_params = []
    for diff in diffs:
        pad_params = ([diff, 0] if left else [0, diff]) + pad_params
    res = F.pad(inputs, pad=pad_params, value=value)
    return res


def left_pad(
    inputs: torch.Tensor, target_size: Union[torch.Size, Sequence[int]], value=0.0
):
    return pad(inputs=inputs, target_size=target_size, value=value, left=True)


def right_pad(
    inputs: torch.Tensor, target_size: Union[torch.Size, Sequence[int]], value=0.0
):
    return pad(inputs=inputs, target_size=target_size, value=value, left=False)


def manual_seed(args_or_seed: Union[int, argparse.Namespace], fix_cudnn=False):
    if hasattr(args_or_seed, "seed"):
        args_or_seed = args_or_seed.seed
    random.seed(args_or_seed)
    np.random.seed(args_or_seed)
    torch.manual_seed(args_or_seed)
    torch.cuda.manual_seed_all(args_or_seed)
    os.environ["PYTHONHASHSEED"] = str(args_or_seed)
    if fix_cudnn:
        torch.backends.cudnn.deterministic = True  # noqa
        torch.backends.cudnn.benchmark = False  # noqa


def make_meta_prompts(meta_prompt_pattern: str):
    meta_prompt_files = glob.glob(meta_prompt_pattern)
    print(f"Found {len(meta_prompt_files)} meta prompts: {meta_prompt_files}")

    meta_prompts = []
    for meta_prompt_file in meta_prompt_files:
        with open(meta_prompt_file, "r", encoding="utf-8") as f:
            meta_prompt = f.readlines()
        meta_prompt = "".join(meta_prompt).strip()
        meta_prompts.append(meta_prompt)
    return meta_prompts


class InfiniteLoader(object):
    """Wraps an existing loader so that it outputs stuff indefinitely; useful for semi-supervised learning."""

    def __init__(self, loader: DataLoader):
        super(InfiniteLoader, self).__init__()
        self.loader = loader
        self.iterator = iter(loader)

    def __next__(self):
        try:
            return next(self.iterator)
        except StopIteration:
            self.iterator = iter(self.loader)
            return next(self.iterator)

training/data_utils/common_utils.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
#
# 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 dataclasses
from typing import Callable, Dict, Optional, List, Sequence

import logging
import pandas as pd

import torch
from torch.utils.data import Dataset

import transformers
import datasets

import data_utils.common_utils as utils


logger = logging.getLogger(__name__)


BASE_PROMPT_DICT = {
    "prompt_input": "{instruction}\n\n{input}",
    "prompt_no_input": "{instruction}",
}


def format_input_and_prompt(
    example: Dict[str, str],
    meta_prompts: List[str],
) -> str:
    assert (
        "instruction" in example and "input" in example
    ), "Internal error: example missing required keys."

    if "example_id" in example:
        total_meta_prompt = len(meta_prompts)
        meta_prompt = meta_prompts[int(example["example_id"]) % total_meta_prompt]
    else:
        meta_prompt = meta_prompts[0]

    if example.get("input", "") != "":
        prompt_format = BASE_PROMPT_DICT["prompt_input"]
    else:
        prompt_format = BASE_PROMPT_DICT["prompt_no_input"]

    formatted_input = prompt_format.format(**example)
    meta_prompt = meta_prompt.split("{Output}")[0]

    formatted_prompt = meta_prompt.format(Input=formatted_input)
    return formatted_input, formatted_prompt


class QueryResponseDataset(Dataset):
    """Dataset that emits tokenized left-padded queries."""

    def __init__(
        self,
        df: pd.DataFrame,
        meta_prompts: List[str],
        tokenizer: transformers.PreTrainedTokenizer,
        query_len: int,
        df_postprocessor: Optional[Callable] = None,
    ):
        super(QueryResponseDataset, self).__init__()

        if df_postprocessor is not None:
            df = df_postprocessor(df)
        list_dict_data = df.to_dict(orient="records")

        # prompts are strings; queries are tensors.
        inputs_and_prompts = [
            format_input_and_prompt(example=dict_data, meta_prompts=meta_prompts)
            for dict_data in list_dict_data
        ]
        formated_inputs, prompts = zip(*inputs_and_prompts)

        length_bonus = [
            dict_data.get("length_bonus", 1.0) for dict_data in list_dict_data
        ]

        # For question_type:
        # 0: general
        # 1: reasoning
        # 2: red-teaming
        question_types = [
            dict_data.get("question_type", 0) for dict_data in list_dict_data
        ]

        pure_queries = [
            tokenizer(
                formated_input, return_tensors="pt", truncation=False
            ).input_ids.squeeze(dim=0)
            for formated_input in formated_inputs
        ]

        queries = [
            tokenizer(prompt, return_tensors="pt", truncation=False).input_ids.squeeze(
                dim=0
            )
            for prompt in prompts
        ]

        filtered_inputs = []
        filtered_queries = []
        filtered_length_bonus = []
        filtered_question_types = []

        for pure_query, query, ex_length_bonus, ex_question_type in zip(
            pure_queries, queries, length_bonus, question_types
        ):
            if len(query) <= query_len:
                filtered_inputs.append(pure_query)
                filtered_queries.append(query)
                filtered_length_bonus.append(ex_length_bonus)
                filtered_question_types.append(ex_question_type)

        logger.warning(
            f"Filtered out {len(queries) - len(filtered_queries)} instances out of {len(queries)} that "
            f"exceed length limit. These examples are not used for training, but will still be used in evaluation. "
        )

        pure_queries = torch.stack(
            [
                utils.left_pad(
                    pure_query, target_size=(query_len,), value=tokenizer.pad_token_id
                )
                for pure_query in filtered_inputs
            ]
        )

        queries = torch.stack(
            [
                utils.left_pad(
                    query, target_size=(query_len,), value=tokenizer.pad_token_id
                )
                for query in filtered_queries
            ]
        )

        self.length_bonus = torch.tensor(
            filtered_length_bonus, dtype=torch.float32
        ).reshape(-1, 1)
        self.question_types = torch.tensor(
            filtered_question_types, dtype=torch.long
        ).reshape(-1, 1)
        self.pure_queries = pure_queries
        self.queries = queries
        self.query_attn_masks = queries.ne(tokenizer.pad_token_id).long()

        assert self.pure_queries.shape[0] == self.queries.shape[0]
        assert self.pure_queries.shape[0] == self.query_attn_masks.shape[0]
        assert self.pure_queries.shape[0] == self.length_bonus.shape[0]
        assert self.pure_queries.shape[0] == self.question_types.shape[0]

        # Auxiliary data.
        self.prompts = prompts
        self.list_dict_data = list_dict_data

    def __getitem__(self, i):
        return_dict = dict(
            pure_queries=self.pure_queries[i],
            queries=self.queries[i],
            query_attn_masks=self.query_attn_masks[i],
            length_bonus=self.length_bonus[i],
            question_types=self.question_types[i],
        )
        return return_dict

    def __len__(self):
        return len(self.queries)


@dataclasses.dataclass
class DataCollatorForQueryResponseDataset(object):
    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        return {
            key: torch.stack([instance[key] for instance in instances])
            for key in instances[0].keys()
        }


def make_rl_data_module(
    tokenizer: transformers.PreTrainedTokenizer,
    data_args,
    training_args,
):
    # prompt_dict = utils.jload(data_args.prompt_dict_path)

    policy_meta_prompts = utils.make_meta_prompts(data_args.policy_meta_prompt_pattern)

    if data_args.dataset_path.endswith("json"):
        train_instructions = datasets.load_dataset(
            "json", data_files=data_args.dataset_path
        )
    else:
        train_instructions = datasets.load_dataset(
            data_args.dataset_path, data_args.dataset_name
        )
    train_df = pd.concat(
        [pd.DataFrame(train_instructions[split]) for split in data_args.train_splits]
    )

    eval_instructions = datasets.load_dataset(
        data_args.eval_dataset_path, data_args.eval_dataset_name
    )
    eval_df = pd.concat(
        [pd.DataFrame(eval_instructions[split]) for split in data_args.eval_splits]
    )

    train_dataset = QueryResponseDataset(
        df=train_df,
        meta_prompts=policy_meta_prompts,
        tokenizer=tokenizer,
        query_len=training_args.query_len,
    )
    eval_dataset = QueryResponseDataset(
        df=eval_df,
        meta_prompts=policy_meta_prompts,
        tokenizer=tokenizer,
        query_len=training_args.query_len,
    )
    return dict(
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=DataCollatorForQueryResponseDataset(),
    )

training/data_utils/data_utils_ppo.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
#
# 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.

from dataclasses import dataclass
import json
import logging
from typing import Callable, Optional, Dict, Sequence, List, Tuple, Union

import tqdm

import einops
import pandas as pd
import torch
import transformers
from torch.nn.utils.rnn import pad_sequence
from datasets import load_dataset
from datasets import Dataset as HFDataset
from torch.utils.data import Dataset

import data_utils.common_utils as utils

IGNORE_INDEX = -100


logger = logging.getLogger(__name__)


BASE_PROMPT_DICT = {
    "prompt_input": "{instruction}\n\n{input}",
    "prompt_no_input": "{instruction}",
}


def format_prompt(
    example: Dict[str, str],
    meta_prompts: List[str],
) -> str:
    assert (
        "instruction" in example and "input" in example
    ), "Internal error: example missing required keys."

    if "example_id" in example:
        total_meta_prompt = len(meta_prompts)
        meta_prompt = meta_prompts[int(example["example_id"]) % total_meta_prompt]
    else:
        meta_prompt = meta_prompts[0]

    if example.get("input", "") != "":
        prompt_format = BASE_PROMPT_DICT["prompt_input"]
    else:
        prompt_format = BASE_PROMPT_DICT["prompt_no_input"]

    formatted_input = prompt_format.format(**example)
    meta_prompt = meta_prompt.split("{Output}")[0]

    formatted_prompt = meta_prompt.format(Input=formatted_input)
    return formatted_prompt


def format_output(
    example: dict,
    meta_prompts: List[str],
    eos_token: Optional[str] = None,
    output_key="output",
) -> str:
    if eos_token is None:
        eos_token = ""

    if "example_id" in example:
        total_meta_prompt = len(meta_prompts)
        meta_prompt = meta_prompts[int(example["example_id"]) % total_meta_prompt]
    else:
        meta_prompt = meta_prompts[0]

    meta_prompt = "{Output}" + meta_prompt.split("{Output}")[1]
    formatted_output = meta_prompt.format(Output=example[output_key]) + eos_token
    return formatted_output


def format_full_prompt(
    example: Dict[str, str],
    meta_prompts: List[str],
    principle_collection: Dict[str, str],
    tokenizer: transformers.PreTrainedTokenizer,
    eos_token: Optional[str] = None,
    query_len: Optional[int] = None,
    response_len: Optional[int] = None,
    output_key: str = "output",
    max_principles: Optional[int] = 5,
) -> str:
    if eos_token is None:
        eos_token = ""
    assert (
        "instruction" in example and "input" in example
    ), "Internal error: example missing required keys."

    if "example_id" in example:
        total_meta_prompt = len(meta_prompts)
        meta_prompt = meta_prompts[int(example["example_id"]) % total_meta_prompt]
    else:
        meta_prompt = meta_prompts[0]

    if "{Dimensions}" in meta_prompt:
        assert principle_collection is not None

        if "dimensions" in example:
            dimensions = eval(example["dimensions"])
        else:
            dimensions = list(principle_collection.keys())[:max_principles]

        if "negative_definitions" in example:
            negative_definitions = eval(example["negative_definitions"])
        else:
            negative_definitions = [0] * len(dimensions)

        principle_str = ""
        for neg_flag, dim in zip(negative_definitions, dimensions):
            if neg_flag:
                principle_str += f"- {principle_collection[dim + '_neg']}\n"
            else:
                principle_str += f"- {principle_collection[dim]}\n"
        principle_str = principle_str.strip()
        meta_prompt = meta_prompt.replace("{Dimensions}", principle_str)

    if example.get("input", "") != "":
        prompt_format = BASE_PROMPT_DICT["prompt_input"]
    else:
        prompt_format = BASE_PROMPT_DICT["prompt_no_input"]

    formatted_input = prompt_format.format(**example)

    if query_len is not None:
        ori_truncation_side = tokenizer.truncation_side
        tokenizer.truncation_side = "left"
        formatted_input_tokens = tokenizer(
            text=formatted_input,
            max_length=query_len,
            add_special_tokens=False,
            truncation=True,
            padding=False,
            return_tensors="pt",
        )
        formatted_input = tokenizer.batch_decode(
            formatted_input_tokens["input_ids"], skip_special_tokens=True
        )[0].strip()
        tokenizer.truncation_side = ori_truncation_side

    if response_len is not None:
        ori_truncation_side = tokenizer.truncation_side
        tokenizer.truncation_side = "right"
        formatted_output_tokens = tokenizer(
            text=example[output_key],
            max_length=response_len,
            add_special_tokens=False,
            truncation=True,
            padding=False,
            return_tensors="pt",
        )
        formatted_output = tokenizer.batch_decode(
            formatted_output_tokens["input_ids"], skip_special_tokens=True
        )[0].strip()
        tokenizer.truncation_side = ori_truncation_side
    else:
        formatted_output = example[output_key]

    formatted_prompt = (
        meta_prompt.format(
            Input=formatted_input,
            Output=formatted_output,
        )
        + eos_token
    )

    return formatted_prompt


def format_prompt_with_data_frame(
    df: pd.DataFrame,
    meta_prompts: dict,
    df_postprocessor: Optional[Callable] = None,
    return_dict=False,
):
    if df_postprocessor is not None:
        df = df_postprocessor(df)
    list_dict_data = df.to_dict(orient="records")

    prompts = [format_prompt(example, meta_prompts) for example in list_dict_data]
    metadata = {"prompt_dict": meta_prompts}

    if return_dict:
        return dict(prompts=prompts, list_dict_data=list_dict_data, metadata=metadata)
    return prompts, list_dict_data, metadata


def _tokenize_fn(
    strings: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
    use_data_frame: bool = True,
) -> dict:
    """Tokenize a list of strings and return the tokenized content as well metadata (e.g., truncation statistics)."""
    padding = getattr(tokenizer, "padding", "max_length")
    return_overflowing_tokens = False

    if use_data_frame:
        strings_ds = HFDataset.from_dict({"full_prompt": strings})
    else:
        strings_ds = strings

    def _tokenize(x):
        tokenized_text = tokenizer(
            x["full_prompt"],
            return_tensors="pt",
            padding=padding,
            max_length=tokenizer.model_max_length,
            truncation=True,
            return_overflowing_tokens=return_overflowing_tokens,
        )
        return {
            "input_ids": tokenized_text.input_ids,
            "attention_mask": tokenized_text.attention_mask,
            "input_length": tokenized_text.attention_mask.sum(),
        }

    tokenized_list = strings_ds.map(
        lambda x: _tokenize(x),
        num_proc=16,
    )
    # logger.warning(f"Example of tokenized text: {tokenized_list[0]['tokenized_text']}")

    tokenized_list.set_format(
        type="torch",
        columns=["input_ids", "attention_mask", "input_length"],
    )
    if padding == "max_length":
        input_ids = labels = torch.cat(
            [
                tokenized["input_ids"]
                for tokenized in tqdm.tqdm(
                    tokenized_list,
                    desc="Concatenating input_ids",
                )
            ]
        )
    else:  # "longest"
        input_ids = labels = [
            tokenized["input_ids"][0]
            for tokenized in tqdm.tqdm(
                tokenized_list,
                desc="Concatenating input_ids",
            )
        ]

    logger.warning(
        "You are using a `transformers` version that does not support `return_overflowing_tokens=True`. "
        "The tokenization metadata will not be recorded."
        "In order to see truncation statistics, please downgrade to `transformers<=4.26.1`."
    )
    input_ids_lens = labels_lens = [
        tokenized["input_length"].item()
        for tokenized in tqdm.tqdm(tokenized_list, desc="Computing lengths")
    ]
    num_truncated_tokens = num_truncated_examples = -1

    return dict(
        input_ids=input_ids,
        labels=labels,
        input_ids_lens=input_ids_lens,
        labels_lens=labels_lens,
        tokenization_metadata=dict(
            num_examples=len(tokenized_list),
            num_truncated_tokens=num_truncated_tokens,
            num_truncated_examples=num_truncated_examples,
            input_ids_avg_len=utils.mean(input_ids_lens),
            input_ids_max_len=max(input_ids_lens),
            input_ids_min_len=min(input_ids_lens),
            labels_avg_len=utils.mean(labels_lens),
            labels_max_len=max(labels_lens),
            labels_min_len=min(labels_lens),
            model_max_length=tokenizer.model_max_length,
        ),
    )


def preprocess_for_reward_modeling(
    data: Sequence[dict],
    meta_prompts: List[str],
    principle_collection: Dict[str, str],
    tokenizer: transformers.PreTrainedTokenizer,
    df_postprocessor: Optional[Callable] = None,
    end_sequence_with_eos: bool = False,
    query_len: Optional[int] = None,
    response_len: Optional[int] = None,
    use_data_frame: bool = True,
    verbose: bool = True,
) -> Dict[str, torch.Tensor]:
    if use_data_frame:
        df = pd.DataFrame(data)
        if df_postprocessor is not None:
            df = df_postprocessor(df)
        list_dict_data = df.to_dict(orient="records")
    else:
        list_dict_data = data

    index_0, index_1 = tuple(
        torch.full(
            size=(len(list_dict_data), 1), fill_value=fill_value, dtype=torch.long
        )
        for fill_value in (0, 1)
    )

    def _get_numeric_preference(example: dict):
        # 1 vs 2 is stored in table, but for modeling we use 0 vs 1; remap here.
        return {1: 0, 2: 1}[example["preference"]]

    choice = torch.tensor(
        [[_get_numeric_preference(dict_data)] for dict_data in list_dict_data]
    )

    def _get_text(example: dict, output_key: str):
        full_prompt = format_full_prompt(
            example,
            meta_prompts=meta_prompts,
            principle_collection=principle_collection,
            tokenizer=tokenizer,
            eos_token=tokenizer.eos_token if end_sequence_with_eos else None,
            output_key=output_key,
            query_len=query_len,
            response_len=response_len,
        )
        return full_prompt

    if use_data_frame:
        text_list_0, text_list_1 = tuple(
            [_get_text(dict_data, key) for dict_data in list_dict_data]
            for key in ("output_1", "output_2")
        )
    else:
        text_list_0 = list_dict_data.map(
            lambda example: {
                "full_prompt": _get_text(example, "output_1"),
            },
            num_proc=16,
        )
        text_list_1 = list_dict_data.map(
            lambda example: {
                "full_prompt": _get_text(example, "output_2"),
            },
            num_proc=16,
        )

    def _merge_tokenization_metadata(metadata_list: Sequence[dict]) -> dict:
        num_examples = sum(metadata["num_examples"] for metadata in metadata_list)
        num_truncated_tokens = sum(
            metadata["num_truncated_tokens"] for metadata in metadata_list
        )
        num_truncated_examples = sum(
            metadata["num_truncated_examples"] for metadata in metadata_list
        )
        input_ids_avg_lens = (
            sum(
                [
                    metadata["input_ids_avg_len"] * metadata["num_examples"]
                    for metadata in metadata_list
                ]
            )
            / num_examples
        )
        input_ids_max_len = max(
            metadata["input_ids_max_len"] for metadata in metadata_list
        )
        input_ids_min_len = min(
            metadata["input_ids_min_len"] for metadata in metadata_list
        )
        labels_avg_lens = (
            sum(
                [
                    metadata["labels_avg_len"] * metadata["num_examples"]
                    for metadata in metadata_list
                ]
            )
            / num_examples
        )
        labels_max_len = max(metadata["labels_max_len"] for metadata in metadata_list)
        labels_min_len = min(metadata["labels_min_len"] for metadata in metadata_list)
        return dict(
            num_examples=num_examples,
            num_truncated_tokens=num_truncated_tokens,
            num_truncated_examples=num_truncated_examples,
            input_ids_avg_len=input_ids_avg_lens,
            input_ids_max_len=input_ids_max_len,
            input_ids_min_len=input_ids_min_len,
            labels_avg_len=labels_avg_lens,
            labels_max_len=labels_max_len,
            labels_min_len=labels_min_len,
        )

    logger.warning(f"Tokenizing {len(list_dict_data)} pairs...")
    tokenized_0, tokenized_1 = tuple(
        _tokenize_fn(text_list, tokenizer, use_data_frame)
        for text_list in (text_list_0, text_list_1)
    )
    # "size" (bsz, 2, seq_len)
    input_ids = [
        list(pair)
        for pair in utils.zip_(tokenized_0["input_ids"], tokenized_1["input_ids"])
    ]
    labels = [
        list(pair) for pair in utils.zip_(tokenized_0["labels"], tokenized_1["labels"])
    ]
    tokenization_metadata = _merge_tokenization_metadata(
        [tokenized_0["tokenization_metadata"], tokenized_1["tokenization_metadata"]]
    )

    packaged_data = dict(
        input_ids=input_ids,
        labels=labels,
        index_0=index_0,
        index_1=index_1,
        choice=choice,
        tokenization_metadata=tokenization_metadata,
        metadata=dict(mean_choice=choice.float().mean().item()),
    )
    if verbose:
        logger.warning(
            f"Tokenization metadata:\n{json.dumps(packaged_data['tokenization_metadata'])}"
        )

    return packaged_data


class BinaryRewardModelingDataset(Dataset):
    def __init__(
        self,
        data: Sequence[dict],
        meta_prompts: List[str],
        principle_collection: Dict[str, str],
        tokenizer: transformers.PreTrainedTokenizer,
        df_postprocessor: Optional[Callable] = None,
        end_sequence_with_eos: bool = False,
        query_len: Optional[int] = None,
        response_len: Optional[int] = None,
        use_data_frame: bool = True,
    ):
        super(BinaryRewardModelingDataset, self).__init__()
        data_dict = preprocess_for_reward_modeling(
            data=data,
            meta_prompts=meta_prompts,
            principle_collection=principle_collection,
            tokenizer=tokenizer,
            df_postprocessor=df_postprocessor,
            end_sequence_with_eos=end_sequence_with_eos,
            query_len=query_len,
            response_len=response_len,
            use_data_frame=use_data_frame,
        )
        self.input_ids = data_dict["input_ids"]
        self.labels = data_dict["labels"]
        self.index_0 = data_dict["index_0"]
        self.index_1 = data_dict["index_1"]
        self.choice = data_dict["choice"]
        self.metadata = data_dict["metadata"]

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        return dict(
            input_ids=self.input_ids[i],
            labels=self.labels[i],
            index_0=self.index_0[i],
            index_1=self.index_1[i],
            choice=self.choice[i],
        )


def _get_generator(seed: int) -> torch.Generator:
    rng = torch.Generator()
    rng.manual_seed(seed)
    return rng


def split_train_into_train_and_eval(
    train_dataset: Dataset, eval_size: int, seed: int
) -> Tuple[Dataset, Dataset]:
    assert eval_size < len(
        train_dataset  # noqa
    ), "Requested eval_size cannot be equal/larger than original train data size."
    new_train_size = len(train_dataset) - eval_size  # noqa
    train_dataset, eval_dataset = torch.utils.data.random_split(
        train_dataset, [new_train_size, eval_size], generator=_get_generator(seed)
    )
    return train_dataset, eval_dataset


def pad_sequence_from_left(
    sequences: Sequence[torch.Tensor],
    batch_first: bool = False,
    padding_value: float = 0.0,
):
    """Mirror of `torch.nn.utils.rnn.pad_sequence`, but pad from left."""
    sequences = tuple(sequence.flip(0) for sequence in sequences)
    padded_sequence = torch._C._nn.pad_sequence(
        sequences, batch_first, padding_value
    )  # noqa
    padded_sequence = padded_sequence.flip(int(batch_first))
    return padded_sequence


@dataclass
class DataCollatorForBinaryRewardModelingDataset(object):
    """
    This collation assumes data preprocessing converts text into *padded* tensors of the same length.
    For autoregressive models like OPT and GPT2, `input_ids` alone is sufficient to produce the rewards.
    For enc-dec models like T5, we need `labels`.

    `input_ids` and `labels` are tensors of size (bsz, num_candidates, max_seq_len), i.e., each batch instance has
    `num_candidates` generations/completions.
    `index_0` and `index_1` are tensors of size (bsz, num_pairs), and are used to index into `input_ids` and
    `labels` to find the first and second sequences in the pair.
    `choice` is a binary int/long tensor of size (bsz, num_pairs) indicating which sequence in the pair is better,
    i.e., 0 means the first sequence is preferred, and 1 means otherwise.
    """

    tokenizer: transformers.PreTrainedTokenizer

    def _left_pad_helper(self, instances: Sequence[dict], key: str):
        # TODO(lxuechen): Potentially replace with `transformers.PretrainedTokenizerBase.prepare_for_model`.
        # `instances` is a list of dicts, each dict has key whose value is a list of tensors, possibly of unequal length.
        input_ids = [seq for instance in instances for seq in instance[key]]  # Flatten.
        input_ids = pad_sequence_from_left(
            input_ids,
            batch_first=True,
            padding_value=self.tokenizer.pad_token_id,
        )
        input_ids = einops.rearrange(
            input_ids,
            "(bsz num_candidates) max_seq_len -> bsz num_candidates max_seq_len",
            num_candidates=len(instances[0][key]),
        )
        return input_ids

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        index_0, index_1, choice = tuple(
            torch.stack([instance[key] for instance in instances])
            for key in ("index_0", "index_1", "choice")
        )
        input_ids = self._left_pad_helper(instances, "input_ids")
        attention_mask = input_ids.ne(self.tokenizer.pad_token_id).long()
        return dict(
            input_ids=input_ids,
            attention_mask=attention_mask,
            index_0=index_0,
            index_1=index_1,
            choice=choice,
        )


def make_binary_reward_modeling_data_module(
    tokenizer: transformers.PreTrainedTokenizer,
    data_args,
    training_args,
):
    meta_prompts = utils.make_meta_prompts(data_args.meta_prompt_pattern)

    principle_collection = None
    if data_args.principle_collection_path:
        principle_collection = {}
        with open(data_args.principle_collection_path, "r", encoding="utf-8") as f:
            principle_collection_list = json.load(f)
            for item in principle_collection_list:
                principle_collection[item["dimension"]] = item["definition"]
                if "negative_definition" in item:
                    principle_collection[item["dimension"] + "_neg"] = item[
                        "negative_definition"
                    ]

    if "alpaca_farm" in data_args.dataset_path:
        train_preference = load_dataset(data_args.dataset_path, data_args.dataset_name)[
            "preference"
        ]
        use_data_frame = True
    elif data_args.dataset_path.endswith("json"):
        train_preference = load_dataset("json", data_files=data_args.dataset_path)[
            "train"
        ]
        use_data_frame = False
    else:
        raise ValueError(
            f"Unsupported dataset_path: {data_args.dataset_path}."
            "Only json and alpaca_farm datasets are supported."
        )

    train_dataset = BinaryRewardModelingDataset(
        data=train_preference,
        meta_prompts=meta_prompts,
        principle_collection=principle_collection,
        tokenizer=tokenizer,
        end_sequence_with_eos=training_args.end_sequence_with_eos,
        query_len=training_args.query_len,
        response_len=training_args.response_len,
        use_data_frame=use_data_frame,
    )

    if (
        data_args.dataset_path == data_args.eval_dataset_path
        and data_args.dataset_name == data_args.eval_dataset_name
    ):
        train_dataset, eval_dataset = split_train_into_train_and_eval(
            train_dataset=train_dataset,
            eval_size=data_args.eval_size,
            seed=training_args.seed,
        )

    else:
        alpaca_human_preference = load_dataset(
            data_args.eval_dataset_path, data_args.eval_dataset_name
        )["preference"]
        eval_dataset = BinaryRewardModelingDataset(
            data=alpaca_human_preference,
            meta_prompts=meta_prompts,
            principle_collection=principle_collection,
            tokenizer=tokenizer,
            end_sequence_with_eos=training_args.end_sequence_with_eos,
            query_len=training_args.query_len,
            response_len=training_args.response_len,
        )
        _, eval_dataset = split_train_into_train_and_eval(
            train_dataset=eval_dataset,
            eval_size=data_args.eval_size,
            seed=training_args.seed,
        )

    data_collator = DataCollatorForBinaryRewardModelingDataset(tokenizer=tokenizer)
    return dict(
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
    )

training/data_utils/data_utils_rm.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
#
# 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 copy

import os
from dataclasses import dataclass, field
from typing import Optional, Dict, Sequence, List
import pandas as pd

import torch
import transformers
from torch.nn.utils.rnn import pad_sequence
from datasets import load_dataset, Dataset

import data_utils.common_utils as utils


IGNORE_INDEX = -100


@dataclass
class DataCollatorForCausalLM(object):
    left_truncated_tokenizer: transformers.PreTrainedTokenizer
    tokenizer: transformers.PreTrainedTokenizer
    source_max_len: int
    target_max_len: int
    train_on_source: bool
    predict_with_generate: bool
    add_eos_to_target: bool

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        # Extract elements
        sources = [example["input"] for example in instances]
        if self.add_eos_to_target:
            targets = [
                f"\n{example['output']}{self.tokenizer.eos_token}"
                for example in instances
            ]
        else:
            targets = [f"\n{example['output']}" for example in instances]

        begin_padding_len = self.tokenizer(
            ["\n"], return_tensors="pt", add_special_tokens=False
        ).input_ids.shape[1]

        # Tokenize
        tokenized_sources_with_prompt = self.left_truncated_tokenizer(
            sources,
            max_length=self.source_max_len,
            truncation=True,
            # add_special_tokens=False,
        )
        tokenized_targets = self.tokenizer(
            targets,
            max_length=self.target_max_len + begin_padding_len,
            truncation=True,
            add_special_tokens=False,
        )
        # Build the input and labels for causal LM
        input_ids = []
        labels = []
        for tokenized_source, tokenized_target in zip(
            tokenized_sources_with_prompt["input_ids"], tokenized_targets["input_ids"]
        ):
            tokenized_target = tokenized_target[begin_padding_len:]
            if not self.predict_with_generate:
                input_ids.append(torch.tensor(tokenized_source + tokenized_target))
                if not self.train_on_source:
                    labels.append(
                        torch.tensor(
                            [IGNORE_INDEX for _ in range(len(tokenized_source))]
                            + copy.deepcopy(tokenized_target)
                        )
                    )
                else:
                    labels.append(
                        torch.tensor(copy.deepcopy(tokenized_source + tokenized_target))
                    )
            else:
                input_ids.append(torch.tensor(tokenized_source))
        # Apply padding
        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_INDEX)
            if not self.predict_with_generate
            else None
        )
        data_dict = {
            "input_ids": input_ids,
            "attention_mask": input_ids.ne(self.tokenizer.pad_token_id),
        }
        if labels is not None:
            data_dict["labels"] = labels
        return data_dict


def extract_unnatural_instructions_data(examples, extract_reformulations=False):
    out = {
        "input": [],
        "output": [],
    }
    for example_instances in examples["instances"]:
        for instance in example_instances:
            out["input"].append(instance["instruction_with_input"])
            out["output"].append(instance["output"])
    if extract_reformulations:
        for example_reformulations in examples["reformulations"]:
            if example_reformulations is not None:
                for instance in example_reformulations:
                    out["input"].append(instance["instruction_with_input"])
                    out["output"].append(instance["output"])
    return out


DROMEDARY_PROMPT_DICT = {
    "prompt_input": (
        "{meta_prompt}\n" "{instruction}\n\n" "{input}\n\n" "### Dromedary"
    ),
    "prompt_no_input": ("{meta_prompt}\n" "{instruction}\n\n" "### Dromedary"),
}


def extract_dromedary_dataset(example, meta_prompts):
    assert "example_id" in example
    total_meta_prompt = len(meta_prompts)
    meta_prompt = meta_prompts[int(example["example_id"]) % total_meta_prompt]

    if example.get("input", "") != "":
        prompt_format = DROMEDARY_PROMPT_DICT["prompt_input"]
    else:
        prompt_format = DROMEDARY_PROMPT_DICT["prompt_no_input"]

    return {
        "input": prompt_format.format(meta_prompt=meta_prompt, **example),
        "output": "\n" + example["output"],
    }


def local_dataset(dataset_name):
    if dataset_name.endswith(".json"):
        full_dataset = load_dataset("json", data_files=dataset_name)
    else:
        raise ValueError(f"Unsupported dataset format: {dataset_name}")

    return full_dataset


def make_sft_data_module(
    left_truncated_tokenizer: transformers.PreTrainedTokenizer,
    tokenizer: transformers.PreTrainedTokenizer,
    args,
) -> Dict:
    """
    Make dataset and collator for supervised fine-tuning.
    Datasets are expected to have the following columns: { `input`, `output` }
    """

    def load_data(dataset_name):
        if dataset_name == "alpaca":
            return load_dataset("tatsu-lab/alpaca")
        elif dataset_name == "alpaca-clean":
            return load_dataset("yahma/alpaca-cleaned")
        elif dataset_name == "chip2":
            return load_dataset("laion/OIG", data_files="unified_chip2.jsonl")
        elif dataset_name == "self-instruct":
            return load_dataset("yizhongw/self_instruct", name="self_instruct")
        elif dataset_name == "hh-rlhf":
            return load_dataset("Anthropic/hh-rlhf")
        elif dataset_name == "longform":
            return load_dataset("akoksal/LongForm")
        elif dataset_name == "oasst1":
            return load_dataset("timdettmers/openassistant-guanaco")
        elif dataset_name == "vicuna":
            raise NotImplementedError("Vicuna data was not released.")
        else:
            if os.path.exists(dataset_name):
                try:
                    args.dataset_format = (
                        args.dataset_format if args.dataset_format else "alpaca"
                    )
                    full_dataset = local_dataset(dataset_name)
                    return full_dataset
                except:
                    raise ValueError(f"Error loading dataset from {dataset_name}")
            else:
                raise NotImplementedError(
                    f"Dataset {dataset_name} not implemented yet."
                )

    def format_dataset(dataset, dataset_format, meta_prompt_pattern):
        if dataset_format == "dromedary":
            assert meta_prompt_pattern is not None

            meta_prompts = utils.make_meta_prompts(meta_prompt_pattern)

            dataset = dataset.map(
                lambda ex: extract_dromedary_dataset(ex, meta_prompts=meta_prompts),
                remove_columns=["instruction", "example_id"],
            )
        elif dataset_format == "chip2" or (
            dataset_format is None and args.dataset == "chip2"
        ):
            dataset = dataset.map(
                lambda x: {
                    "input": x["text"].split("\n<bot>: ")[0].replace("<human>: ", ""),
                    "output": x["text"].split("\n<bot>: ")[1],
                }
            )
        elif dataset_format == "self-instruct" or (
            dataset_format is None and args.dataset == "self-instruct"
        ):
            for old, new in [["prompt", "input"], ["completion", "output"]]:
                dataset = dataset.rename_column(old, new)
        elif dataset_format == "hh-rlhf" or (
            dataset_format is None and args.dataset == "hh-rlhf"
        ):
            dataset = dataset.map(lambda x: {"input": "", "output": x["chosen"]})
        elif dataset_format == "oasst1" or (
            dataset_format is None and args.dataset == "oasst1"
        ):
            dataset = dataset.map(
                lambda x: {
                    "input": "",
                    "output": x["text"],
                }
            )
        # Remove unused columns.
        dataset = dataset.remove_columns(
            [
                col
                for col in dataset.column_names["train"]
                if col not in ["input", "output"]
            ]
        )
        return dataset

    # Load dataset.
    dataset = load_data(args.dataset)
    dataset = format_dataset(dataset, args.dataset_format, args.meta_prompt_pattern)

    # Split train/eval, reduce size
    if args.do_eval or args.do_predict:
        if "eval" in dataset:
            eval_dataset = dataset["eval"]
        else:
            print(
                "Splitting train dataset in train and validation according to `eval_dataset_size`"
            )
            dataset = dataset["train"].train_test_split(
                test_size=args.eval_dataset_size, shuffle=True, seed=42
            )
            eval_dataset = dataset["test"]
        if (
            args.max_eval_samples is not None
            and len(eval_dataset) > args.max_eval_samples
        ):
            eval_dataset = eval_dataset.select(range(args.max_eval_samples))
        if args.group_by_length:
            eval_dataset = eval_dataset.map(
                lambda x: {"length": len(x["input"]) + len(x["output"])}
            )
    if args.do_train:
        train_dataset = dataset["train"]
        if (
            args.max_train_samples is not None
            and len(train_dataset) > args.max_train_samples
        ):
            train_dataset = train_dataset.select(range(args.max_train_samples))
        if args.group_by_length:
            train_dataset = train_dataset.map(
                lambda x: {"length": len(x["input"]) + len(x["output"])}
            )

    data_collator = DataCollatorForCausalLM(
        left_truncated_tokenizer=left_truncated_tokenizer,
        tokenizer=tokenizer,
        source_max_len=args.source_max_len,
        target_max_len=args.target_max_len,
        train_on_source=args.train_on_source,
        predict_with_generate=args.predict_with_generate,
        add_eos_to_target=args.add_eos_to_target,
    )
    return dict(
        train_dataset=train_dataset if args.do_train else None,
        eval_dataset=eval_dataset if args.do_eval else None,
        predict_dataset=eval_dataset if args.do_predict else None,
        data_collator=data_collator,
    )

training/data_utils/data_utils_sft.py
IBM/SALMON
# coding=utf-8
# Copyright 2023 Self-Align, EleutherAI, and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# 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.
""" LLaMA model configuration. """

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging


logger = logging.get_logger(__name__)

LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = {}


class LlamaConfig(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA
    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
    defaults will yield a similar configuration to that of the LLaMA-7B.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information."""
    model_type = "llama"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size=32000,
        hidden_size=4096,
        intermediate_size=11008,
        num_hidden_layers=32,
        num_attention_heads=32,
        num_key_value_heads=None,
        hidden_act="silu",
        max_position_embeddings=2048,
        initializer_range=0.02,
        rms_norm_eps=1e-6,
        use_cache=True,
        pad_token_id=0,
        bos_token_id=1,
        eos_token_id=2,
        tie_word_embeddings=False,
        rope_scaling=None,
        **kwargs,
    ):
        # for backward compatibility
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads

        self.num_key_value_heads = num_key_value_heads
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.rope_scaling = rope_scaling
        self._rope_scaling_validation()
        self.use_cache = use_cache
        self.cache_shape = None
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

    def _rope_scaling_validation(self):
        """
        Validate the `rope_scaling` configuration.
        """
        if self.rope_scaling is None:
            return

        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
            raise ValueError(
                "`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, "
                f"got {self.rope_scaling}"
            )
        rope_scaling_type = self.rope_scaling.get("type", None)
        rope_scaling_factor = self.rope_scaling.get("factor", None)
        if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
            raise ValueError(
                f"`rope_scaling`'s name field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
            )
        if (
            rope_scaling_factor is None
            or not isinstance(rope_scaling_factor, float)
            or rope_scaling_factor <= 1.0
        ):
            raise ValueError(
                f"`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}"
            )

training/models/configuration_llama.py
IBM/SALMON
# Copyright 2023 The Alpaca Team
#
# 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.

"""Utilities for PyTorch's distributed training.

Compatible with torchrun / elastic.

Internal map:
    https://github.com/lxuechen/ml-swissknife/blob/main/ml_swissknife/distributed_utils.py
"""

import os
import sys
from typing import Optional

import torch
import torch.distributed as dist


def setup(rank: Optional[int] = None, world_size: Optional[int] = None):
    if rank is None:
        rank = get_rank()
    if world_size is None:
        world_size = get_world_size()

    if world_size <= 1:
        return rank, world_size

    if not dist.is_initialized():
        if sys.platform == "win32":
            # Distributed package only covers collective communications with Gloo
            # backend and FileStore on Windows platform. Set init_method parameter
            # in init_process_group to a local file.
            # Example init_method="file:///f:/libtmp/some_file"
            init_method = "file:///f:/libtmp/dist-tmp"
            dist.init_process_group(
                backend="gloo",
                init_method=init_method,
                rank=rank,
                world_size=world_size,
            )
        elif torch.cuda.is_available():
            dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
        else:
            dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)

    return rank, world_size


def cleanup():
    dist.destroy_process_group()


def get_rank():
    return int(os.getenv("RANK", 0))


def get_local_rank():
    return int(os.getenv("LOCAL_RANK", 0))


def get_world_size():
    return int(os.getenv("WORLD_SIZE", 1))


def should_save():
    """Return True if the current process is the main process."""
    return get_rank() <= 0


def all_gather_and_cat(tensor: torch.Tensor, dim=0):
    if get_world_size() > 1:
        tensor_list = [torch.empty_like(tensor) for _ in range(get_world_size())]
        dist.all_gather(tensor_list, tensor)
        tensor = torch.cat(tensor_list, dim=dim)
    return tensor


is_main_process = should_save

training/models/distributed_utils.py
IBM/SALMON
# coding=utf-8
# Copyright 2023 Self-Align, EleutherAI, and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# 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.
""" PyTorch LLaMA model with Flash-Attention support. """

import math
from typing import List, Optional, Tuple, Union

import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
    SequenceClassifierOutputWithPast,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import (
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)

from models.configuration_llama import LlamaConfig

from einops import rearrange


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "LlamaConfig"


# Copied from transformers.models.bart.modeling_bart._make_causal_mask
def _make_causal_mask(
    input_ids_shape: torch.Size,
    dtype: torch.dtype,
    device: torch.device,
    past_key_values_length: int = 0,
):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz, tgt_len = input_ids_shape
    mask = torch.full(
        (tgt_len, tgt_len),
        torch.tensor(torch.finfo(dtype).min, device=device),
        device=device,
    )
    mask_cond = torch.arange(mask.size(-1), device=device)
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat(
            [
                torch.zeros(
                    tgt_len, past_key_values_length, dtype=dtype, device=device
                ),
                mask,
            ],
            dim=-1,
        )
    return mask[None, None, :, :].expand(
        bsz, 1, tgt_len, tgt_len + past_key_values_length
    )


# Copied from transformers.models.bart.modeling_bart._expand_mask
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(
        inverted_mask.to(torch.bool), torch.finfo(dtype).min
    )


class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return (self.weight * hidden_states).to(input_dtype)


class LlamaRotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq)

        # Build here to make `torch.jit.trace` work.
        self.max_seq_len_cached = max_position_embeddings
        t = torch.arange(
            self.max_seq_len_cached,
            device=self.inv_freq.device,
            dtype=self.inv_freq.dtype,
        )
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        dtype = torch.get_default_dtype()
        self.register_buffer(
            "cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False
        )
        self.register_buffer(
            "sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False
        )

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
        if seq_len > self.max_seq_len_cached:
            self.max_seq_len_cached = seq_len
            t = torch.arange(
                self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype
            )
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            # Different from paper, but it uses a different permutation in order to obtain the same calculation
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self.register_buffer(
                "cos_cached", emb.cos()[None, None, :, :].to(x.dtype), persistent=False
            )
            self.register_buffer(
                "sin_cached", emb.sin()[None, None, :, :].to(x.dtype), persistent=False
            )
        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )


def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    q_embed = q_embed.to(q.dtype)
    k_embed = k_embed.to(k.dtype)
    return q_embed, k_embed


class LlamaMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
    ):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.act_fn = ACT2FN[hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, num_key_value_heads, n_rep, slen, head_dim
    )
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


class LlamaAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: LlamaConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.max_position_embeddings = config.max_position_embeddings

        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )
        self.q_proj = nn.Linear(
            self.hidden_size, self.num_heads * self.head_dim, bias=False
        )
        self.k_proj = nn.Linear(
            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.o_proj = nn.Linear(
            self.num_heads * self.head_dim, self.hidden_size, bias=False
        )
        self.rotary_emb = LlamaRotaryEmbedding(
            self.head_dim, max_position_embeddings=self.max_position_embeddings
        )

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return (
            tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

    def orig_forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_shape: Optional[Tuple[int]] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        bsz, q_len, _ = hidden_states.size()
        assert output_attentions is False, "output_attentions not implemented"

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        assert query_states.dtype == hidden_states.dtype

        query_states = query_states.view(
            bsz, q_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key_states = key_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)
        value_states = value_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)

        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            kv_seq_len += past_key_value[2]
        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        query_states, key_states = apply_rotary_pos_emb(
            query_states, key_states, cos, sin, position_ids
        )
        # [bsz, nh, t, hd]

        assert query_states.dtype == hidden_states.dtype

        if use_cache:
            if cache_shape is None:
                if past_key_value is not None:
                    key_states = torch.cat([past_key_value[0], key_states], dim=2)
                    value_states = torch.cat([past_key_value[1], value_states], dim=2)
            else:
                if past_key_value is None:
                    past_key_value = (
                        torch.zeros(
                            bsz,
                            self.num_key_value_heads,
                            cache_shape[0],
                            self.head_dim,
                            device=hidden_states.device,
                            dtype=key_states.dtype,
                        ),
                        torch.zeros(
                            bsz,
                            self.num_key_value_heads,
                            cache_shape[0],
                            self.head_dim,
                            device=hidden_states.device,
                            dtype=value_states.dtype,
                        ),
                        0,
                    )
                past_key = past_key_value[0]
                past_value = past_key_value[1]
                num_past_tokens = past_key_value[2]

                past_key[:, :, num_past_tokens:kv_seq_len] = key_states
                past_value[:, :, num_past_tokens:kv_seq_len] = value_states

                key_states = past_key[:, :, :kv_seq_len]
                value_states = past_value[:, :, :kv_seq_len]

                past_key_value = (
                    past_key,
                    past_value,
                    kv_seq_len,
                )
        else:
            past_key_value = None

        # repeat k/v heads if n_kv_heads < n_heads
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        # [bsz, nh, t, hd]
        with torch.backends.cuda.sdp_kernel(
            enable_flash=False, enable_mem_efficient=False
        ):
            if attention_mask is not None:
                if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                    raise ValueError(
                        f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                    )
                output = F.scaled_dot_product_attention(
                    query_states, key_states, value_states, attn_mask=attention_mask
                )
            else:
                output = F.scaled_dot_product_attention(
                    query_states,
                    key_states,
                    value_states,
                    is_causal=True,
                )

        attn_output = self.o_proj(rearrange(output, "b h s d -> b s (h d)"))
        return attn_output, None, past_key_value

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_shape: Optional[Tuple[int]] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        """Input shape: Batch x Time x Channel

        attention_mask: [bsz, q_len]
        """

        if past_key_value is not None or use_cache or output_attentions:
            return self.orig_forward(
                hidden_states=hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_value=past_key_value,
                output_attentions=output_attentions,
                use_cache=use_cache,
                cache_shape=cache_shape,
            )

        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(
            bsz, q_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key_states = key_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)
        value_states = value_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)
        # [bsz, q_len, nh, hd]
        # [bsz, nh, q_len, hd]

        assert query_states.dtype == hidden_states.dtype

        kv_seq_len = key_states.shape[-2]
        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        query_states, key_states = apply_rotary_pos_emb(
            query_states, key_states, cos, sin, position_ids
        )

        assert query_states.dtype == hidden_states.dtype

        # repeat k/v heads if n_kv_heads < n_heads
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        # [bsz, nh, t, hd]
        with torch.backends.cuda.sdp_kernel():
            if attention_mask is not None:
                if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                    raise ValueError(
                        f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                    )
                output = F.scaled_dot_product_attention(
                    query_states, key_states, value_states, attn_mask=attention_mask
                )
            else:
                output = F.scaled_dot_product_attention(
                    query_states,
                    key_states,
                    value_states,
                    is_causal=True,
                )

        return self.o_proj(rearrange(output, "b h s d -> b s (h d)")), None, None


class LlamaDecoderLayer(nn.Module):
    def __init__(self, config: LlamaConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = LlamaAttention(config=config)
        self.mlp = LlamaMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
        )
        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = LlamaRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = False,
        cache_shape: Optional[Tuple[int]] = None,
    ) -> Tuple[
        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
    ]:
        """
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
            use_cache (`bool`, *optional*):
                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
                (see `past_key_values`).
            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
        """

        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            cache_shape=cache_shape,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights,)

        if use_cache:
            outputs += (present_key_value,)

        return outputs


LLAMA_START_DOCSTRING = r"""
    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`LlamaConfig`]):
            Model configuration class with all the parameters of the model. Initializing with a config file does not
            load the weights associated with the model, only the configuration. Check out the
            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""


@add_start_docstrings(
    "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
    LLAMA_START_DOCSTRING,
)
class LlamaPreTrainedModel(PreTrainedModel):
    config_class = LlamaConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["LlamaDecoderLayer"]
    _skip_keys_device_placement = "past_key_values"

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, LlamaModel):
            module.gradient_checkpointing = value


LLAMA_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
            it.

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
            `past_key_values`).

            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
            information on the default strategy.

            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
            config.n_positions - 1]`.

            [What are position IDs?](../glossary#position-ids)
        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
            `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.

            Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
            blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.

            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
            model's internal embedding lookup matrix.
        use_cache (`bool`, *optional*):
            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
            `past_key_values`).
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
            tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""


@add_start_docstrings(
    "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
    LLAMA_START_DOCSTRING,
)
class LlamaModel(LlamaPreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]

    Args:
        config: LlamaConfig
    """

    def __init__(self, config: LlamaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)]
        )
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
    def _prepare_decoder_attention_mask(
        self, attention_mask, input_shape, inputs_embeds, past_key_values_length
    ):
        # create causal mask
        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        combined_attention_mask = None
        if input_shape[-1] > 1:
            combined_attention_mask = _make_causal_mask(
                input_shape,
                inputs_embeds.dtype,
                device=inputs_embeds.device,
                past_key_values_length=past_key_values_length,
            )

        if attention_mask is not None:
            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
            expanded_attn_mask = _expand_mask(
                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
            ).to(inputs_embeds.device)
            combined_attention_mask = (
                expanded_attn_mask
                if combined_attention_mask is None
                else expanded_attn_mask + combined_attention_mask
            )

        return combined_attention_mask

    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        cache_shape = self.config.cache_shape

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        # retrieve input_ids and inputs_embeds
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError(
                "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
            )
        elif input_ids is not None:
            batch_size, seq_length = input_ids.shape
        elif inputs_embeds is not None:
            batch_size, seq_length, _ = inputs_embeds.shape
        else:
            raise ValueError(
                "You have to specify either decoder_input_ids or decoder_inputs_embeds"
            )

        seq_length_with_past = seq_length
        past_key_values_length = 0

        if past_key_values is not None:
            past_key_values_length = past_key_values[0][2]
            seq_length_with_past = seq_length_with_past + past_key_values_length

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length,
                seq_length + past_key_values_length,
                dtype=torch.long,
                device=device,
            )
            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
        else:
            position_ids = position_ids.view(-1, seq_length).long()

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        # embed positions
        if attention_mask is None:
            attention_mask = torch.ones(
                (batch_size, seq_length_with_past),
                dtype=torch.bool,
                device=inputs_embeds.device,
            )
        attention_mask = self._prepare_decoder_attention_mask(
            attention_mask,
            (batch_size, seq_length),
            inputs_embeds,
            past_key_values_length,
        )

        hidden_states = inputs_embeds

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        next_decoder_cache = () if use_cache else None

        for idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            past_key_value = (
                past_key_values[idx] if past_key_values is not None else None
            )

            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # None for past_key_value
                        return module(*inputs, output_attentions, None)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(decoder_layer),
                    hidden_states,
                    attention_mask,
                    position_ids,
                    None,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    cache_shape=cache_shape,
                )

            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
                if v is not None
            )
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )


class LlamaForCausalLM(LlamaPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.model = LlamaModel(config)

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
    @replace_return_docstrings(
        output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
    )
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        r"""
        Args:
            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Returns:

        Example:

        ```python
        >>> from transformers import AutoTokenizer, LlamaForCausalLM

        >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""

        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        logits = logits.float()

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]

        position_ids = kwargs.get("position_ids", None)
        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
            }
        )
        return model_inputs

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        reordered_past = ()
        for layer_past in past_key_values:
            reordered_past += (
                tuple(
                    past_state.index_select(0, beam_idx.to(past_state.device))
                    for past_state in layer_past
                ),
            )
        return reordered_past


@add_start_docstrings(
    """
    The LLaMa Model transformer with a sequence classification head on top (linear layer).

    [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
    (e.g. GPT-2) do.

    Since it does classification on the last token, it requires to know the position of the last token. If a
    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
    each row of the batch).
    """,
    LLAMA_START_DOCSTRING,
)
class LlamaForSequenceClassification(LlamaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = LlamaModel(config)
        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        transformer_outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]
        logits = self.score(hidden_states)

        if input_ids is not None:
            batch_size = input_ids.shape[0]
        else:
            batch_size = inputs_embeds.shape[0]

        if self.config.pad_token_id is None and batch_size != 1:
            raise ValueError(
                "Cannot handle batch sizes > 1 if no padding token is defined."
            )
        if self.config.pad_token_id is None:
            sequence_lengths = -1
        else:
            if input_ids is not None:
                sequence_lengths = (
                    torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1
                ).to(logits.device)
            else:
                sequence_lengths = -1

        pooled_logits = logits[
            torch.arange(batch_size, device=logits.device), sequence_lengths
        ]

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (
                    labels.dtype == torch.long or labels.dtype == torch.int
                ):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(pooled_logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(
                    pooled_logits.view(-1, self.num_labels), labels.view(-1)
                )
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(pooled_logits, labels)
        if not return_dict:
            output = (pooled_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutputWithPast(
            loss=loss,
            logits=pooled_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

training/models/llama_with_flash_attn.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
# Copyright 2023 The Alpaca Team
#
# 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 json
import gc
import glob
from itertools import chain
import logging
import os
import pathlib
import re
from typing import Callable, Dict, List, Optional, Tuple, Union

import accelerate
import numpy as np
import pandas as pd
import torch
import tqdm
import transformers
from transformers import LlamaTokenizerFast

from peft.utils import WEIGHTS_NAME, get_peft_model_state_dict

from data_utils.data_utils_ppo import QueryResponseDataset

import data_utils.common_utils as common_utils
import models.rl_models as rl_models

from models.qlora_model import load_4bit_model_for_inference
from models.reward_model import load_4bit_reward_model_for_inference
from models.rl_trainer import AlpacaAccelerator, RLTrainer, truncate_after_eos
from models.llama_with_flash_attn import LlamaForCausalLM


AnyPath = Union[str, os.PathLike, pathlib.Path]
AnyPathOrNone = Optional[AnyPath]

logger = logging.getLogger(__name__)

if torch.__version__ < "2.0.0":
    LRScheduler = torch.optim.lr_scheduler._LRScheduler  # noqa
else:
    LRScheduler = torch.optim.lr_scheduler.LRScheduler


# Name of the files used for checkpointing
ADAPTER_MODEL_DIR = "adapter_model"
TRAINING_ARGS_NAME = "training_args.bin"
TRAINER_STATE_NAME = "trainer_state.json"
OPTIMIZER_NAME = "optimizer.pt"
SCHEDULER_NAME = "scheduler.pt"
VALUE_HEAD_NAME = "value_head.pt"
SCALER_NAME = "scaler.pt"


class PPOTrainer(RLTrainer):
    def __init__(
        self,
        args,
        train_dataset: QueryResponseDataset,
        eval_dataset: QueryResponseDataset,
        data_collator: Callable,
        policy: rl_models.ActorCritic,
        ref_policy: rl_models.Policy,
        reward_model,
        tokenizer: transformers.PreTrainedTokenizer,
        accelerator: AlpacaAccelerator,
        optimizer: Optional[torch.optim.Optimizer] = None,
        lr_scheduler: Optional[LRScheduler] = None,
        reward_tokenizer: Optional[transformers.PreTrainedTokenizer] = None,
    ):
        super(PPOTrainer, self).__init__(
            args=args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            data_collator=data_collator,
            policy=policy,
            ref_policy=ref_policy,
            reward_model=reward_model,
            tokenizer=tokenizer,
            accelerator=accelerator,
            optimizer=optimizer,
            lr_scheduler=lr_scheduler,
            reward_tokenizer=reward_tokenizer,
        )

    def _shape_reward(
        self,
        rewards: torch.Tensor,
        responses: torch.Tensor,
        logprobs: torch.Tensor,
        ref_logprobs: torch.Tensor,
        length_bonus: torch.Tensor,
    ) -> Dict[str, torch.Tensor]:
        # For some reason, line below doesn't work.
        # kl = (logits.softmax(dim=-1) * (logits.log_softmax(dim=-1) - ref_logits.log_softmax(dim=-1))).sum(dim=-1)

        if self.args.kl_approximator == "k1":
            # KL (q | p) = sum_i q_i (log q_i - log p_i)
            kl = torch.clamp(logprobs - ref_logprobs, min=0.0)
        elif self.args.kl_approximator == "k3":
            # r = p / q, log r = log p - log q
            # KL (q | p) = (r - 1) - log r = e ^ log r - 1 - log r
            log_r = ref_logprobs - logprobs
            kl = torch.exp(log_r) - 1.0 - log_r
        else:
            raise ValueError(f"Unknown KL approximator: {self.args.kl_approximator}")

        non_score_rewards = -self.kl_ctl.value * kl
        shaped_rewards = non_score_rewards.clone()
        # This introduces a small index off by one bug if pad_token_id == eos_token_id.
        # terminal_positions = (responses != self.tokenizer.pad_token_id).sum(dim=1) - 1
        # shaped_rewards[list(range(rewards.size(0))), terminal_positions] += rewards

        shaped_rewards[:, -1] += rewards + (
            torch.clamp(length_bonus, max=self.args.length_bonus_upper_bound)
            * self.args.length_bonus_score
        )
        return dict(
            shaped_rewards=shaped_rewards, non_score_rewards=non_score_rewards, kl=kl
        )

    def _estimate_advantage(
        self, rewards: torch.Tensor, values: torch.Tensor
    ) -> Dict[str, torch.Tensor]:
        """Generalized advantage estimation.

        Reference:
            https://arxiv.org/abs/1506.02438
        """
        if self.args.whiten_rewards:
            rewards = whiten(
                rewards, shift_mean=False, async_stats=self.args.whitening_async_stats
            )
        else:
            rewards = rewards * 10.0
        lastgaelam = 0
        advantages_reversed = []
        gen_length = self.args.response_len
        for t in reversed(range(gen_length)):
            nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0
            delta = rewards[:, t] + self.args.gamma * nextvalues - values[:, t]
            lastgaelam = delta + self.args.gamma * self.args.lam * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1], dim=1)
        returns = advantages + values
        advantages = whiten(
            advantages, shift_mean=True, async_stats=self.args.whitening_async_stats
        )
        return dict(returns=returns, advantages=advantages)

    @torch.inference_mode()
    def rollout(self, queries_data) -> Dict[str, torch.Tensor]:
        """Rollout trajectories with policy.

        Args:
            queries_data: Sequence of batches or DataLoader.
                Each batch is a dict with keys 'queries' and 'query_attn_masks'.

        Returns:
            Dictionary with keys
                'queries', 'query_attn_masks', 'responses',
                'logprobs', 'ref_logprobs', 'values',
                'rewards', 'non_score_rewards', 'shaped_rewards'.
        """
        # Give up dropout throughout.
        self.policy.eval()
        # `keep_fp32_wrapper` retains the autocast wrapper of model.forward created by accelerate:
        #  recall one sets mixed precision options with accelerator.
        # The precise value of this arg doesn't matter here, since we use the unwrapped model only for respond.
        # Generally, try to use the wrapped model as much as you can, since it's got the autocast/cast-back wrappers.
        unwrapped_policy = self.accelerator.unwrap_model(
            self.policy, keep_fp32_wrapper=True
        )

        self.ref_policy.eval()
        self.reward_model.eval()

        rollouts = []
        for batch_idx, batch in tqdm.tqdm(
            enumerate(queries_data),
            total=len(queries_data),
            disable=not self.accelerator.is_main_process,
            desc="rollout",
        ):
            gc.collect()
            torch.cuda.empty_cache()
            # Sample rollouts.
            (
                pure_queries,
                queries,
                query_attn_masks,
                length_bonus_multiplier,
                question_types,
            ) = common_utils.unpack_dict(
                common_utils.prepare_inputs(batch, device=self.accelerator.device),
                keys=(
                    "pure_queries",
                    "queries",
                    "query_attn_masks",
                    "length_bonus",
                    "question_types",
                ),
            )

            respond_outputs = unwrapped_policy.respond(
                queries, query_attn_masks, temperature=self.args.temperature
            )
            (responses,) = common_utils.unpack_dict(respond_outputs, ("responses",))

            rollouts_batch = {
                "queries": queries,
                "query_attn_masks": query_attn_masks,
                "responses": responses,
            }

            # Evaluate reward of the samples.
            text_pure_queries = self.tokenizer.batch_decode(
                pure_queries,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            )

            text_responses = self.tokenizer.batch_decode(
                responses,
                skip_special_tokens=False,
                clean_up_tokenization_spaces=False,
            )

            if isinstance(self.tokenizer, LlamaTokenizerFast):
                llama_fast_special_tokens = [self.tokenizer.eos_token]
                if self.tokenizer.eos_token != self.tokenizer.bos_token:
                    llama_fast_special_tokens.append(self.tokenizer.bos_token)
                for special_token in llama_fast_special_tokens:
                    if special_token is not None:
                        text_responses = [
                            text_response.replace(special_token, f" {special_token}")
                            for text_response in text_responses
                        ]

            if self.args.data_config.stop_token is not None:
                parsed_stop_token = self.args.data_config.stop_token.replace(
                    r"\n", "\n"
                )
            else:
                parsed_stop_token = self.tokenizer.eos_token

            text_responses = [
                truncate_after_stop_token(text_response, parsed_stop_token).split(
                    self.tokenizer.pad_token
                )[0]
                for text_response in text_responses
            ]

            has_stop_token = [
                parsed_stop_token in text_response for text_response in text_responses
            ]

            del pure_queries, queries, responses  # Prevent mistakes.

            # We retokenizer, since policy and reward model might not have the same tokenizer.
            # TODO(lxuechen): Avoid retokenization when policy and reward tokenizer are the same.
            # text_sequences = [
            #     q + r for q, r in common_utils.zip_(text_queries, text_responses)
            # ]
            # TODO(lxuechen): This response retokenization has issues with OPT, since the tokenizer always prepend
            #  <bos_token>. But the issue is local to post_reward, which isn't an issue if we don't penalize.

            question_type_flags = []
            for question_type in question_types.reshape(-1).tolist():
                if question_type == 1:
                    if self.args.enable_reasoning_principles:
                        question_type_flags.append("REASONING")
                    else:
                        question_type_flags.append("GENERAL")
                elif question_type == 2:
                    if self.args.enable_redteaming_principles:
                        question_type_flags.append("REDTEAMING")
                    else:
                        question_type_flags.append("GENERAL")
                else:
                    question_type_flags.append("GENERAL")

            text_sequences = [
                self.prepare_reward_inputs(
                    inputs=q, outputs=clean_after_stop_token(r), question_type_flag=f
                )
                for q, r, f in common_utils.zip_(
                    text_pure_queries, text_responses, question_type_flags
                )
            ]

            # gc.collect()
            # torch.cuda.empty_cache()
            if self.accelerator.is_main_process:
                print("=" * 20)
                # print(rollouts_batch["responses"][0])
                print(text_sequences[0])
                print(
                    self.tokenizer.pad_token_id,
                    self.tokenizer.pad_token,
                    unwrapped_policy.policy.base_tokenizer.pad_token,
                )
                print("Question Type Flag:", question_type_flags[0])
                print(
                    self.args.enable_reasoning_principles,
                    length_bonus_multiplier.reshape(-1),
                )
                print("=" * 20)

            begin_padding_len = self.tokenizer(
                ["\n"], return_tensors="pt", add_special_tokens=False
            ).input_ids.shape[1]
            text_responses = ["\n" + r for r in text_responses]

            sequences = self.reward_tokenizer(
                text_sequences, return_tensors="pt", padding=True, truncation=True
            )

            # Pad on the right to match the "beginning" of the sequence.
            ori_padding_side = self.tokenizer.padding_side
            self.tokenizer.padding_side = "right"
            responses = self.tokenizer(
                text_responses,
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=self.args.response_len + begin_padding_len,
                add_special_tokens=False,
            )
            self.tokenizer.padding_side = ori_padding_side

            responses = responses.input_ids[:, begin_padding_len:]

            non_pad_mask = responses.ne(self.tokenizer.pad_token_id)
            non_pad_seq_len = (
                non_pad_mask.sum(dim=1).float().to(self.accelerator.device)
            )
            length_bonus = non_pad_seq_len / float(self.args.response_len)
            length_bonus = length_bonus * length_bonus_multiplier.reshape(
                length_bonus.shape
            )

            sequences, responses = common_utils.prepare_inputs(
                (sequences, responses), device=self.accelerator.device
            )

            if self.args.clean_tokens_after_eos:
                rollouts_batch["responses"] = responses

            # Evaluate logprobs of the samples.
            batch_size_per_device = rollouts_batch["responses"].shape[0]
            sub_batch_size = self.args.reward_model_per_device_batch_size

            if sub_batch_size is None or sub_batch_size == batch_size_per_device:
                ref_policy_outputs = self.ref_policy(
                    **rollouts_batch, temperature=self.args.temperature
                )
            else:
                assert batch_size_per_device % sub_batch_size == 0

                ref_policy_outputs_list = []

                for sub_batch_idx in range(batch_size_per_device // sub_batch_size):
                    sub_batch = {
                        key: value[
                            sub_batch_idx
                            * sub_batch_size : (sub_batch_idx + 1)
                            * sub_batch_size
                        ]
                        for key, value in rollouts_batch.items()
                    }
                    sub_batch_ref_policy_outputs = self.ref_policy(
                        **sub_batch, temperature=self.args.temperature
                    )
                    ref_policy_outputs_list.append(sub_batch_ref_policy_outputs)

                ref_policy_outputs = common_utils.merge_dict(
                    ref_policy_outputs_list, merge_fn=torch.cat
                )
                del sub_batch_ref_policy_outputs
                del ref_policy_outputs_list
                del sub_batch

            if sub_batch_size is None or sub_batch_size == batch_size_per_device:
                policy_outputs = self.policy(
                    **rollouts_batch,
                    temperature=self.args.temperature,
                )
            else:
                assert batch_size_per_device % sub_batch_size == 0

                policy_outputs_list = []

                for sub_batch_idx in range(batch_size_per_device // sub_batch_size):
                    sub_batch = {
                        key: value[
                            sub_batch_idx
                            * sub_batch_size : (sub_batch_idx + 1)
                            * sub_batch_size
                        ]
                        for key, value in rollouts_batch.items()
                    }
                    sub_batch_policy_outputs = self.policy(
                        **sub_batch,
                        temperature=self.args.temperature,
                    )
                    policy_outputs_list.append(sub_batch_policy_outputs)

                policy_outputs = common_utils.merge_dict(
                    policy_outputs_list, merge_fn=torch.cat
                )
                del sub_batch_policy_outputs
                del policy_outputs_list
                del sub_batch

            policy_outputs = common_utils.unpack_dict(
                policy_outputs,
                keys=("logprobs", "entropies", "values"),
                return_type=dict,
            )
            ref_policy_outputs = common_utils.unpack_dict(
                ref_policy_outputs, keys=("logprobs", "entropies"), return_type=dict
            )

            if self.accelerator.is_main_process:
                # We print log if nan is detected.
                if torch.isnan(policy_outputs["logprobs"]).any():
                    print("logprobs", policy_outputs["logprobs"][0])
                if torch.isnan(policy_outputs["values"]).any():
                    print("values", policy_outputs["values"][0])
                if torch.isnan(ref_policy_outputs["logprobs"]).any():
                    print("ref_logprobs", ref_policy_outputs["logprobs"][0])
                print("=" * 20)

            rollouts_batch.update(policy_outputs)
            rollouts_batch.update(
                {f"ref_{key}": value for key, value in ref_policy_outputs.items()}
            )
            rollouts_batch["length_bonus"] = length_bonus

            if sub_batch_size is None or sub_batch_size == batch_size_per_device:
                reward_outputs = self.reward_model(
                    input_ids=sequences.input_ids,
                    attention_mask=sequences.attention_mask,
                )
            else:
                assert batch_size_per_device % sub_batch_size == 0

                reward_outputs_list = []

                for sub_batch_idx in range(batch_size_per_device // sub_batch_size):
                    idx_start = sub_batch_idx * sub_batch_size
                    idx_end = (sub_batch_idx + 1) * sub_batch_size
                    sub_batch_reward_outputs = self.reward_model(
                        input_ids=sequences.input_ids[idx_start:idx_end],
                        attention_mask=sequences.attention_mask[idx_start:idx_end],
                    )
                    reward_outputs_list.append(sub_batch_reward_outputs)

                reward_outputs = common_utils.merge_dict(
                    reward_outputs_list, merge_fn=torch.cat
                )
                del reward_outputs_list
                del sub_batch_reward_outputs

            if self.accelerator.is_main_process:
                print(reward_outputs)
                print("=" * 20)

            reward_outputs = self.post_reward(
                reward_outputs,
                responses,
                penalize_no_stop_token=self.args.penalize_no_stop_token,
                relative_stop_token_penalty=self.args.relative_stop_token_penalty,
                has_stop_token=has_stop_token,
            )

            rollouts_batch.update(reward_outputs)

            # Shape reward with KL penalty.
            shape_reward_outputs = self._shape_reward(
                rewards=rollouts_batch["rewards"],
                responses=rollouts_batch["responses"],
                logprobs=rollouts_batch["logprobs"],
                ref_logprobs=rollouts_batch["ref_logprobs"],
                length_bonus=rollouts_batch["length_bonus"],
            )
            rollouts_batch.update(shape_reward_outputs)

            rollouts_batch_cpu = {
                key: value.cpu() for key, value in rollouts_batch.items()
            }
            rollouts.append(rollouts_batch_cpu)

        # Items in dict need to be of same shape.
        rollouts = common_utils.merge_dict(rollouts, merge_fn=torch.cat)

        # Estimating advantages outside the loop gives more samples for reward normalization.
        advantages = self._estimate_advantage(
            rewards=rollouts["shaped_rewards"].to(self.accelerator.device),
            values=rollouts["values"].to(self.accelerator.device),
        )
        advantages = {key: value.cpu() for key, value in advantages.items()}
        return {**rollouts, **advantages}

    def post_reward(
        self,
        reward_outputs: Dict[str, torch.Tensor],
        responses: torch.Tensor,
        penalize_no_stop_token: bool,
        relative_stop_token_penalty: bool,
        has_stop_token: List[bool],
    ) -> Dict[str, torch.Tensor]:
        """Assign bad reward values to sequences which didn't stop properly."""
        if penalize_no_stop_token:
            has_stop_token = torch.tensor(has_stop_token, device=responses.device)
            rewards = reward_outputs["rewards"]
            if relative_stop_token_penalty:
                rewards = (
                    rewards + (~has_stop_token).float() * self.args.penalty_reward_value
                )
            else:
                rewards[~has_stop_token] = self.args.penalty_reward_value
            reward_outputs["rewards"] = rewards
            return reward_outputs

        if self.args.truncate_token_ids is None:
            return reward_outputs

        def get_validity_mask(
            sequences: torch.Tensor, end_token_id: int
        ) -> torch.Tensor:
            """Mark a batch element as False if the sequence doesn't end with `end_token_id` after `truncate_after`."""
            assert sequences.dim() == 2
            validity_mask = []
            for sequence in sequences:
                (nonzeros,) = (sequence == end_token_id).nonzero(as_tuple=True)
                if len(nonzeros) == 0:
                    validity_mask.append(False)
                else:
                    validity_mask.append(
                        self.args.truncate_after is None
                        or
                        # Last occurrence of `end_token_id` is after `truncate_after`.
                        nonzeros[-1] > self.args.truncate_after
                    )
            return torch.tensor(validity_mask, device=sequences.device)

        validity_masks = [
            get_validity_mask(responses, end_token_id)
            for end_token_id in self.args.truncate_token_ids
        ]
        validity_mask = torch.stack(validity_masks).any(
            dim=0
        )  # Sequence is valid if it ends with any end token.
        rewards = reward_outputs["rewards"]
        rewards[~validity_mask] = self.args.penalty_reward_value
        return reward_outputs

    def compute_loss(
        self, rollouts: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict]:
        (
            values,
            old_logprob,
            returns,
            advantages,
            queries,
            query_attn_masks,
            responses,
        ) = common_utils.prepare_inputs(
            common_utils.unpack_dict(
                rollouts,
                keys=(
                    "values",
                    "logprobs",
                    "returns",
                    "advantages",
                    "queries",
                    "query_attn_masks",
                    "responses",
                ),
            ),
            device=self.accelerator.device,
        )

        # Enable training mode for graident checkpointing.
        self.policy.train()

        outputs = self.policy(
            queries, query_attn_masks, responses, temperature=self.args.temperature
        )

        vpred = outputs["values"]
        vpredclipped = torch.clamp(
            vpred,
            min=values - self.args.cliprange_value,
            max=values + self.args.cliprange_value,
        )
        vf_losses1 = (vpred - returns) ** 2.0
        vf_losses2 = (vpredclipped - returns) ** 2.0
        vf_loss = 0.5 * torch.maximum(vf_losses1, vf_losses2).mean()
        vf_clipfrac = (vf_losses2 > vf_losses1).to(torch.get_default_dtype()).mean()

        logprob = outputs["logprobs"]
        ratio = torch.exp(logprob - old_logprob)
        # When current policy is close to the old policy, the KL component of this advantage is approximately correct.
        pg_losses = -advantages * ratio
        pg_losses2 = -advantages * torch.clamp(
            ratio, min=1.0 - self.args.cliprange, max=1.0 + self.args.cliprange
        )
        pg_loss = torch.maximum(pg_losses, pg_losses2).mean()
        pg_clipfrac = (
            (pg_losses2 > pg_losses).to(torch.get_default_dtype()).mean()
        )  # noqa

        loss = pg_loss + self.args.vf_coef * vf_loss

        entropy = outputs["entropies"].mean()
        approxkl = 0.5 * ((logprob - old_logprob) ** 2.0).mean()

        return_mean, return_var = returns.mean(), returns.var(unbiased=False)
        value_mean, value_var = values.mean(), values.var(unbiased=False)

        stats = dict(
            loss=dict(policy=pg_loss, value=vf_loss, total=loss),
            policy=dict(entropy=entropy, approxkl=approxkl, clipfrac=pg_clipfrac),
            returns=dict(mean=return_mean, var=return_var),
            val=dict(
                vpred=vpred.mean(),
                error=((vpred - returns) ** 2).mean(),
                clipfrac=vf_clipfrac,
                mean=value_mean,
                var=value_var,
            ),
        )
        return loss, common_utils.flatten_dict(
            stats, sep="/", postprocess_fn=lambda x: x.detach()
        )

    def compute_policy_loss(
        self, rollouts: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict]:
        (
            values,
            old_logprob,
            returns,
            advantages,
            queries,
            query_attn_masks,
            responses,
        ) = common_utils.prepare_inputs(
            common_utils.unpack_dict(
                rollouts,
                keys=(
                    "values",
                    "logprobs",
                    "returns",
                    "advantages",
                    "queries",
                    "query_attn_masks",
                    "responses",
                ),
            ),
            device=self.accelerator.device,
        )

        # Enable training mode for graident checkpointing.
        self.policy.train()

        outputs = self.policy(
            queries,
            query_attn_masks,
            responses,
            temperature=self.args.temperature,
            mode="policy",
        )

        logprob = outputs["logprobs"]
        ratio = torch.exp(logprob - old_logprob)
        # When current policy is close to the old policy, the KL component of this advantage is approximately correct.
        pg_losses = -advantages * ratio
        pg_losses2 = -advantages * torch.clamp(
            ratio, min=1.0 - self.args.cliprange, max=1.0 + self.args.cliprange
        )
        pg_loss = torch.maximum(pg_losses, pg_losses2).mean()
        pg_clipfrac = (
            (pg_losses2 > pg_losses).to(torch.get_default_dtype()).mean()
        )  # noqa

        loss = pg_loss + outputs["dummy_loss"]

        entropy = outputs["entropies"].mean()
        approxkl = 0.5 * ((logprob - old_logprob) ** 2.0).mean()

        return_mean, return_var = returns.mean(), returns.var(unbiased=False)

        stats = dict(
            loss=dict(policy=pg_loss),
            policy=dict(entropy=entropy, approxkl=approxkl, clipfrac=pg_clipfrac),
            returns=dict(mean=return_mean, var=return_var),
        )
        return loss, common_utils.flatten_dict(
            stats, sep="/", postprocess_fn=lambda x: x.detach()
        )

    def compute_value_loss(
        self, rollouts: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict]:
        (
            values,
            old_logprob,
            returns,
            advantages,
            queries,
            query_attn_masks,
            responses,
        ) = common_utils.prepare_inputs(
            common_utils.unpack_dict(
                rollouts,
                keys=(
                    "values",
                    "logprobs",
                    "returns",
                    "advantages",
                    "queries",
                    "query_attn_masks",
                    "responses",
                ),
            ),
            device=self.accelerator.device,
        )

        # Enable training mode for graident checkpointing.
        self.policy.train()

        outputs = self.policy(
            queries,
            query_attn_masks,
            responses,
            temperature=self.args.temperature,
            mode="value",
        )

        vpred = outputs["values"]
        vpredclipped = torch.clamp(
            vpred,
            min=values - self.args.cliprange_value,
            max=values + self.args.cliprange_value,
        )
        vf_losses1 = (vpred - returns) ** 2.0
        vf_losses2 = (vpredclipped - returns) ** 2.0
        vf_loss = 0.5 * torch.maximum(vf_losses1, vf_losses2).mean()
        vf_clipfrac = (vf_losses2 > vf_losses1).to(torch.get_default_dtype()).mean()

        loss = self.args.vf_coef * vf_loss + outputs["dummy_loss"]

        value_mean, value_var = values.mean(), values.var(unbiased=False)

        stats = dict(
            loss=dict(value=vf_loss),
            val=dict(
                vpred=vpred.mean(),
                error=((vpred - returns) ** 2).mean(),
                clipfrac=vf_clipfrac,
                mean=value_mean,
                var=value_var,
            ),
        )
        return loss, common_utils.flatten_dict(
            stats, sep="/", postprocess_fn=lambda x: x.detach()
        )

    def record_step_stats(self, train_stats, rollouts, step_idx, **kwargs):
        kl = rollouts["kl"]
        kl_sum_seq, kl_avg_seq = kl.sum(dim=1).mean(dim=0), kl.mean()
        shaped_rewards = rollouts["shaped_rewards"].sum(dim=1).mean(dim=0)
        non_score_rewards = rollouts["non_score_rewards"].sum(dim=1).mean(dim=0)
        rewards = rollouts["rewards"].mean(dim=0)
        stats = {
            f"objective/kl_coef": kwargs["kl_coef"],
            f"objective/kl_sum_seq": kl_sum_seq,
            f"objective/kl_avg_seq": kl_avg_seq,
            f"objective/length_bonus": rollouts["length_bonus"].mean(),
            f"objective/shaped_rewards": shaped_rewards,
            f"objective/non_score_rewards": non_score_rewards,
            f"objective/rewards": rewards,  # Original model reward.
            f"objective/lr": self.optimizer.param_groups[0]["lr"],
            f"objective/entropies": rollouts["entropies"].mean(),
            f"objective/ref_entropies": rollouts["ref_entropies"].mean(),
        }
        for k, v in train_stats.items():
            stats[f"ppo/{k}"] = v.mean(dim=0)
        stats = {
            key: value.item() if torch.is_tensor(value) else value
            for key, value in stats.items()
        }
        if self.accelerator.is_main_process:
            self.accelerator.log(stats, step=step_idx)
            if self.args.output_dir is not None:
                # Store rollout data to disk to debug.
                rollouts_to_disk = {
                    key: self.tokenizer.batch_decode(
                        truncate_after_eos(
                            tensor,
                            self.tokenizer.eos_token_id,
                        ),
                        skip_special_tokens=False,
                        clean_up_tokenization_spaces=False,
                    )
                    for key, tensor in common_utils.unpack_dict(
                        rollouts, keys=("queries", "responses"), return_type=dict
                    ).items()
                }

                rewards = [str(_) for _ in rollouts["rewards"].tolist()]
                rollouts_to_disk["rewards"] = rewards

                rollouts_to_disk = pd.DataFrame(rollouts_to_disk).to_dict(
                    orient="records"
                )
                rollout_log_dir = os.path.join(self.args.output_dir, "rollouts")
                os.makedirs(rollout_log_dir, exist_ok=True)
                with open(
                    os.path.join(rollout_log_dir, f"step_{step_idx}.json"),
                    "w",
                    encoding="utf-8",
                ) as f:
                    json.dump(rollouts_to_disk, f, indent=4)
        return stats

    @torch.inference_mode()
    def save_model(
        self,
        output_dir: Optional[str] = None,
        give_rw_access=True,
        check_corrupted=True,
    ):
        output_dir = self.args.output_dir if output_dir is None else output_dir

        global_rank = int(os.environ.get("RANK", 0))

        if global_rank == 0:
            os.makedirs(output_dir, exist_ok=True)
            print("Saving model checkpoint to %s" % output_dir)

            # Save policy model.
            unwrapped_policy = self.accelerator.unwrap_model(
                self.policy, keep_fp32_wrapper=True
            )

            policy_model = unwrapped_policy.policy.base_model

            peft_model_path = os.path.join(output_dir, ADAPTER_MODEL_DIR)

            # policy_model.save_pretrained(peft_model_path)
            save_adapters(
                policy_model,
                peft_model_path,
                adapter_names=["lora_policy"],
            )

            pytorch_model_paths = glob.glob(
                os.path.join(output_dir, "pytorch_model*.bin")
            )
            for pytorch_model_path in pytorch_model_paths:
                if os.path.exists(pytorch_model_path):
                    os.remove(pytorch_model_path)

            # Save value model.
            value_model = unwrapped_policy.value_model

            save_adapters(
                value_model.base_model,
                peft_model_path,
                adapter_names=["lora_value"],
            )

            torch.save(
                value_model.value_head.state_dict(),
                os.path.join(output_dir, VALUE_HEAD_NAME),
            )

            pytorch_model_paths = glob.glob(
                os.path.join(output_dir, "pytorch_model*.bin")
            )
            for pytorch_model_path in pytorch_model_paths:
                if os.path.exists(pytorch_model_path):
                    os.remove(pytorch_model_path)

            # Save optimizer.
            torch.save(
                self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME)
            )
            # Save scheduler.
            torch.save(
                self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME)
            )

            # Delete other optimizer checkpoints to save disk space.
            # glob pattern to match all optimizer.pt files in the father directory
            pattern = os.path.join(os.path.dirname(output_dir), "*/optimizer.pt")

            # get a list of all matching paths
            optimizer_files = glob.glob(pattern)

            # iterate over the optimizer files
            for file in optimizer_files:
                # if the file is not in the output_dir, delete it
                if output_dir not in file:
                    os.remove(file)

        else:
            print("Skipping PEFT checkpoint save on rank %d" % global_rank)

    @torch.inference_mode()
    def resume_training(self, checkpoint_dir):
        # Load optimizer.
        optimizer_path = os.path.join(checkpoint_dir, OPTIMIZER_NAME)
        if os.path.exists(optimizer_path):
            load_paged_optimizer_state_dict(
                self.optimizer.optimizer,
                torch.load(
                    optimizer_path,
                    map_location="cpu",
                ),
            )

        # Unpage optimizer state.
        torch.cuda.synchronize()
        gc.collect()
        torch.cuda.empty_cache()

        # Load scheduler.
        scheduler_path = os.path.join(checkpoint_dir, SCHEDULER_NAME)
        if os.path.exists(scheduler_path):
            self.lr_scheduler.load_state_dict(
                torch.load(
                    scheduler_path,
                    map_location="cpu",
                )
            )

        spattern = re.compile(r"checkpoint-(\d+)")
        skipping_steps = int(spattern.search(checkpoint_dir).group(1))
        return skipping_steps


def smart_tokenizer_and_embedding_resize(
    num_new_tokens: int,
    tokenizer: transformers.PreTrainedTokenizer,
    model: transformers.PreTrainedModel,
):
    if num_new_tokens > 0:
        model.resize_token_embeddings(len(tokenizer))
        input_embeddings = model.get_input_embeddings().weight.data
        output_embeddings = model.get_output_embeddings().weight.data

        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
            dim=0, keepdim=True
        )
        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
            dim=0, keepdim=True
        )

        input_embeddings[-num_new_tokens:] = input_embeddings_avg
        output_embeddings[-num_new_tokens:] = output_embeddings_avg

        model.get_input_embeddings().requires_grad_(True)
        model.get_output_embeddings().requires_grad_(True)


def make_models(
    tokenizer: transformers.PreTrainedTokenizer,
    reward_tokenizer: transformers.PreTrainedTokenizer,
    args,
    accelerator: accelerate.Accelerator,
    num_new_tokens: int = 0,
    reward_num_new_tokens: int = 0,
    resume_from_checkpoint: Optional[str] = None,
) -> dict:
    def make_generative_policy(
        adapter_name,
        is_trainable,
        reuse_base_model=True,
        resume_path=None,
        fully_initialize=False,
        policy_model_bits=4,
    ):
        model = load_4bit_model_for_inference(
            checkpoint_dir=resume_path or args.policy_model_name_or_path,
            bits=policy_model_bits,
            fp16=args.fp16,
            bf16=args.bf16,
            gradient_checkpointing=args.gradient_checkpointing,
            adapter_name=adapter_name,
            is_trainable=is_trainable,
            reuse_base_model=reuse_base_model,
            trust_remote_code=args.trust_remote_code,
            base_model_mapping=args.base_model_mapping,
            fully_initialize=fully_initialize,
        )

        if args.bf16 and not isinstance(model, LlamaForCausalLM):
            model = model.to(torch.bfloat16)

        return model

    def make_reward_model(
        adapter_name, is_trainable, reuse_base_model=True, resume_path=None
    ):
        model = load_4bit_reward_model_for_inference(
            checkpoint_dir=resume_path or args.reward_model_name_or_path,
            bits=args.reward_model_bits,
            fp16=args.fp16,
            bf16=args.bf16,
            gradient_checkpointing=args.gradient_checkpointing,
            adapter_name=adapter_name,
            is_trainable=is_trainable,
            reuse_base_model=reuse_base_model,
            trust_remote_code=args.trust_remote_code,
            base_model_mapping=args.base_model_mapping,
        )

        if args.bf16 and not isinstance(model.backbone_model, LlamaForCausalLM):
            model = model.to(torch.bfloat16)

        return model

    # Model construction below seems convoluted, but it's made to trade time for RAM efficiency.
    # For large models, object creation could be extremely RAM intensive.
    # Especially so for multiple processes on single node, each starting off with a copy of the model.
    # General strategy is to 1) create a model, 2) move it to target device / shard it, 3) then start next model,
    # as opposed to creating all needed models on CPU first, and separately moving / sharding each.

    policy_resume_path = None
    if resume_from_checkpoint:
        policy_resume_path = os.path.join(
            resume_from_checkpoint, ADAPTER_MODEL_DIR, "lora_policy"
        )

    ref_policy = rl_models.make_policy_with_base_model(
        args,
        make_generative_policy(
            adapter_name="lora_ref_policy",
            is_trainable=False,
        ),
        tokenizer,
        adapter_name="lora_ref_policy",
    )

    policy = rl_models.make_policy_with_base_model(
        args,
        make_generative_policy(
            adapter_name="lora_policy",
            is_trainable=True,
            resume_path=policy_resume_path,
            fully_initialize=args.fully_initialize_policy,
        ),
        tokenizer,
        adapter_name="lora_policy",
    )

    value_resume_path = None
    value_head_resume_path = None
    if resume_from_checkpoint:
        value_resume_path = os.path.join(
            resume_from_checkpoint, ADAPTER_MODEL_DIR, "lora_value"
        )
        value_head_resume_path = os.path.join(resume_from_checkpoint, VALUE_HEAD_NAME)

    if args.init_value_with_reward:
        # Initialize value from reward model a la OAI.
        logger.warning("Initializing value model with reward model.")
        value_model = rl_models.make_value_with_base_model(
            args,
            make_reward_model(
                adapter_name="lora_value",
                is_trainable=True,
                resume_path=value_resume_path,
            ).backbone_model,
            reward_tokenizer,
            adapter_name="lora_value",
        )
    else:
        logger.warning("Initializing value model with policy model.")
        # Initialize value from policy. Works for sanity, but generally performs worse in instruction-following.
        value_model = rl_models.make_value_with_base_model(
            args,
            make_generative_policy(
                adapter_name="lora_value",
                is_trainable=True,
                # reuse_base_model=False,
                # policy_model_bits=4,
                resume_path=value_resume_path,
                fully_initialize=args.fully_initialize_policy,
            ),
            tokenizer,
            adapter_name="lora_value",
        )

    if value_head_resume_path:
        value_model.value_head.load_state_dict(
            torch.load(value_head_resume_path, map_location="cpu")
        )

    actor_critic = rl_models.ActorCritic(policy=policy, value_model=value_model)
    # We cast how respond should run. It's important the dtypes be consistent with training, since a bf16
    # fine-tuned model might not work with fp16 inference.
    # Cast step below must precede accelerator.prepare(), since wrapped model might not have `respond` method.
    # actor_critic = common.prepare_model_for_custom_fn(
    #     model=actor_critic, fn_name="respond", accelerator=accelerator
    # )

    reward_model = make_reward_model(
        adapter_name="lora_reward",
        is_trainable=False,
    )

    if num_new_tokens > 0:
        smart_tokenizer_and_embedding_resize(
            num_new_tokens, tokenizer, policy.base_model
        )
    if reward_num_new_tokens > 0:
        smart_tokenizer_and_embedding_resize(
            reward_num_new_tokens, reward_tokenizer, reward_model.backbone_model
        )

    if args.fully_initialize_policy:
        actor_critic = accelerator.prepare(actor_critic)  # noqa

    # ref_policy.requires_grad_(False)
    # ref_policy = accelerator.prepare(ref_policy)  # noqa
    # reward_model.requires_grad_(False)
    # if not args.init_value_with_reward:
    #     reward_model = accelerator.prepare(reward_model)

    # # TODO: This is a hack to get FSDP running. Remove in the future when we figure things out.
    # if accelerator.distributed_type == accelerate.DistributedType.FSDP:
    #     inputs = tokenizer("fsdp are you happy now??? :)" * 50, return_tensors="pt")
    #     inputs = {key: value.to(accelerator.device) for key, value in inputs.items()}
    #     actor_critic(inputs["input_ids"], inputs["attention_mask"], inputs["input_ids"])
    return dict(policy=actor_critic, ref_policy=ref_policy, reward_model=reward_model)


def whiten(
    values: torch.Tensor, shift_mean=True, epsilon=1e-8, async_stats="full_batch"
) -> torch.Tensor:
    assert async_stats in ["full_batch", "per_gpu", "none"]

    values_for_statistics = values
    if async_stats == "full_batch":
        if not values_for_statistics.is_cuda:
            raise ValueError("SyncWhiten expected input tensor to be on GPU")

        need_sync = (
            torch.distributed.is_available() and torch.distributed.is_initialized()
        )

        if need_sync:
            process_group = torch.distributed.group.WORLD
            world_size = torch.distributed.get_world_size(process_group)
            need_sync = world_size > 1

        if need_sync:
            tensor_list = [
                torch.zeros_like(values_for_statistics) for _ in range(world_size)
            ]
            torch.distributed.all_gather(tensor_list, values_for_statistics)
            values_for_statistics = torch.cat(tensor_list, dim=0)

    if async_stats in ["full_batch", "per_gpu"]:
        assert (
            values_for_statistics.size(0) >= 8
        ), f"Internal error: Minibatch size {values.size(0)} is insufficient for whitening."
        mean = values_for_statistics.mean()  # noqa
        std = values_for_statistics.std(unbiased=False)  # noqa

    else:
        mean = values.mean(dim=-1, keepdim=True)
        std = values.std(dim=-1, unbiased=False, keepdim=True)

    whitened = (values - mean) / (std + epsilon)
    if not shift_mean:
        whitened = whitened + mean
    return whitened


def clean_after_stop_token(
    response: str,
    additional_stop_token: Optional[List[str]] = None,
) -> str:
    if additional_stop_token is None:
        additional_stop_token = ["<s>", "</s>", "###", "<|endoftext|>"]

    for token in additional_stop_token:
        response = response.split(token)[0]

    return response.strip()


def truncate_after_stop_token(
    response: str,
    stop_token: Optional[str] = None,
    additional_stop_token: Optional[List[str]] = None,
) -> str:
    if stop_token is None:
        return response

    if additional_stop_token is None:
        additional_stop_token = [
            "<s>",
            "</s>",
            "\nUpdate",
            "\nNote",
            "\nComment",
        ]

    for token in additional_stop_token + [stop_token]:
        if len(response.split(token)) > 1:
            response = response.split(token)[0] + token

    if "###" in response:
        response = (
            response.split("###")[0]
            + "###"
            + " ".join(response.split("###")[1].split(" ")[:2])
        )

    return response.strip()


def save_adapters(model, save_directory, adapter_names, **kwargs):
    r"""
    This function saves the adapter model and the adapter configuration files to a directory, so that it can be
    reloaded using the [`LoraModel.from_pretrained`] class method, and also used by the [`LoraModel.push_to_hub`]
    method.

    Args:
        model: The model to save.
        save_directory (`str`):
            Directory where the adapter model and configuration files will be saved (will be created if it does not
            exist).
        adapter_name (`str`):
            Name of the adapter to save.
        kwargs (additional keyword arguments, *optional*):
            Additional keyword arguments passed along to the `push_to_hub` method.
    """
    if os.path.isfile(save_directory):
        raise ValueError(
            f"Provided path ({save_directory}) should be a directory, not a file"
        )
    os.makedirs(save_directory, exist_ok=True)
    # model.create_or_update_model_card(save_directory)

    for adapter_name, peft_config in model.peft_config.items():
        if adapter_name in adapter_names:
            # save only the trainable weights
            output_state_dict = get_peft_model_state_dict(
                model,
                state_dict=kwargs.get("state_dict", None),
                adapter_name=adapter_name,
            )
            output_dir = (
                os.path.join(save_directory, adapter_name)
                if adapter_name != "default"
                else save_directory
            )
            os.makedirs(output_dir, exist_ok=True)

            torch.save(output_state_dict, os.path.join(output_dir, WEIGHTS_NAME))

            # save the config and change the inference mode to `True`
            if peft_config.base_model_name_or_path is None:
                peft_config.base_model_name_or_path = (
                    model.base_model.model.__dict__.get("name_or_path", None)
                )
            inference_mode = peft_config.inference_mode
            peft_config.inference_mode = True
            peft_config.save_pretrained(output_dir)
            peft_config.inference_mode = inference_mode


def load_paged_optimizer_state_dict(optimizer, state_dict):
    """
    Load an optimizer state dict that was saved.
    """

    # Validate the state_dict
    groups = optimizer.param_groups
    saved_groups = state_dict["param_groups"]

    if len(groups) != len(saved_groups):
        raise ValueError(
            "loaded state dict has a different number of " "parameter groups"
        )
    param_lens = (len(g["params"]) for g in groups)
    saved_lens = (len(g["params"]) for g in saved_groups)
    if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
        raise ValueError(
            "loaded state dict contains a parameter group "
            "that doesn't match the size of optimizer's group"
        )

    # Update the state
    id_map = {
        p: old_id
        for old_id, p in zip(
            chain.from_iterable(g["params"] for g in saved_groups),
            chain.from_iterable(g["params"] for g in groups),
        )
    }

    for g in groups:
        for p in g["params"]:
            if p in optimizer.state:
                values = optimizer.state[p]
                for k, v in values.items():
                    if isinstance(v, torch.Tensor):
                        v.copy_(state_dict["state"][id_map[p]][k])
                        optimizer.state[p][k] = v.to("cpu")
                    else:
                        optimizer.state[p][k] = state_dict["state"][id_map[p]][k]

training/models/ppo_trainer.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
#
# 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.

from argparse import Namespace
from typing import Optional
from os.path import join, exists

import torch
import bitsandbytes as bnb
from transformers import (
    AutoModelForCausalLM,
    BitsAndBytesConfig,
)

from peft import (
    prepare_model_for_kbit_training,
    LoraConfig,
    PeftModel,
    PeftModelForCausalLM,
)
from peft.tuners.lora import LoraLayer
from falcon_modeling.modelling_RW import RWForCausalLM
from models.llama_with_flash_attn import LlamaForCausalLM

REGISTERED_BASE_MODELS = {}


def find_all_linear_names(
    args: Namespace,
    model: torch.nn.Module,
):
    cls = (
        bnb.nn.Linear4bit
        if args.bits == 4
        else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
    )
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split(".")
            if "lora" not in names[-1]:
                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")
    return list(lora_module_names)


def get_accelerate_model(
    args: Namespace,
    checkpoint_dir: Optional[str] = None,
    adapter_name="lora_default",
    is_trainable=True,
    reuse_base_model=False,
):
    global REGISTERED_BASE_MODELS

    compute_dtype = (
        torch.float16 if args.fp16 else (torch.bfloat16 if args.bf16 else torch.float32)
    )

    if checkpoint_dir is not None:
        if exists(join(checkpoint_dir, "adapter_model")):
            checkpoint_dir = join(checkpoint_dir, "adapter_model")

        if exists(join(checkpoint_dir, "lora_default")):
            checkpoint_dir = join(checkpoint_dir, "lora_default")

    if args.model_name_or_path in REGISTERED_BASE_MODELS and reuse_base_model:
        config = {
            "load_in_4bit": args.bits == 4,
            "load_in_8bit": args.bits == 8,
            "llm_int8_threshold": 6.0,
            "llm_int8_has_fp16_weight": False,
            "bnb_4bit_compute_dtype": compute_dtype,
            "bnb_4bit_use_double_quant": args.double_quant,
            "bnb_4bit_quant_type": args.quant_type,
        }

        registered_model, registered_config = REGISTERED_BASE_MODELS[
            args.model_name_or_path
        ]
        if registered_config == config and not args.full_finetune:
            print(f"loading registered model {args.model_name_or_path}...")
            model = registered_model

            if checkpoint_dir is not None:
                model.load_adapter(
                    checkpoint_dir,
                    adapter_name=adapter_name,
                    is_trainable=is_trainable,
                )
            else:
                modules = args.lora_modules or find_all_linear_names(args, model)
                print("adding LoRa modules: ", modules)
                config = LoraConfig(
                    r=args.lora_r,
                    lora_alpha=args.lora_alpha,
                    target_modules=modules,
                    lora_dropout=args.lora_dropout,
                    bias="none",
                    task_type="CAUSAL_LM",
                )
                model.add_adapter(adapter_name, peft_config=config)
            return model
        else:
            raise ValueError(
                f"Model {args.model_name_or_path} is already registered with a different config."
                f"{registered_config} != {config}"
            )

    current_device = torch.cuda.current_device()
    if args.full_finetune:
        assert args.bits in [16, 32]

    print(f"loading base model {args.model_name_or_path}...")

    CausalLM = AutoModelForCausalLM

    if "falcon" in args.model_name_or_path.lower():
        CausalLM = RWForCausalLM
    elif (
        "llama" in args.model_name_or_path.lower()
        or "vicuna" in args.model_name_or_path.lower()
        or "dromedary" in args.model_name_or_path.lower()
    ) and torch.__version__ >= "2.0.0":
        CausalLM = LlamaForCausalLM

    model = CausalLM.from_pretrained(
        args.model_name_or_path,
        load_in_4bit=args.bits == 4,
        load_in_8bit=args.bits == 8,
        device_map={"": current_device},
        # max_memory=max_memory,
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=args.bits == 4,
            load_in_8bit=args.bits == 8,
            llm_int8_threshold=6.0,
            llm_int8_has_fp16_weight=False,
            bnb_4bit_compute_dtype=compute_dtype,
            bnb_4bit_use_double_quant=args.double_quant,
            bnb_4bit_quant_type=args.quant_type,
        ),
        torch_dtype=(
            torch.float16
            if args.fp16
            else (torch.bfloat16 if args.bf16 else torch.float32)
        ),
        trust_remote_code=args.trust_remote_code,
    )
    if compute_dtype == torch.float16 and args.bits == 4:
        major, minor = torch.cuda.get_device_capability()
        if major >= 8:
            print("=" * 80)
            print(
                "Your GPU supports bfloat16, you can accelerate training with the argument --bf16"
            )
            print("=" * 80)

    setattr(model, "model_parallel", True)
    setattr(model, "is_parallelizable", True)

    model.config.torch_dtype = (
        torch.float16 if args.fp16 else (torch.bfloat16 if args.bf16 else torch.float32)
    )

    if not args.full_finetune:
        model = prepare_model_for_kbit_training(
            model, use_gradient_checkpointing=args.gradient_checkpointing
        )
    if args.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    if not args.full_finetune:
        if checkpoint_dir is not None:
            print("Loading adapters from checkpoint.")

            model = PeftModel.from_pretrained(
                model,
                checkpoint_dir,
                adapter_name=adapter_name,
                is_trainable=is_trainable,
            )
        else:
            # print(f'adding LoRA modules...')
            modules = args.lora_modules or find_all_linear_names(args, model)
            print("adding LoRa modules: ", modules)
            config = LoraConfig(
                r=args.lora_r,
                lora_alpha=args.lora_alpha,
                target_modules=modules,
                lora_dropout=args.lora_dropout,
                bias="none",
                task_type="CAUSAL_LM",
            )
            model = get_peft_model(model, config, adapter_name=adapter_name)

        if args.model_name_or_path not in REGISTERED_BASE_MODELS:
            config = {
                "load_in_4bit": args.bits == 4,
                "load_in_8bit": args.bits == 8,
                "llm_int8_threshold": 6.0,
                "llm_int8_has_fp16_weight": False,
                "bnb_4bit_compute_dtype": compute_dtype,
                "bnb_4bit_use_double_quant": args.double_quant,
                "bnb_4bit_quant_type": args.quant_type,
            }
            REGISTERED_BASE_MODELS[args.model_name_or_path] = (model, config)

    for name, module in model.named_modules():
        if isinstance(module, LoraLayer):
            if args.bf16:
                module = module.to(torch.bfloat16)
            else:
                module = module.to(torch.float32)
        if "lm_head" in name or "embed_tokens" in name:
            if hasattr(module, "weight"):
                if args.bf16 and module.weight.dtype == torch.float32:
                    module = module.to(torch.bfloat16)
                # if not args.bf16:
                #     module = module.to(torch.float32)
    return model


def load_4bit_model_for_inference(
    checkpoint_dir: str,
    bits: int = 4,
    fp16: bool = False,
    bf16: bool = False,
    double_quant: bool = True,
    quant_type: str = "nf4",
    gradient_checkpointing: bool = False,
    adapter_name="lora_default",
    is_trainable=True,
    reuse_base_model=False,
    trust_remote_code=False,
    base_model_mapping=None,
    fully_initialize=False,
):
    if checkpoint_dir is not None:
        if exists(join(checkpoint_dir, "adapter_model")):
            checkpoint_dir = join(checkpoint_dir, "adapter_model")

        if exists(join(checkpoint_dir, "lora_default")):
            checkpoint_dir = join(checkpoint_dir, "lora_default")

    config = LoraConfig.from_pretrained(checkpoint_dir)
    base_model_name_or_path = config.base_model_name_or_path

    if base_model_mapping is not None:
        dict_base_model_mapping = eval(base_model_mapping)
        if (
            dict_base_model_mapping is not None
            and base_model_name_or_path in dict_base_model_mapping
        ):
            base_model_name_or_path = dict_base_model_mapping[base_model_name_or_path]

    args = Namespace(
        model_name_or_path=base_model_name_or_path,
        bits=bits,
        fp16=fp16,
        bf16=bf16,
        double_quant=double_quant,
        quant_type=quant_type,
        gradient_checkpointing=gradient_checkpointing,
        trust_remote_code=trust_remote_code,
        full_finetune=False,
        lora_r=64 if fully_initialize else None,
        lora_alpha=16 if fully_initialize else None,
        lora_dropout=0.0 if fully_initialize else None,
        lora_modules=None,
    )

    if fully_initialize:
        print("Fully initializing qlora model.")

    model = get_accelerate_model(
        args,
        checkpoint_dir=None if fully_initialize else checkpoint_dir,
        adapter_name=adapter_name,
        is_trainable=is_trainable,
        reuse_base_model=reuse_base_model,
    )
    return model


def get_peft_model(model, peft_config, adapter_name="default"):
    """
    Returns a Peft model object from a model and a config.

    Args:
        model ([`transformers.PreTrainedModel`]): Model to be wrapped.
        peft_config ([`PeftConfig`]): Configuration object containing the parameters of the Peft model.
    """
    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)
    return PeftModelForCausalLM(model, peft_config, adapter_name=adapter_name)

training/models/qlora_model.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
# Copyright 2023 The Alpaca Team
#
# 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.

from argparse import Namespace
import os
from typing import Optional, Dict, Sequence, Union

import einops
import torch
from torch import Tensor, nn
import torch.nn.functional as F

import transformers
from transformers.trainer_utils import EvalPrediction
from transformers.utils.generic import ModelOutput

from peft import PeftModel, LoraModel, LoraConfig

from models.qlora_model import get_accelerate_model
from models.llama_with_flash_attn import LlamaModel as FlashAttnLlamaModel


def unpack_dict(
    d: Dict, keys: Sequence[str], return_type: type = tuple
) -> Union[Sequence, Dict]:
    if return_type in (tuple, list):
        return return_type(d[key] for key in keys)
    elif return_type == dict:
        return {key: d[key] for key in keys}
    else:
        raise ValueError(f"Unknown return_type: {return_type}")


def batch_select(input: Tensor, index: Tensor):
    """Select elements from a batched tensor with a batched index tensor.

    Example:
        input = torch.tensor([
            [0, 1, 2],
            [3, 0, 9],
            [6, 7, 8],
        ])
        index = torch.tensor([[0, 1], [1, 0], [0, 0]])
        batch_select(input, index) = tensor([
            [0, 1],
            [0, 3],
            [6, 6]
        ])
    """
    dummy_index = torch.arange(input.size(0), device=input.device).unsqueeze(-1)
    return input[dummy_index, index]


def make_generative_lm(
    args: Namespace,
    model_name_or_path: str,
    qlora: bool = False,
    checkpoint_dir: Optional[str] = None,
    adapter_name="lora_default",
    is_trainable=True,
    reuse_base_model=False,
    **kwargs,
):
    model_cls = transformers.LlamaForCausalLM

    if qlora:
        if checkpoint_dir is None or checkpoint_dir in ["scratch", "none"]:
            return get_accelerate_model(args, None)
        else:
            return get_accelerate_model(
                args,
                checkpoint_dir=checkpoint_dir,
                adapter_name=adapter_name,
                is_trainable=is_trainable,
                reuse_base_model=reuse_base_model,
            )
    return model_cls.from_pretrained(model_name_or_path, **kwargs)


def get_transformer_hidden_size(model: transformers.PreTrainedModel):
    if isinstance(model, PeftModel):
        return get_transformer_hidden_size(model.base_model)

    if isinstance(model, LoraModel):
        return get_transformer_hidden_size(model.model)

    if isinstance(model, transformers.GPT2LMHeadModel):
        hidden_size_attr_name = "n_embd"
    elif isinstance(model, transformers.OPTForCausalLM):
        hidden_size_attr_name = "word_embed_proj_dim"
    elif isinstance(model, transformers.T5ForConditionalGeneration):
        hidden_size_attr_name = "d_model"
    elif isinstance(model, transformers.GPTBigCodeForCausalLM):
        hidden_size_attr_name = "n_embd"
    elif "modelling_RW.RWModel" in str(
        type(model)
    ) or "modelling_RW.RWForCausalLM" in str(type(model)):
        # TODO(zhiqings): Hack to add support for Falcon.
        hidden_size_attr_name = "hidden_size"
    else:
        # Hack to deal with the fact that transformers library changed the LLaMA model name.
        llama_cls = getattr(
            transformers,
            "LLaMAForCausalLM"
            if hasattr(transformers, "LLaMAForCausalLM")
            else "LlamaForCausalLM",
        )
        if isinstance(model, llama_cls) or "LlamaForCausalLM" in str(type(model)):
            hidden_size_attr_name = "hidden_size"
        else:
            raise ValueError(f"Unknown base_model type: {type(model)}")
        from typing import Any, Mapping
    return getattr(model.config, hidden_size_attr_name)


class RewardConfig(transformers.PretrainedConfig):
    model_type = "reward_model"

    # Huggingface doesn't allow non-kwargs for `__init__`.
    def __init__(self, backbone_model_name_or_path=None, **kwargs):
        super(RewardConfig, self).__init__(**kwargs)
        self.backbone_model_name_or_path = backbone_model_name_or_path


class RewardModelOutput(ModelOutput):
    rewards: Tensor = None


class RewardModel(transformers.PreTrainedModel):
    config_class = RewardConfig
    supports_gradient_checkpointing = True

    def __init__(
        self,
        args: Namespace,
        config: RewardConfig,
        checkpoint_dir: Optional[str] = None,
        adapter_name="lora_default",
        **kwargs,
    ):
        super(RewardModel, self).__init__(config)
        self.adapter_name = adapter_name
        self.backbone_model = make_generative_lm(
            args,
            config.backbone_model_name_or_path,
            checkpoint_dir=checkpoint_dir,
            adapter_name=adapter_name,
            **kwargs,
        )
        hidden_size = get_transformer_hidden_size(self.backbone_model)
        reward_head = nn.Linear(hidden_size, 1)
        torch.nn.init.zeros_(reward_head.bias)
        device = next(self.backbone_model.parameters()).device
        self.reward_head = reward_head.to(device)

        if checkpoint_dir is not None:
            reward_head_path = os.path.join(checkpoint_dir, "reward_head")
            if os.path.exists(reward_head_path):
                self.reward_head.load_state_dict(
                    torch.load(
                        reward_head_path,
                        map_location="cpu",
                    )
                )
            else:
                print(f"Warning: reward head not found at {reward_head_path}")

        self.reward_head.requires_grad_(kwargs.get("is_trainable", True))

    def forward(self, input_ids, attention_mask=None, return_dict=True, **kwargs):
        # We only compute the rewards and don't compute the logistic regression loss in this function so that it's
        # easier to use for later stages of reranking / RL training.
        self.backbone_model.set_adapter(self.adapter_name)
        self.backbone_model.config.use_cache = False

        outputs = self.backbone_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
            output_hidden_states=True,
            **kwargs,
        )
        last_hidden_state = outputs.hidden_states[-1]
        assert isinstance(last_hidden_state, torch.Tensor), f"{outputs}"
        # last_hidden_state = outputs.last_hidden_state
        # TODO(zhiqings): Hacking to make sure every parameter is used in the backward pass.
        logits = outputs.logits
        last_hidden_state = last_hidden_state + 0.0 * torch.mean(logits)

        last_hidden_state_at_the_end = last_hidden_state[:, -1, :]
        # TODO(lxuechen): Make returning rewards at all positions and last_hidden_state an option.
        final_dtype_tensor = next(self.reward_head.parameters())
        last_hidden_state_at_the_end = last_hidden_state_at_the_end.type_as(
            final_dtype_tensor
        )

        if final_dtype_tensor.dtype == torch.float16:
            last_hidden_state_at_the_end = last_hidden_state_at_the_end.to(
                torch.float32
            )

        rewards = self.reward_head(last_hidden_state_at_the_end).squeeze(-1)
        return RewardModelOutput(rewards=rewards) if return_dict else (rewards,)

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, transformers.LlamaModel):
            module.gradient_checkpointing = value
        elif isinstance(module, FlashAttnLlamaModel):
            module.gradient_checkpointing = value
        if isinstance(module, transformers.GPTBigCodeModel):
            module.gradient_checkpointing = value
        # TODO(zhiqings): Hack to add support for Falcon.
        if "RWModel" in str(type(module)):
            module.gradient_checkpointing = value


class RewardModelTrainer(transformers.Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # input_ids, attention_mask each of size (bsz, num_candidates, seq_len).
        # index_0, index_1 each of size (bsz, num_pairs); indexes into input_ids.
        # choice of size (bsz, num_pairs); 1 if index_1's seq is chosen, 0 otherwise.
        input_ids, attention_mask, index_0, index_1, choice = unpack_dict(
            inputs, keys=("input_ids", "attention_mask", "index_0", "index_1", "choice")
        )
        num_candidates, num_pairs = input_ids.size(1), choice.size(1)
        input_ids_flat, attention_mask_flat = tuple(
            einops.rearrange(x, "b c l -> (b c) l") for x in (input_ids, attention_mask)
        )
        outputs = model(input_ids=input_ids_flat, attention_mask=attention_mask_flat)
        rewards_flat = outputs.rewards
        rewards = einops.rearrange(
            rewards_flat, "(b c) -> b c", c=num_candidates
        )  # Size: (bsz, num_candidates).

        rewards_0, rewards_1 = tuple(
            batch_select(rewards, index) for index in (index_0, index_1)
        )  # Size: (bsz, num_pairs).
        logits = rewards_1 - rewards_0  # Size: (bsz, num_pairs).
        # Type casting of `choice` is due to amp.autocast context manager.
        loss = F.binary_cross_entropy_with_logits(
            logits, choice.to(logits.dtype), reduction="mean"
        )
        return (loss, dict(logits=logits)) if return_outputs else loss


def compute_reward_modeling_metrics(eval_prediction: EvalPrediction) -> Dict:
    # eval_prediction.label_ids is a tuple that matches up with `training_args.label_names`.
    logits = torch.tensor(eval_prediction.predictions).squeeze(-1)
    labels = torch.tensor(eval_prediction.label_ids[-1]).squeeze(-1)
    predictions = (logits >= 0.0).long()
    accuracy = predictions.eq(labels).float().mean().item()
    label_positive_rate = (labels == 1).float().mean().item()
    return dict(
        accuracy=accuracy,
        label_positive_rate=label_positive_rate,
    )


def load_4bit_reward_model_for_inference(
    checkpoint_dir: str,
    bits: int = 4,
    fp16: bool = False,
    bf16: bool = False,
    double_quant: bool = True,
    quant_type: str = "nf4",
    gradient_checkpointing: bool = False,
    adapter_name="lora_default",
    is_trainable=True,
    reuse_base_model=False,
    trust_remote_code=False,
    base_model_mapping=None,
):
    # Load the model.
    lora_checkpoint_dir = checkpoint_dir
    if os.path.exists(os.path.join(lora_checkpoint_dir, "adapter_model")):
        lora_checkpoint_dir = os.path.join(lora_checkpoint_dir, "adapter_model")
    if os.path.exists(os.path.join(lora_checkpoint_dir, "lora_default")):
        lora_checkpoint_dir = os.path.join(lora_checkpoint_dir, "lora_default")

    lora_config = LoraConfig.from_pretrained(lora_checkpoint_dir)
    base_model_name_or_path = lora_config.base_model_name_or_path

    if base_model_mapping is not None:
        dict_base_model_mapping = eval(base_model_mapping)
        if (
            dict_base_model_mapping is not None
            and base_model_name_or_path in dict_base_model_mapping
        ):
            base_model_name_or_path = dict_base_model_mapping[base_model_name_or_path]

    config = RewardConfig(backbone_model_name_or_path=base_model_name_or_path)

    args = Namespace(
        model_name_or_path=config.backbone_model_name_or_path,
        bits=bits,
        fp16=fp16,
        bf16=bf16,
        double_quant=double_quant,
        quant_type=quant_type,
        trust_remote_code=trust_remote_code,
        full_finetune=False,
        gradient_checkpointing=gradient_checkpointing,
    )

    model = RewardModel(
        args,
        config,
        checkpoint_dir=checkpoint_dir,
        qlora=bits == 4 or bits == 8,
        adapter_name=adapter_name,
        is_trainable=is_trainable,
        reuse_base_model=reuse_base_model,
    )
    return model

training/models/reward_model.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
# Copyright 2023 The Alpaca Team
#
# 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.

"""Model classes that are shared across different algorithms.

WARNING:
    Do not tamper with the state_dict function for any of these classes.
    If you tamper, make sure the keys are the same, otherwise FSDP will get confused.
"""

import abc
import logging
from typing import Dict, Optional

import torch
import transformers
from torch import Tensor, nn

from data_utils.common_utils import right_pad, compute_logprobs
from models.reward_model import get_transformer_hidden_size


logger = logging.getLogger(__name__)


class Policy(nn.Module, abc.ABC):
    def __init__(
        self,
        args,
        base_model: transformers.PreTrainedModel,
        base_tokenizer: transformers.PreTrainedTokenizer,
        adapter_name: Optional[str] = None,
    ):
        super().__init__()
        self.args = args
        self.base_model = base_model
        self.base_tokenizer = base_tokenizer
        self.adapter_name = adapter_name

    @abc.abstractmethod
    def forward(
        self,
        queries: Tensor,
        query_attn_masks: Tensor,
        responses: Tensor,
        temperature: Optional[float] = None,
    ) -> Dict[str, Tensor]:
        raise NotImplementedError

    def respond(
        self,
        queries: Tensor,
        query_attn_masks: Tensor,
        temperature: Optional[float] = None,
        num_return_sequences=1,
    ) -> Dict[str, Tensor]:
        assert not self.training, "Policy must be in eval model for generation."
        return self._post_respond(
            self._respond(queries, query_attn_masks, temperature, num_return_sequences)
        )

    @abc.abstractmethod
    def _respond(
        self,
        queries: Tensor,
        query_attn_masks: Tensor,
        temperature: Optional[float] = None,
        num_return_sequences=1,
    ) -> Dict[str, Tensor]:
        raise NotImplementedError

    def _post_respond(self, respond_outputs: Dict[str, Tensor]) -> Dict[str, Tensor]:
        return respond_outputs


class AutoregressivePolicy(Policy):
    def forward(
        self,
        queries: Tensor,
        query_attn_masks: Tensor,
        responses: Tensor,
        temperature: Optional[float] = None,
    ) -> Dict[str, Tensor]:
        # TODO(lxuechen): Refactor attention mask. Here query_attn_masks overrides padding-based attention mask.

        if self.adapter_name is not None:
            self.base_model.set_adapter(self.adapter_name)
        self.base_model.config.use_cache = False

        if temperature is None:
            temperature = self.args.temperature
        input_ids = torch.cat([queries, responses], dim=1)
        attention_mask = input_ids.ne(self.base_tokenizer.pad_token_id)
        attention_mask[:, : queries.size(1)] = query_attn_masks
        # Fix position id issues and ensure consistency with `respond` for GPT and OPT.
        inputs = self.base_model.prepare_inputs_for_generation(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=False,
        )
        outputs = self.base_model(**inputs, output_hidden_states=True)
        original_logits = outputs.logits[:, -self.args.response_len - 1 : -1]
        logits = original_logits / temperature
        labels = input_ids[:, -self.args.response_len :]
        logprobs = compute_logprobs(
            logits, labels, ignore_index=self.base_tokenizer.pad_token_id
        )
        entropies = -(logits.softmax(dim=-1) * logits.log_softmax(dim=-1)).sum(dim=-1)
        last_hidden_state = outputs.hidden_states[-1][
            :, -self.args.response_len - 1 : -1
        ]
        return dict(
            original_logits=original_logits,
            logits=logits,
            logprobs=logprobs,
            entropies=entropies,
            last_hidden_state=last_hidden_state,
        )

    def _respond(
        self,
        queries: Tensor,
        query_attn_masks: Tensor,
        temperature: Optional[float] = None,
        num_return_sequences=1,
    ) -> Dict[str, Tensor]:
        if self.adapter_name is not None:
            self.base_model.set_adapter(self.adapter_name)
        self.base_model.config.use_cache = True
        self.base_model.config.cache_shape = (
            queries.shape[-1] + self.args.response_len,
        )

        if temperature is None:
            temperature = self.args.temperature
        sequences = self.base_model.generate(
            inputs=queries,
            attention_mask=query_attn_masks,
            do_sample=True,
            max_new_tokens=self.args.response_len,
            pad_token_id=self.base_tokenizer.pad_token_id,
            eos_token_id=self.base_tokenizer.eos_token_id,
            top_p=1.0,
            top_k=0,
            temperature=temperature,
            num_return_sequences=num_return_sequences,
            synced_gpus=True,
        )
        responses = right_pad(
            sequences[:, queries.size(1) :],
            target_size=(sequences.size(0), self.args.response_len),
            value=self.base_tokenizer.pad_token_id,
        )
        return dict(
            responses=responses
        )  # Size (bsz * num_return_sequences, response_len).


class Value(nn.Module, abc.ABC):
    def __init__(
        self,
        args,
        base_model: transformers.PreTrainedModel,
        base_tokenizer: transformers.PreTrainedTokenizer,
        adapter_name: Optional[str] = None,
    ):
        super().__init__()
        self.args = args
        self.base_model = base_model
        self.base_tokenizer = base_tokenizer
        hidden_size = get_transformer_hidden_size(base_model)
        value_head = torch.nn.Linear(hidden_size, 1)
        value_head.weight.data.zero_()
        value_head.bias.data.zero_()
        self.value_head = value_head.to(next(base_model.parameters()).device)
        self.adapter_name = adapter_name

    @abc.abstractmethod
    def forward(
        self, queries: Tensor, query_attn_masks: Tensor, responses: Tensor
    ) -> Dict[str, Tensor]:
        raise NotImplementedError


class AutoregressiveValue(Value):
    def forward(
        self, queries: Tensor, query_attn_masks: Tensor, responses: Tensor
    ) -> Dict[str, Tensor]:
        if self.adapter_name is not None:
            self.base_model.set_adapter(self.adapter_name)
        self.base_model.config.use_cache = False

        sequences = torch.cat([queries, responses], dim=1)
        sequence_attn_masks = sequences.ne(self.base_tokenizer.pad_token_id)

        inputs = self.base_model.prepare_inputs_for_generation(
            input_ids=sequences,
            attention_mask=sequence_attn_masks,
            use_cache=False,
        )
        outputs = self.base_model(
            **inputs,
            return_dict=True,
            output_hidden_states=True,
        )
        # value[t]: \hat{V}(sequences_{:t-1}); must align with `_estimate_advantage`.

        last_hidden_state = outputs.hidden_states[-1]
        assert isinstance(last_hidden_state, torch.Tensor), f"{outputs}"
        logits = outputs.logits
        # TODO(zhiqings): Hacking to make sure every parameter is used in the backward pass.
        last_hidden_state = last_hidden_state + 0.0 * torch.mean(logits)
        last_hidden_state = last_hidden_state[:, queries.size(1) - 1 : -1]

        # TODO(zhiqings): now we just manully convert output types
        last_hidden_state = last_hidden_state.type_as(
            next(self.value_head.parameters())
        )
        values = self.value_head(last_hidden_state).squeeze(-1)
        return dict(values=values)


class ActorCritic(nn.Module):
    def __init__(self, policy: Policy, value_model: Value):
        super(ActorCritic, self).__init__()
        self.policy = policy
        self.value_model = value_model

    def forward(
        self,
        queries: Tensor,
        query_attn_masks: Tensor,
        responses: Tensor,
        temperature: Optional[float] = None,
        mode: Optional[str] = None,
    ) -> Dict[str, Tensor]:
        # Assume the policy and value model share the same tokenizer.

        if mode is None:
            o1 = self.policy(queries, query_attn_masks, responses, temperature)
            o2 = self.value_model(queries, query_attn_masks, responses)

        elif mode == "policy":
            o1 = self.policy(queries, query_attn_masks, responses, temperature)
            # Add dummy loss to make sure every parameter is used in the backward pass.
            o2 = {
                "dummy_loss": 0.0
                * torch.sum(
                    torch.stack(
                        [
                            torch.mean(value)
                            for key, value in self.named_parameters()
                            if "lora_value" in key
                        ]
                    )
                )
            }
        elif mode == "value":
            o2 = self.value_model(queries, query_attn_masks, responses)
            # Add dummy loss to make sure every parameter is used in the backward pass.
            o1 = {
                "dummy_loss": 0.0
                * torch.sum(
                    torch.stack(
                        [
                            torch.mean(value)
                            for key, value in self.named_parameters()
                            if "lora_policy" in key
                        ]
                    )
                )
            }
        else:
            raise ValueError(f"Unknown mode: {mode}")

        return {**o1, **o2}

    def respond(
        self,
        queries: Tensor,
        query_attn_masks: Tensor,
        temperature: Optional[float] = None,
    ) -> Dict[str, Tensor]:
        return self.policy.respond(
            queries=queries, query_attn_masks=query_attn_masks, temperature=temperature
        )


def make_policy_with_base_model(
    args,
    base_model: transformers.PreTrainedModel,
    base_tokenizer: transformers.PreTrainedTokenizer,
    adapter_name: Optional[str] = "default",
) -> Policy:
    if base_model.config.is_encoder_decoder:
        raise NotImplementedError
    else:
        return AutoregressivePolicy(
            args, base_model, base_tokenizer, adapter_name=adapter_name
        )


def make_value_with_base_model(
    args,
    base_model: transformers.PreTrainedModel,
    base_tokenizer: transformers.PreTrainedTokenizer,
    adapter_name: Optional[str] = "default",
) -> Value:
    if base_model.config.is_encoder_decoder:
        raise NotImplementedError
    else:
        return AutoregressiveValue(
            args, base_model, base_tokenizer, adapter_name=adapter_name
        )

training/models/rl_models.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
# Copyright 2023 The Alpaca Team
#
# 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 abc
import copy
import dataclasses
import json
import logging
import math
import os
from pathlib import Path
import random
import sys
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union

import einops
import tqdm

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

import accelerate
from accelerate.optimizer import AcceleratedOptimizer
from accelerate.utils import convert_outputs_to_fp32

import transformers
from transformers.trainer_utils import enable_full_determinism, set_seed

from data_utils.data_utils_ppo import QueryResponseDataset
import data_utils.common_utils as utils
import models.distributed_utils as distributed_utils
from models.trainer_utils import create_optimizer, create_scheduler


logger = logging.getLogger(__name__)

if torch.__version__ < "2.0.0":
    LRScheduler = torch.optim.lr_scheduler._LRScheduler  # noqa
else:
    LRScheduler = torch.optim.lr_scheduler.LRScheduler

FIRST_STEP_IDX = 1


class KLController(abc.ABC):
    value: Union[int, float]

    def step(self, *args, **kwargs):
        pass


class FixedKLController(KLController):
    def __init__(self, kl_coef):
        super(FixedKLController, self).__init__()
        self.value = kl_coef


@dataclasses.dataclass
class HFDecodingArguments:
    """Only the core args for decoding with HF models."""

    top_p: float = 0.9
    top_k: int = 0
    temperature: float = 1.0
    do_sample: bool = True
    num_beams: int = 1
    max_new_tokens: int = (
        100  # This is aligned with `openai_utils.OpenAIDecodingArguments`.
    )
    num_return_sequences: int = 1


class AlpacaAccelerator(accelerate.Accelerator):
    """Thin wrapper for accelerate.Accelerator."""

    def __repr__(self):
        return (
            f"Accelerator(\n"
            f"  state={self.state}, \n"
            f"  gradient_accumulation_steps={self.gradient_accumulation_steps:.6f}, \n"
            f"  split_batches={self.split_batches}, \n"
            f"  step_scheduler_with_optimizer={self.step_scheduler_with_optimizer},\n"
            f")"
        )

    def unwrap_optimizer(self, optimizer: accelerate.accelerator.AcceleratedOptimizer):
        return optimizer.optimizer


class RLTrainer(object):
    def __init__(
        self,
        args,
        train_dataset: QueryResponseDataset,
        eval_dataset: QueryResponseDataset,
        data_collator: Callable,
        tokenizer: transformers.PreTrainedTokenizer,
        policy: nn.Module,
        accelerator: AlpacaAccelerator,
        ref_policy: Optional[nn.Module] = None,
        reward_model: Optional[nn.Module] = None,
        optimizer: Optional[torch.optim.Optimizer] = None,
        lr_scheduler: Optional[LRScheduler] = None,
        reward_tokenizer: Optional[transformers.PreTrainedTokenizer] = None,
    ):
        super(RLTrainer, self).__init__()
        self.args = args
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.data_collator = data_collator
        self.policy = policy
        self.ref_policy = ref_policy
        self.reward_model = reward_model
        self.tokenizer = tokenizer
        self.optimizer = optimizer
        self.accelerator = accelerator
        self.lr_scheduler = lr_scheduler
        self.kl_ctl = FixedKLController(kl_coef=args.kl_coef)
        self.log_history = []
        self.args.set_truncate_token_ids(self.tokenizer)
        enable_full_determinism(
            self.args.seed
        ) if self.args.full_determinism else set_seed(self.args.seed)
        self.policy_meta_prompts = None
        self.reward_meta_prompts = None
        self.principle_collection = None

        if reward_tokenizer is None:
            self.reward_tokenizer = self.tokenizer
        else:
            self.reward_tokenizer = reward_tokenizer

    @abc.abstractmethod
    @torch.inference_mode()
    def rollout(self, queries_data) -> Dict[str, torch.Tensor]:
        raise NotImplementedError

    @abc.abstractmethod
    def compute_loss(
        self, rollouts: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict]:
        raise NotImplementedError

    @abc.abstractmethod
    @torch.inference_mode()
    def record_step_stats(self, train_stats, rollouts, step_idx, **kwargs):
        raise NotImplementedError

    @property
    def optimizable_params(self):
        return [
            p
            for p in self.policy.parameters()
            if p.requires_grad and p.grad is not None
        ]

    @torch.inference_mode()
    def _compute_grad_norm(self):
        grad_norm = torch.stack([p.grad.norm(2) for p in self.optimizable_params]).norm(
            2
        )
        return grad_norm

    @torch.inference_mode()
    def _compute_param_norm(self):
        param_norm = torch.stack([p.norm(2) for p in self.optimizable_params]).norm(2)
        return param_norm

    def step_with_rollouts(self, rollouts):
        """Based on fixed rollouts, run PPO for multiple epochs."""
        assert isinstance(self.optimizer, AcceleratedOptimizer), (
            "`optimizer` must be pushed through `accelerator.prepare`. "
            "Otherwise the `accelerator.accumulate` context manager won't correctly disable `zero_grad` or `step`."
        )
        rollouts_dataloader = self.get_rollouts_dataloader(rollouts=rollouts)
        stats_list = []
        for epoch_idx in range(self.args.noptepochs):
            for batch_idx, rollouts_batch in tqdm.tqdm(
                enumerate(rollouts_dataloader, 1),
                total=len(rollouts_dataloader),
                disable=not self.accelerator.is_main_process,
                desc="gradstep",
            ):
                with self.accelerator.accumulate(self.policy):
                    stats_for_this_step = {}
                    with self.accelerator.no_sync(self.policy):
                        policy_loss, policy_stats = self.compute_policy_loss(
                            rollouts_batch
                        )
                        stats_for_this_step.update(policy_stats)
                        self.accelerator.backward(policy_loss)

                    value_loss, value_stats = self.compute_value_loss(rollouts_batch)
                    stats_for_this_step.update(value_stats)
                    self.accelerator.backward(value_loss)

                    if self.accelerator.sync_gradients:
                        # Gradient norm almost blows up at some point, but stabilizes eventually, even w/o clipping.
                        if self.args.max_grad_norm is not None:
                            self.accelerator.clip_grad_norm_(
                                self.policy.parameters(), self.args.max_grad_norm
                            )
                        stats_for_this_step[
                            "loss/grad_norm"
                        ] = self._compute_grad_norm()
                        stats_list.append(stats_for_this_step)
                    self.optimizer.step()
                    self.optimizer.zero_grad(set_to_none=True)

        return utils.merge_dict(
            stats_list, torch.stack
        )  # list of dict -> dict: str -> 1-D tensor

    def step(self, train_dataloader, step_idx: int):
        queries_batches = [
            next(train_dataloader) for _ in range(self.args.rollout_accumulation_steps)
        ]
        rollouts = self.rollout(queries_batches)
        train_stats = self.step_with_rollouts(rollouts)
        if self.lr_scheduler is not None:
            self.lr_scheduler.step()
        stats = self.record_step_stats(
            rollouts=rollouts,
            train_stats=train_stats,
            step_idx=step_idx,
            kl_coef=self.kl_ctl.value,
        )
        self.kl_ctl.step(stats["objective/kl_sum_seq"])
        return stats

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        optimizer = create_optimizer(
            args=self.args, model=self.policy, optimizer=self.optimizer
        )
        lr_scheduler = create_scheduler(
            args=self.args,
            optimizer=optimizer,
            lr_scheduler=self.lr_scheduler,
            num_training_steps=num_training_steps,
        )
        self.optimizer, self.lr_scheduler = self.accelerator.prepare(
            optimizer, lr_scheduler
        )
        self.accelerator.register_for_checkpointing(
            self.lr_scheduler
        )  # LR scheduler needs another call to save.
        return self.optimizer, self.lr_scheduler

    def train(self, resume_training_ckpt: Optional[str] = None):
        """Entry point for training."""
        total_epochs = self.args.total_epochs
        total_episodes = len(self.train_dataset) * total_epochs  # noqa
        total_steps = total_episodes // self.args.rollout_batch_size  # noqa
        logger.warning(
            f"***Training starts***\n"
            f"Total epochs: {total_epochs} => Total episodes: {total_episodes} => Total steps: {total_steps}"
        )

        self.create_optimizer_and_scheduler(total_steps)

        skipping_steps = 0
        if resume_training_ckpt is not None:
            skipping_steps = self.resume_training(resume_training_ckpt)
            print(
                f"Resuming training from {resume_training_ckpt} at step {skipping_steps}."
            )

        infinite_train_dataloader = self.get_train_dataloader()
        for step_idx in tqdm.tqdm(
            range(FIRST_STEP_IDX, total_steps + FIRST_STEP_IDX),
            disable=not self.accelerator.is_main_process,
            desc="steps",
            total=total_steps,
        ):
            if step_idx < skipping_steps:
                for _ in range(self.args.rollout_accumulation_steps):
                    next(infinite_train_dataloader)
                continue

            if (
                step_idx % self.args.save_steps == 0
                or step_idx in self.args.save_steps_extra_list
            ):
                if step_idx > skipping_steps:
                    self.save_model(
                        os.path.join(self.args.output_dir, f"checkpoint-{step_idx}")
                    )
            if (
                self.args.eval_steps is not None
                and step_idx % self.args.eval_steps == 0
            ):
                self.evaluate(step_idx)
            self.log_history.append(self.step(infinite_train_dataloader, step_idx))
        return self.log_history

    @torch.inference_mode()
    def evaluate(self, step_idx: int, unwrapped_policy=None):
        """Evaluate by generating sequences with test prefixes.

        FSDP compat: all devices should do the forward pass, since sharded params need to be summoned.
                     only write results in the main process.
        """
        # TODO: unhardcode inference args.
        logger.warning(
            f"Start evaluation at step: {step_idx}",
            # main_process_only=True,
        )

        pure_queries, prompts, list_dict_data = (
            self.eval_dataset.pure_queries,
            self.eval_dataset.prompts,
            self.eval_dataset.list_dict_data,
        )
        if any(item is None for item in (pure_queries, prompts, list_dict_data)):
            logger.warning(
                "No evaluation data, skipping evaluation.",
                # main_process_only=True
            )
            return

        # Constants.
        model_name = Path(
            self.args.output_dir
        ).stem  # Don't use the helper in common, as no checkpoint is saved yet.
        model_name_at_step = f"{model_name}_ckpt_{step_idx}"
        temperature = 0.7
        del model_name

        # Start evaluation.
        self.policy.eval()
        if unwrapped_policy is None:
            unwrapped_policy = self.accelerator.unwrap_model(
                self.policy, keep_fp32_wrapper=True
            )
            unwrapped_policy = unwrapped_policy.policy.base_model

        outputs = decode_prompts_with_huggingface_given_model(
            model=unwrapped_policy,
            tokenizer=self.tokenizer,
            prompts=prompts,
            decoding_args=HFDecodingArguments(
                max_new_tokens=self.args.response_len, temperature=temperature
            ),
            per_device_batch_size=self.args.per_device_eval_batch_size,
            divide_work=False,
        )

        # sequences = [i + o for i, o in utils.zip_(prompts, outputs)]
        sequences = [
            self.prepare_reward_inputs(inputs=q, outputs=r)
            for q, r in utils.zip_(pure_queries, outputs)
        ]

        rewards = score_sequences_with_huggingface_given_model(
            model=self.reward_model,
            tokenizer=self.tokenizer,
            sequences=sequences,
            per_device_batch_size=self.args.rollout_per_device_batch_size,
            divide_work=False,
        )

        if self.accelerator.is_main_process:
            results = [
                {"reward": reward, model_name_at_step: output, **example}
                for reward, output, example in utils.zip_(
                    rewards, outputs, list_dict_data
                )
            ]
            if self.args.output_dir is not None:
                with open(
                    os.path.join(self.args.output_dir, f"eval_results_{step_idx}.json"),
                    "w",
                    encoding="utf-8",
                ) as f:
                    json.dump(results, f, indent=2)

            logger.warning(
                f"End evaluation at step: {step_idx}. Processed {len(results)} examples"
            )

    @abc.abstractmethod
    @torch.inference_mode()
    def save_model(self, output_dir: Optional[str] = None):
        raise NotImplementedError

    @abc.abstractmethod
    @torch.inference_mode()
    def resume_training(self, checkpoint_dir: str):
        raise NotImplementedError

    def _log_batch_size(self, loader: DataLoader, loader_name):
        batch = next(iter(loader))
        if isinstance(batch, torch.Tensor):
            batch_size = batch.shape[0]
        elif isinstance(batch, (list, tuple)):
            batch_size = batch[0]
        else:
            tensor = list(batch.values())[0]
            batch_size = tensor.size(0)
        logger.warning(
            f"Batch size of {loader_name} dataloader: {batch_size}",
            # main_process_only=True,
        )

    def get_train_dataloader(self):
        logger.warning(
            f"Train dataset size: {len(self.train_dataset)}",
            # main_process_only=True
        )  # noqa
        train_dataloader = DataLoader(
            dataset=self.train_dataset,
            collate_fn=self.data_collator,
            batch_size=self.args.rollout_per_device_batch_size,
            shuffle=True,
            drop_last=True,
        )
        train_dataloader = self.accelerator.prepare(train_dataloader)  # noqa
        self._log_batch_size(train_dataloader, "train_dataloader")
        return utils.InfiniteLoader(train_dataloader)

    def get_rollouts_dataloader(
        self, rollouts: Dict[str, torch.Tensor], shuffle=True, drop_last=True, keys=None
    ):
        if keys is None:
            keys = tuple(rollouts.keys())

        def collate_rollouts(instances: Sequence[tuple]):
            return {
                key: torch.stack([instance[idx] for instance in instances])
                for idx, key in enumerate(keys)
            }

        rollouts_dataset = TensorDataset(*[rollouts[key] for key in keys])
        rollouts_dataloader = DataLoader(
            dataset=rollouts_dataset,
            batch_size=self.args.step_per_device_batch_size,
            collate_fn=collate_rollouts,
            shuffle=shuffle,
            drop_last=drop_last,
        )
        # Do not prepare, since we don't need to shard the rollouts sampled on each batch.
        return rollouts_dataloader

    def get_policy_meta_prompts(self):
        if self.policy_meta_prompts is not None:
            return self.policy_meta_prompts

        policy_meta_prompt_pattern = self.args.data_config.policy_meta_prompt_pattern
        policy_meta_prompts = utils.make_meta_prompts(
            policy_meta_prompt_pattern,
        )
        self.policy_meta_prompts = policy_meta_prompts
        return policy_meta_prompts

    def get_reward_meta_prompts(self):
        if self.reward_meta_prompts is not None:
            return self.reward_meta_prompts

        reward_meta_prompt_pattern = self.args.data_config.reward_meta_prompt_pattern
        reward_meta_prompts = utils.make_meta_prompts(
            reward_meta_prompt_pattern,
        )
        self.reward_meta_prompts = reward_meta_prompts
        return reward_meta_prompts

    def get_principle_collection(self):
        if self.principle_collection is not None:
            return self.principle_collection

        principle_collection_path = self.args.data_config.principle_collection_path
        assert os.path.exists(principle_collection_path)
        print("Loading principle collection from", principle_collection_path)
        with open(principle_collection_path, "r", encoding="utf-8") as f:
            principle_collection = json.load(f)
        self.principle_collection = principle_collection
        return principle_collection

    def prepare_reward_inputs(self, inputs, outputs, question_type_flag="GENERAL"):
        reward_meta_prompt = self.get_reward_meta_prompts()[0]

        reasoning_keywords = ["reason", "numerical"]
        reasoning_bonus_weight = 128.0

        redteaming_keywords = ["ethical", "harmless"]
        redteaming_bonus_weight = 128.0

        if "{Dimensions}" in reward_meta_prompt:
            principle_collection = self.get_principle_collection()
            random.shuffle(principle_collection)
            dimension_str = []
            for item in principle_collection:
                dimension_str.append(f"- {item['definition']}")
            if self.args.data_config.max_principles is not None:
                if "weight" not in principle_collection[0]:
                    dimension_str = dimension_str[
                        : self.args.data_config.max_principles
                    ]
                else:
                    remaining_weights = []

                    for item in principle_collection:
                        assert "weight" in item
                        reasoning_principle_flag = False
                        for keyword in reasoning_keywords:
                            if keyword in item["dimension"].lower():
                                reasoning_principle_flag = True

                        redteaming_principle_flag = False
                        for keyword in redteaming_keywords:
                            if keyword in item["dimension"].lower():
                                redteaming_principle_flag = True

                        if (
                            question_type_flag == "REASONING"
                            and reasoning_principle_flag
                        ):
                            remaining_weights.append(
                                item["weight"] + reasoning_bonus_weight
                            )
                        elif (
                            question_type_flag == "REDTEAMING"
                            and redteaming_principle_flag
                        ):
                            remaining_weights.append(
                                item["weight"] + redteaming_bonus_weight
                            )
                        else:
                            remaining_weights.append(item["weight"])

                    remaining_idx = list(range(len(dimension_str)))

                    sampled_dimension_str = []
                    while (
                        len(sampled_dimension_str)
                        < self.args.data_config.max_principles
                    ):
                        sampled_idx = random.choices(
                            list(range(len(remaining_idx))), weights=remaining_weights
                        )[0]
                        sampled_dimension_str.append(
                            dimension_str[remaining_idx[sampled_idx]]
                        )
                        remaining_idx.pop(sampled_idx)
                        remaining_weights.pop(sampled_idx)

                    dimension_str = sampled_dimension_str

            dimension_str = "\n".join(dimension_str)
            reward_meta_prompt = reward_meta_prompt.replace(
                "{Dimensions}", dimension_str
            )

        return reward_meta_prompt.format(
            Input=inputs,
            Output=outputs,
        )


def cast_with_native_amp(
    func: Callable, mixed_precision: Optional[str] = None
) -> Callable:
    """Almost like how huggingface accelerate cast `model.forward`."""
    if mixed_precision not in ("fp16", "bf16"):
        logger.warning(
            f"Unknown mixed precision mode: {mixed_precision}, falling back to fp32."
        )
        return func

    if mixed_precision == "fp16":
        output_func = torch.cuda.amp.autocast(dtype=torch.float16)(func)
    else:
        device_type = "cuda" if torch.cuda.is_available() else "cpu"
        output_func = torch.autocast(device_type=device_type, dtype=torch.bfloat16)(
            func
        )
    output_func = convert_outputs_to_fp32(output_func)
    return output_func


def truncate_after_eos(completions, eos_token_id):
    # We truncate tokens after eos_token_id
    clean_completions = completions.tolist()
    for idx, completion in enumerate(clean_completions):
        try:
            end_idx = completion.index(eos_token_id)
            clean_completions[idx] = completion[:end_idx]
        except ValueError:
            pass
    return clean_completions

training/models/rl_trainer.py
IBM/SALMON
# Copyright 2023 The Self-Align Team
# Copyright 2023 The Alpaca Team
# Copyright 2022 The HuggingFace 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.

from typing import Optional

from torch import nn, optim
from transformers import Trainer
from transformers.optimization import get_scheduler
from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
from transformers.trainer_pt_utils import get_parameter_names


def create_optimizer(
    args, model: nn.Module, optimizer: Optional[optim.Optimizer] = None
):
    """Create optimizer for trainer.

    This is detached version of the `Trainer.create_optimizer` method.
    We don't support sagemaker and fairscale for simplicity.

    Reference:
        https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py
    """
    opt_model = model

    if optimizer is None:
        decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
        decay_parameters = [name for name in decay_parameters if "bias" not in name]
        optimizer_grouped_parameters = [
            {
                "params": [
                    p
                    for n, p in opt_model.named_parameters()
                    if (n in decay_parameters and p.requires_grad)
                ],
                "weight_decay": args.weight_decay,
            },
            {
                "params": [
                    p
                    for n, p in opt_model.named_parameters()
                    if (n not in decay_parameters and p.requires_grad)
                ],
                "weight_decay": 0.0,
            },
        ]

        optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(args)

        optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
    return optimizer


def create_scheduler(args, optimizer, lr_scheduler, num_training_steps):
    """Create scheduler for trainer.

    This is detached version of the `Trainer.create_scheduler` method.

    Reference:
        https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py
    """
    if lr_scheduler is None:
        lr_scheduler = get_scheduler(
            args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=args.get_warmup_steps(num_training_steps),
            num_training_steps=num_training_steps,
        )
    return lr_scheduler

training/models/trainer_utils.py
IBM/SALMON
import glob
import os
from os.path import exists, join, isdir
import shutil
import sys
from typing import Optional, Dict, Sequence, List

import torch
import transformers
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR

from models.reward_model import RewardModel

DEFAULT_PAD_TOKEN = "[PAD]"


class SavePeftModelCallback(transformers.TrainerCallback):
    def save_model(self, args, state, kwargs):
        print("Saving PEFT checkpoint...")

        global_rank = int(os.environ.get("RANK", 0))

        if global_rank == 0:
            print("Saving model checkpoint to %s" % args.output_dir)
            if state.best_model_checkpoint is not None:
                checkpoint_folder = state.best_model_checkpoint
            else:
                checkpoint_folder = os.path.join(
                    args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
                )

            peft_model_path = os.path.join(checkpoint_folder, "adapter_model")
            reward_head_path = os.path.join(checkpoint_folder, "reward_head")

            if isinstance(kwargs["model"], RewardModel):
                kwargs["model"].backbone_model.save_pretrained(peft_model_path)
                torch.save(
                    kwargs["model"].reward_head.state_dict(),
                    reward_head_path,
                )
            else:
                kwargs["model"].save_pretrained(peft_model_path)

            pytorch_model_paths = glob.glob(
                os.path.join(checkpoint_folder, "pytorch_model*.bin")
            )
            for pytorch_model_path in pytorch_model_paths:
                if os.path.exists(pytorch_model_path):
                    os.remove(pytorch_model_path)

            optimizer_path = os.path.join(checkpoint_folder, "optimizer.pt")
            if os.path.exists(optimizer_path):
                os.remove(optimizer_path)

        else:
            print("Skipping PEFT checkpoint save on rank %d" % global_rank)

    def on_save(self, args, state, control, **kwargs):
        self.save_model(args, state, kwargs)
        return control

    def on_train_end(self, args, state, control, **kwargs):
        def touch(fname, times=None):
            global_rank = int(os.environ.get("RANK", 0))
            if global_rank == 0:
                with open(fname, "a"):
                    os.utime(fname, times)

        touch(join(args.output_dir, "completed"))
        self.save_model(args, state, kwargs)


def print_trainable_parameters(args, model):
    """
    Prints the number of trainable parameters in the model.
    """
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    if args.bits == 4:
        trainable_params /= 2
    print(
        f"trainable params: {trainable_params} || "
        f"all params: {all_param} || "
        f"trainable: {100 * trainable_params / all_param}"
    )


def get_last_checkpoint(checkpoint_dir):
    if isdir(checkpoint_dir):
        is_completed = exists(join(checkpoint_dir, "completed"))
        if is_completed:
            return None, True  # already finished
        max_step = 0
        for filename in os.listdir(checkpoint_dir):
            if isdir(join(checkpoint_dir, filename)) and filename.startswith(
                "checkpoint"
            ):
                max_step = max(max_step, int(filename.replace("checkpoint-", "")))
        if max_step == 0:
            return None, is_completed  # training started, but no checkpoint
        checkpoint_dir = join(checkpoint_dir, f"checkpoint-{max_step}")
        print(f"Found a previous checkpoint at: {checkpoint_dir}")
        return checkpoint_dir, is_completed  # checkpoint found!
    return None, False  # first training

training/qlora_utils.py
IBM/SALMON
# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the GNU General Public License version 3.

from typing import Optional
import math
import os
import fire
import time
import tqdm
import json

from pathlib import Path

import torch

from llama_dromedary import Llama


def main(
    ckpt_dir: str,
    tokenizer_path: str,
    temperature: float = 1.0,
    top_p: float = 1.0,
    max_seq_len: int = 512,
    max_batch_size: int = 32,
    max_shared_seq_len: int = 512,
    generate_max_len: int = 128,
    group_rank: int = -1,
    group_size: int = -1,
    input_file: str = None,
    output_file: str = None,
    meta_prompt_file: str = None,
    prompt_style: str = "dromedary",
    seed: Optional[int] = None,
):
    assert group_rank >= 0, "Must specify group rank"
    assert group_size >= 0, "Must specify group size"
    assert (
        input_file is not None and output_file is not None
    ), "Must specify input and output files"
    assert meta_prompt_file is not None, "Must specify meta prompt file"

    with open(input_file, "r") as f:
        inputs = json.load(f)
    inputs = inputs[group_rank::group_size]

    if prompt_style == "dromedary":
        generate_prompt_fn = generate_dromedary_prompt
    else:
        raise ValueError(f"Unknown prompt style: {prompt_style}")

    with open(meta_prompt_file, "r") as f:
        meta_prompt = f.read().strip()

    generator = Llama.build(
        ckpt_dir=ckpt_dir,
        tokenizer_path=tokenizer_path,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
        max_shared_seq_len=max_shared_seq_len,
    )

    if seed is not None:
        torch.manual_seed(seed)

    results = []
    # record current progress

    if "shards" not in output_file and group_size > 1:
        output_file = output_file.replace(
            ".json", f"_{group_size}shards_{group_rank}.json"
        )

    if Path(output_file).exists():
        with open(output_file, "r") as f:
            results = f.readlines()
            results = [line for line in results if len(line.strip()) > 0]

    inputs = inputs[len(results) :]
    print("Skipping %d examples" % len(results))

    global_rank = int(os.environ.get("RANK", "0"))
    batching_inputs = tqdm.tqdm(
        BatchIterator(inputs, max_batch_size),
        desc="Batched inference",
        disable=global_rank > 0,
    )
    total_iters = len(inputs) // max_batch_size

    output_handler = None
    if global_rank == 0:
        output_handler = open(output_file, "a")

    # prepare inputs with batch size $max_batch_size
    for iter, batched_inputs in enumerate(batching_inputs):
        t0 = time.time()
        prompts = [
            generate_prompt_fn(
                meta_prompt, ex_input["instruction"].strip(), ex_input["input"].strip()
            )
            for ex_input in batched_inputs
        ]

        if prompt_style == "llama_2_chat":
            outputs = generator.chat_completion(
                prompts,  # type: ignore
                max_gen_len=generate_max_len,
                temperature=temperature,
                top_p=top_p,
            )
        else:
            outputs = generator.text_completion(
                prompts,
                max_gen_len=generate_max_len,
                temperature=temperature,
                top_p=top_p,
            )

        t1 = time.time()

        results = []
        for ex_input, output in zip(batched_inputs, outputs):
            results.append(
                {
                    "instruction": ex_input["instruction"],
                    "input": ex_input["input"],
                    "output": output,
                }
            )

        if group_rank == 0:
            for ex_input, output, _ in zip(batched_inputs, outputs, range(8)):
                print("=" * 20, "iter: ", iter, "/", total_iters, "latency: ", t1 - t0)
                print(f"Input: {ex_input['instruction']}: {ex_input['input']}")
                print(f"Output: {output}")
                print()

        if output_handler is not None:
            for result in results:
                output_handler.write(json.dumps(result) + "\n")
            output_handler.flush()

    if output_handler is not None:
        output_handler.close()


class BatchIterator:
    def __init__(self, data, batch_size=1):
        self.data = data
        self.batch_size = batch_size
        self.index = 0

    def __iter__(self):
        for i in range(0, len(self.data), self.batch_size):
            yield self.data[i : i + self.batch_size]

    def __len__(self):
        return math.ceil(len(self.data) / self.batch_size)


def generate_dromedary_prompt(meta_prompt, instruction, input=None):
    if input:
        return f"""{meta_prompt}
{instruction}

{input}

### Dromedary"""
    else:
        return f"""{meta_prompt}
{instruction}

### Dromedary"""


if __name__ == "__main__":
    fire.Fire(main)

training/step1_synthetic_preference_collection/batch_generation.py
IBM/SALMON
import random
import tqdm
import re
import json
from datasets import load_dataset
from transformers import AutoTokenizer
import fire


def load_oasst_data():
    oasst_dataset = load_dataset("OpenAssistant/oasst1")["train"]

    def create_message_trees(dataset):
        """Create message trees from dataset."""
        # Organize data into dictionary based on parent_id
        organized_data = {}
        for message in dataset:
            parent_id = message["parent_id"]
            if parent_id not in organized_data:
                organized_data[parent_id] = []
            organized_data[parent_id].append(message)

        # Traverse data to create trees
        message_trees = []
        for root_messages in organized_data[None]:
            tree = []
            current_message = root_messages
            while current_message is not None:
                tree.append(current_message)
                children = organized_data.get(current_message["message_id"])
                current_message = children[0] if children else None
            message_trees.append(tree)

        return message_trees

    oasst_message_trees = create_message_trees(oasst_dataset)
    oasst_examples = []

    count = 0
    for oasst_tree in oasst_message_trees:
        if len(oasst_tree) >= 2:
            count += 1
            oasst_examples.append(
                {
                    "instruction": oasst_tree[0]["text"],
                    "input": "",
                    "output": "",
                }
            )
    print("OASST examples:", count)
    return oasst_examples


def main(
    output_file: str = "/path/to/oasst1_prompts.json",
):
    oasst_examples = load_oasst_data()

    with open(output_file, "w") as f:
        json.dump(oasst_examples, f, indent=2)


if __name__ == "__main__":
    fire.Fire(main)

training/step1_synthetic_preference_collection/clean_oasst1_prompts.py
IBM/SALMON
import fcntl
import glob
import json
import os
import random

import numpy as np
import torch
import tqdm
import fire

from peft import PeftModel
from transformers import (
    BitsAndBytesConfig,
    AutoModelForCausalLM,
    LlamaTokenizerFast,
)
from datasets import load_dataset


def main(
    model_name: str,
    adapters_name: str,
    preferece_prompt: str,
    rm_principles: str,
    response_pattern: str,
    output_file: str,
    tokenizer_name: str = "TheBloke/dromedary-65b-lora-HF",  # a random llama-based model
):
    with open(preferece_prompt, "r") as f:
        PREFERENCE_META_PROMPT = f.read().strip()

    with open(rm_principles, "r") as f:
        principle_definitions = json.load(f)

    data_files = glob.glob(response_pattern)
    cache = ""
    raw_data = {}
    for data_file in data_files:
        type = data_file.split("/")[-1].split("_")[-1].split(".")[0]
        raw_data[type] = []
        with open(data_file, "r") as f:
            for line in f:
                try:
                    raw_data[type].append(json.loads(cache + line))
                    cache = ""
                except json.decoder.JSONDecodeError:
                    cache += line.strip()
        raw_data[type] = sorted(raw_data[type], key=lambda x: x["instruction"])

    print(raw_data.keys())

    preference_dataset = []

    for type_idx_a in range(len(raw_data.keys())):
        for type_idx_b in range(type_idx_a + 1, len(raw_data.keys())):
            type_a = list(raw_data.keys())[type_idx_a]
            type_b = list(raw_data.keys())[type_idx_b]
            print(
                f"len {type_a} and {type_b}",
                len(raw_data[type_a]),
                len(raw_data[type_b]),
            )

            for idx in range(len(raw_data[type_a])):
                data_a = raw_data[type_a][idx]
                data_b = raw_data[type_b][idx]

                if data_a["instruction"] != data_b["instruction"]:
                    print("Instruction mismatch!")
                    print(data_a["instruction"])
                    print(data_b["instruction"])
                    exit(0)

                if (
                    len(data_a["output"]["generation"]) < 16
                    or len(data_b["output"]["generation"]) < 16
                ):
                    continue

                preference_dataset.append(
                    {
                        "instruction": data_a["instruction"],
                        "input": data_a["input"],
                        "output_1": data_a["output"]["generation"]
                        .split("###")[0]
                        .strip(),
                        "output_2": data_b["output"]["generation"]
                        .split("###")[0]
                        .strip(),
                        "preference": 0,
                    }
                )

    random.Random(42).shuffle(preference_dataset)

    print(f"Starting to load the model {model_name} into memory")

    rank = int(os.environ.get("RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    torch.cuda.set_device(rank)

    print(preference_dataset[0 + rank])

    batch_size = 4  # For a single 80GB GPU
    m = AutoModelForCausalLM.from_pretrained(
        model_name,
        load_in_4bit=True,
        torch_dtype=torch.float16,
        device_map={"": torch.cuda.current_device()},
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            load_in_8bit=False,
            llm_int8_threshold=6.0,
            llm_int8_has_fp16_weight=False,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
        ),
    )
    m = PeftModel.from_pretrained(
        m,
        adapters_name,
        adapter_name="default",
        is_trainable=False,
    )
    tok = LlamaTokenizerFast.from_pretrained(
        tokenizer_name,
        padding_side="left",
        truncation_side="left",
        model_max_length=1536,
    )
    tok.pad_token_id = 0

    print(f"Successfully loaded the model {model_name} into memory")
    print_flag = True

    idxs = []
    dimensions = []
    prompts = []
    preferences = []

    for idx in tqdm.tqdm(range(len(preference_dataset))):
        if idx % world_size != rank:
            continue

        for principle_definition in principle_definitions:
            dimension = principle_definition["dimension"]
            definition = principle_definition["definition"]

            data = preference_dataset[idx]
            instruction = data["instruction"]
            instruction_input = data["input"]

            if instruction_input:
                instruction += "\n\n" + instruction_input

            output_1 = data["output_1"]
            output_2 = data["output_2"]
            preference = data["preference"]
            preferences.append(preference)
            idxs.append(idx)
            dimensions.append(dimension)

            for a, b in [(output_1, output_2), (output_2, output_1)]:
                prompt_for_score = PREFERENCE_META_PROMPT.format(
                    UserInstruction=instruction,
                    OutputA=a,
                    OutputB=b,
                    Dimension=dimension,
                    Definition=definition,
                )

                if print_flag:
                    print_flag = False
                    print(prompt_for_score)

                prompts.append(tok.bos_token + prompt_for_score)

            if len(prompts) == batch_size * 2:
                tokenized_input = tok(
                    prompts,
                    return_tensors="pt",
                    add_special_tokens=False,
                    padding="max_length",
                    return_attention_mask=True,
                    truncation=True,
                )

                input_ids = tokenized_input["input_ids"].to(m.device)
                attention_mask = tokenized_input["attention_mask"].to(m.device)

                with torch.inference_mode():
                    output = m(
                        input_ids,
                        attention_mask=attention_mask,
                    )

                logits = output.logits

                token_id_a = tok.encode("\n (a", add_special_tokens=False)[-1]
                token_id_b = tok.encode("\n (b", add_special_tokens=False)[-1]

                relative_scores = []
                prompt_preferences = []

                for ex_idx in range(batch_size):
                    score_a_for_1_2 = logits[ex_idx * 2 + 0, -1, token_id_a]
                    score_b_for_1_2 = logits[ex_idx * 2 + 0, -1, token_id_b]
                    score_a_for_2_1 = logits[ex_idx * 2 + 1, -1, token_id_a]
                    score_b_for_2_1 = logits[ex_idx * 2 + 1, -1, token_id_b]

                    relative_score_1_2 = (score_a_for_1_2 - score_b_for_1_2).item()
                    relative_score_2_1 = (score_b_for_2_1 - score_a_for_2_1).item()

                    if relative_score_1_2 > 0.0 and relative_score_2_1 > 0.0:
                        prompt_preference = 1
                    elif relative_score_1_2 < 0.0 and relative_score_2_1 < 0.0:
                        prompt_preference = 2
                    else:
                        prompt_preference = 0

                    relative_scores.append((relative_score_1_2, relative_score_2_1))
                    prompt_preferences.append(prompt_preference)

                outputs = []

                for ex_idx in range(batch_size):
                    outputs.append(
                        {
                            "example_idx": idxs[ex_idx],
                            "dimension": dimensions[ex_idx],
                            "preference": preferences[ex_idx],
                            "prompt_preference": prompt_preferences[ex_idx],
                            "relative_score": relative_scores[ex_idx],
                        }
                    )

                with open(output_file, "a") as f:
                    fcntl.flock(f, fcntl.LOCK_EX)
                    for output in outputs:
                        f.write(json.dumps(output) + "\n")
                    fcntl.flock(f, fcntl.LOCK_UN)

                idxs = []
                dimensions = []
                prompts = []
                preferences = []


if __name__ == "__main__":
    fire.Fire(main)

training/step1_synthetic_preference_collection/synthetic_preference.py
IBM/SALMON
import fcntl
import glob
import json
import os
import random

import numpy as np
import tqdm
from ftlangdetect import detect


ALL_DIMENSIONS = [
    # "Overall Helpful",
    "Concise",
    "Honest and Accurate",
    "Ethical",
    # "Natural and Fluent",
    "Specific",
    "Educational and Engaging",
    "Methodical",
    # "Multilingual",
    "Creative",
    "Comprehensive",
]


def main(
    response_pattern: str,
    preference_file: str,
    output_file: str,
    noisy_ratio: float = 0.1,
    max_principles: int = 8,
    tokenizer_name: str = "TheBloke/dromedary-65b-lora-HF",  # a random llama-based model
):
    data_files = glob.glob(response_pattern)

    cache = ""
    raw_data = {}
    for data_file in data_files:
        type = data_file.split("/")[-1].split("_")[-1].split(".")[0]
        raw_data[type] = []
        with open(data_file, "r") as f:
            for line in f:
                try:
                    raw_data[type].append(json.loads(cache + line))
                    cache = ""
                except json.decoder.JSONDecodeError:
                    cache += line.strip()
        raw_data[type] = sorted(raw_data[type], key=lambda x: x["instruction"])

    print(raw_data.keys(), len(raw_data[type]))
    preference_dataset = []

    for type_idx_a in range(len(raw_data.keys())):
        for type_idx_b in range(type_idx_a + 1, len(raw_data.keys())):
            type_a = list(raw_data.keys())[type_idx_a]
            type_b = list(raw_data.keys())[type_idx_b]

            for idx in range(len(raw_data[type_a])):
                data_a = raw_data[type_a][idx]
                data_b = raw_data[type_b][idx]

                if data_a["instruction"] != data_b["instruction"]:
                    print("Instruction mismatch!")
                    print(data_a["instruction"])
                    print(data_b["instruction"])
                    exit(0)

                if (
                    len(data_a["output"]["generation"]) < 16
                    or len(data_b["output"]["generation"]) < 16
                ):
                    continue

                preference_dataset.append(
                    {
                        "instruction": data_a["instruction"],
                        "input": data_a["input"],
                        "output_1": data_a["output"]["generation"]
                        .split("###")[0]
                        .strip(),
                        "output_2": data_b["output"]["generation"]
                        .split("###")[0]
                        .strip(),
                        "preference": 0,
                    }
                )

    random.Random(42).shuffle(preference_dataset)

    for example_idx in range(len(preference_dataset)):
        preference_dataset[example_idx]["example_idx"] = example_idx

    with open(preference_file, "r") as f:
        fcntl.flock(f, fcntl.LOCK_EX)
        for line in f:
            data = json.loads(line)
            example_idx = data["example_idx"]

            if "dimension_scores" not in preference_dataset[example_idx]:
                preference_dataset[example_idx]["dimension_scores"] = {}

            dimension = data["dimension"]
            score = data["relative_score"]

            preference_dataset[example_idx]["dimension_scores"][dimension] = score
        fcntl.flock(f, fcntl.LOCK_UN)

    all_dimensions = ALL_DIMENSIONS
    print(preference_dataset[0])

    # print variance
    variances = {}
    for dimension in all_dimensions:
        scores = []
        for example_idx in range(len(preference_dataset)):
            if "dimension_scores" not in preference_dataset[example_idx]:
                continue
            if dimension not in preference_dataset[example_idx]["dimension_scores"]:
                continue
            two_scores = preference_dataset[example_idx]["dimension_scores"][dimension]
            scores.append(two_scores[0])
            scores.append(two_scores[1])
        variances[dimension] = np.var(scores)

        print("Variance of {}: {}".format(dimension, np.var(scores)))

    clean_preference_dataset = []

    for example_idx in range(len(preference_dataset)):
        if "dimension_scores" not in preference_dataset[example_idx]:
            continue

        continue_flag = False

        for dimension in all_dimensions:
            if dimension not in preference_dataset[example_idx]["dimension_scores"]:
                continue_flag = True

        if continue_flag:
            continue

        clean_preference_dataset.append(preference_dataset[example_idx])

    print(len(clean_preference_dataset))

    synthetic_data = []

    main_reasons = {}
    for dimension in all_dimensions:
        main_reasons[dimension] = 0

    random.seed(42)

    def mean(a, b):
        return (a + b) / 2.0

    rule_count = 0

    for i in range(4):
        for example_idx in range(len(clean_preference_dataset)):
            random.shuffle(all_dimensions)

            # Compute the probabilities
            num_dimensions = max_principles + 1
            while num_dimensions > max_principles:
                num_dimensions = np.random.geometric(0.2, 1)[0]

            data = clean_preference_dataset[example_idx]

            dimensions = all_dimensions[:num_dimensions]

            scores = []
            for dimension in dimensions:
                two_scores = data["dimension_scores"][dimension]
                if two_scores[0] > 0.0 and two_scores[1] > 0.0:
                    scores.append(
                        mean(two_scores[0], two_scores[1]) / variances[dimension]
                    )
                elif two_scores[0] < 0.0 and two_scores[1] < 0.0:
                    scores.append(
                        mean(two_scores[0], two_scores[1]) / variances[dimension]
                    )
                else:
                    scores.append(0.0)

            # negative_definitions is random benoulli of length num_dimensions
            negative_definitions = np.random.binomial(1, 0.5, num_dimensions)

            scores = [
                -score if negative_definitions[idx] == 1 else score
                for idx, score in enumerate(scores)
            ]

            # The preference is based on the absolute value of the scores.
            doinate_dimension_idx = np.argmax(np.abs(scores)).item()

            score = scores[doinate_dimension_idx]

            if score > 0.0:
                preference = 1
            elif score < 0.0:
                preference = 2
            else:
                continue

            order = random.randint(0, 1)

            rule_flag = False

            # We ban "as an AI language model"
            if "As an AI language model" in data["output_1"]:
                if "As an AI language model" in data["output_2"]:
                    pass
                else:
                    if random.randint(0, 1) > 0.5:
                        preference = 2
                        rule_flag = True
            else:
                if "As an AI language model" in data["output_2"]:
                    if random.randint(0, 1) > 0.5:
                        preference = 1
                        rule_flag = True
                else:
                    pass

            lang_instruction = detect(data["instruction"].replace("\n", " "))["lang"]
            lang_output_1 = detect(data["output_1"].replace("\n", " "))["lang"]
            lang_output_2 = detect(data["output_2"].replace("\n", " "))["lang"]

            # We ban different languages
            if lang_instruction != lang_output_1:
                if lang_instruction != lang_output_2:
                    continue
                else:
                    preference = 2
                    rule_flag = True
            else:
                if lang_instruction != lang_output_2:
                    preference = 1
                    rule_flag = True
                else:
                    pass

            data_tuple = (data["instruction"], data["output_1"], data["output_1"])

            if rule_flag:
                rule_count += 1
                flip = random.random() < noisy_ratio * 2.5  # 2.5 is a magic number
            else:
                flip = random.random() < noisy_ratio
                main_reasons[dimensions[doinate_dimension_idx]] += 1

            preference = 3 - preference if flip else preference

            synthetic_data.append(
                {
                    "doinate_dimension_idx": doinate_dimension_idx,
                    "dimensions": str(tuple(dimensions)),
                    "dimension_scores": str(
                        tuple(scores if order == 0 else [-score for score in scores])
                    ),
                    "negative_definitions": str(tuple(negative_definitions)),
                    "instruction": data["instruction"],
                    "input": data["input"],
                    "output_1": data["output_1"] if order == 0 else data["output_2"],
                    "output_2": data["output_2"] if order == 0 else data["output_1"],
                    "preference": preference if order == 0 else 3 - preference,
                    "flip": flip,
                    "lang": (lang_instruction, lang_output_1, lang_output_2)
                    if order == 0
                    else (lang_instruction, lang_output_2, lang_output_1),
                    "rule_flag": rule_flag,
                }
            )

    print(rule_count)
    print(len(synthetic_data))
    print(main_reasons)

    with open(output_file, "w") as f:
        f.write(
            json.dumps(
                synthetic_data,
                indent=4,
            )
        )

training/step2_rm_training/aggregate_synthetic_preference.py
IBM/SALMON
import json
import random
import tqdm
from datasets import load_dataset
import fire


def find_shared_prefix(str1, str2):
    """
    Find the shared prefix between two strings.

    Args:
        str1 (str): The first string.
        str2 (str): The second string.

    Returns:
        int: The index position where the shared prefix ends in the first string.
    """
    # Find the minimum length of both strings
    min_length = min(len(str1), len(str2))

    max_common_prefix = 0
    # Loop through characters to find where they start to differ
    for i in range(min_length):
        if str1[i] != str2[i]:
            max_common_prefix = i
            break

    while (
        not str1[:max_common_prefix].endswith("Assistant:")
        # Force the prefix to end with "Assistant:"
    ) and max_common_prefix > 0:
        max_common_prefix -= 1

    return max_common_prefix


def split_example(example):
    """
    Split an example into a shared prefix and the remaining parts of the strings.

    Args:
        example (dict): A dictionary containing 'chosen' and 'rejected' keys with strings as values.

    Returns:
        dict: A dictionary with keys 'query', 'output_1', and 'output_2'.
    """
    chosen = example["chosen"]
    rejected = example["rejected"]

    # Find the index where the shared prefix ends
    shared_index = find_shared_prefix(chosen, rejected)

    # Split the strings
    query = chosen[:shared_index].strip()
    output_1 = chosen[shared_index:].strip()
    output_2 = rejected[shared_index:].strip()

    # Return the result as a dictionary
    return {"query": query, "output_1": output_1, "output_2": output_2}


def main(
    output_file: str,
):
    shp_dataset = load_dataset("stanfordnlp/SHP")["train"]

    merged_data = []

    for ex in tqdm.tqdm(shp_dataset):
        qid = ex["post_id"]
        score = ex["score_ratio"]

        ex_dp = {
            "score": score,
            "instruction": ex["history"],
            "input": "",
            "output_1": ex["human_ref_A"],
            "output_2": ex["human_ref_B"],
            "preference": 1 if ex["labels"] == 1 else 2,
        }

        if score > 2.0:
            merged_data.append(ex_dp)

    print(merged_data[-1])
    print(len(merged_data))

    hh_dataset = load_dataset("Anthropic/hh-rlhf")["train"]

    for ex in hh_dataset:
        processed_ex = split_example(ex)

        random_swap = random.random() < 0.5

        merged_data.append(
            {
                "instruction": processed_ex["query"].strip(),
                "input": "",
                "output_1": processed_ex["output_1"]
                if random_swap
                else processed_ex["output_2"],
                "output_2": processed_ex["output_2"]
                if random_swap
                else processed_ex["output_1"],
                "preference": 1 if random_swap else 2,
            }
        )

    print(merged_data[-1])
    print(len(merged_data))

    random.shuffle(merged_data)

    with open(output_file, "w") as f:
        f.write(
            json.dumps(
                merged_data,
            )
        )


if __name__ == "__main__":
    fire.Fire(main)

training/step2_rm_training/clean_pmp_data.py
IBM/SALMON
import json
import fire
from datasets import load_dataset

DEFAULT_SHAREGPT_DATA = [
    # https://huggingface.co/datasets/zetavg/ShareGPT-Processed
    "zetavg/ShareGPT-Processed",
    # https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/blob/main/HTML_cleaned_raw_dataset/sg_90k_part1.json
    "../../outputs/data/sharegpt/sg_90k_part1.json",
    # https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/blob/main/HTML_cleaned_raw_dataset/sg_90k_part2.json
    "../../outputs/data/sharegpt/sg_90k_part2.json",
    # "../../outputs/data/sharegpt/ShareGPT.jsonl",
    # "../../outputs/data/sharegpt/ShareGPT_V3_unfiltered_cleaned_split.json",
]


def extract_from_dataset(dataset, key="markdown", prefix="human"):
    data = []

    for item in dataset:
        conversations = item["conversations"]
        if len(conversations) >= 2:
            user_input = conversations[0]

            if user_input["from"] == prefix:
                data.append(
                    {
                        "instruction": user_input[key],
                        "input": "",
                        "output": "",
                    }
                )

    return data


def extract_from_json_file(filepath, key="value", prefix="human"):
    with open(filepath, "r") as f:
        dataset = json.load(f)

    return extract_from_dataset(dataset, key, prefix)


def extract_from_jsonl_file(filepath):
    data = []

    with open(filepath, "r") as f:
        for line in f:
            line_data = json.loads(line)
            user_input = line_data["text"].split("<bot>")[0][len("<human>:") :].strip()

            data.append(
                {
                    "instruction": user_input,
                    "input": "",
                    "output": "",
                }
            )

    return data


def main(data_files: list = None, output_file: str = "path/to/sharegpt_prompts.json"):
    if data_files is None:
        data_files = DEFAULT_SHAREGPT_DATA

    all_data = []
    merged_data = []
    merged_data_set = set()

    for filepath in data_files:
        if filepath.endswith(".json"):
            all_data.append(extract_from_json_file(filepath))
        elif filepath.endswith(".jsonl"):
            all_data.append(extract_from_jsonl_file(filepath))
        else:
            train_dataset = load_dataset(filepath)["train"]
            all_data.append(extract_from_dataset(train_dataset))

    # Merge and filter unique instructions
    for data in all_data:
        filtered_data_count = 0
        for item in data:
            if item["instruction"] not in merged_data_set:
                merged_data_set.add(item["instruction"])
                merged_data.append(item)
            else:
                filtered_data_count += 1

        print("Filtered data count:", filtered_data_count)
        print("Merged data size:", len(merged_data))

    with open(output_file, "w") as f:
        json.dump(merged_data, f, indent=2)


if __name__ == "__main__":
    fire.Fire(main)

training/step3_ppo_training/aggregate_sharegpt_prompts.py
IBM/SALMON
import random
import tqdm
import re
import json
from datasets import load_dataset
from transformers import AutoTokenizer
import fire

# QUESTION TYPE
# 0: GENERAL (ShareGPT + Dolly + OASST)
# 1: REASONING (OpenOrca + MATH)
# 2: REDTEAMING (Anthropic Red Teaming)

CHATGPT_LANGUAGES = {
    "it": "Italian",
    "pl": "Polish",
    "ru": "Russian",
    "sk": "Slovak",
    "pt": "Portuguese",
    "ro": "Romanian",
    "da": "Danish",
    "sv": "Swedish",
    "no": "Norwegian",
    "en": "English",
    "es": "Spanish",
    "fr": "French",
    "cs": "Czech",
    "de": "German",
    "fi": "Finnish",
    "et": "Estonian",
    "lv": "Latvian",
    "lt": "Lithuanian",
    "fa": "Persian",
    "hu": "Hungarian",
    "he": "Hebrew",
    "el": "Greek",
    "ar": "Arabic",
    "kr": "Korean",
    "ja": "Japanese",
    "zh": "Chinese",
    "zh-traditional": "Chinese (Traditional)",
    "zh-simplified": "Chinese (Simplified)",
    "th": "Thai",
    "vi": "Vietnamese",
}


def remove_leading_fraction(input_string):
    # remove leading fraction
    cleaned_string = re.sub(r"^\d+\s*/\s*\d+", "", input_string)
    cleaned_string = re.sub(r"\d+\s*/\s*\d+$", "", cleaned_string)
    cleaned_string = cleaned_string.split("1 / 1", 1)[-1]

    # \uc9c0\uae08 \ubc88\uc5ed\ud558\uae30
    cleaned_string = cleaned_string.split("지금 번역하기")[0]

    # Language: English
    cleaned_string = cleaned_string.split("\n \n Language: ")[0]

    cleaned_string = cleaned_string.strip()

    if cleaned_string.endswith("Share Prompt"):
        cleaned_string = cleaned_string[: -len("Share Prompt")].strip()

    if cleaned_string.endswith("Translate now"):
        cleaned_string = cleaned_string[: -len("Translate now")].strip()

    for lang_code in CHATGPT_LANGUAGES:
        lang_suffix = f"Language: {CHATGPT_LANGUAGES[lang_code]}"
        if cleaned_string.endswith(lang_suffix):
            cleaned_string = cleaned_string[: -len(lang_suffix)].strip()

    # ~The following is a conversation with Bing, not ChatGPT.~
    if cleaned_string.startswith(
        "~The following is a conversation with Bing, not ChatGPT.~"
    ):
        cleaned_string = cleaned_string[
            len("~The following is a conversation with Bing, not ChatGPT.~") :
        ].strip()

    return cleaned_string


def load_dolly_data(length_bonus):
    dataset = load_dataset("databricks/databricks-dolly-15k")["train"]
    category_to_examples = {}
    for example in dataset:
        category = example["category"]
        if category not in category_to_examples:
            category_to_examples[category] = []
        category_to_examples[category].append(example)

    merged_examples = []
    for data in [
        category_to_examples["creative_writing"],
        category_to_examples["brainstorming"],
        category_to_examples["open_qa"],
        category_to_examples["general_qa"],
        category_to_examples["classification"],
    ]:
        for i in range(len(data)):
            assert data[i]["context"] == ""
            merged_examples.append(
                {
                    "instruction": data[i]["instruction"],
                    "input": "",
                    "output": "",
                    "length_bonus": length_bonus,
                    "question_type": 0,
                }
            )
    print("Dolly examples:", len(merged_examples))
    return merged_examples


def load_oasst_data(length_bonus):
    oasst_dataset = load_dataset("OpenAssistant/oasst1")["train"]

    def create_message_trees(dataset):
        """Create message trees from dataset."""
        # Organize data into dictionary based on parent_id
        organized_data = {}
        for message in dataset:
            parent_id = message["parent_id"]
            if parent_id not in organized_data:
                organized_data[parent_id] = []
            organized_data[parent_id].append(message)

        # Traverse data to create trees
        message_trees = []
        for root_messages in organized_data[None]:
            tree = []
            current_message = root_messages
            while current_message is not None:
                tree.append(current_message)
                children = organized_data.get(current_message["message_id"])
                current_message = children[0] if children else None
            message_trees.append(tree)

        return message_trees

    oasst_message_trees = create_message_trees(oasst_dataset)
    oasst_examples = []

    count = 0
    for oasst_tree in oasst_message_trees:
        if len(oasst_tree) >= 2:
            count += 1
            oasst_examples.append(
                {
                    "instruction": oasst_tree[0]["text"],
                    "input": "",
                    "output": "",
                    "length_bonus": length_bonus,
                    "question_type": 0,
                }
            )
    print("OASST examples:", count)
    return oasst_examples


def load_math_data(length_bonus):
    dataset = load_dataset("competition_math")
    train_data = dataset["train"]
    max_samples_per_dataset = 10000

    math_data = []

    for i in range(min(len(train_data), max_samples_per_dataset)):
        ex_instruction = train_data[i]["problem"].strip()
        math_data.append(
            {
                "instruction": ex_instruction,
                "input": "",
                "output": "",
                "length_bonus": length_bonus,
                "question_type": 1,
            }
        )
    return math_data


def filter_and_clean_examples(merged_examples):
    tokenizer = AutoTokenizer.from_pretrained("TheBloke/dromedary-65b-lora-HF")
    max_seq_length = 256
    filtered_examples = []
    set_of_unique_instructions = set()

    # Filter out examples with non-ascci characters
    merged_examples = [
        {
            "instruction": remove_leading_fraction(example["instruction"]),
            "input": "",
            "output": "",
            "length_bonus": example["length_bonus"],
            "question_type": example["question_type"],
        }
        for example in merged_examples
    ]

    for example in tqdm.tqdm(merged_examples):
        instruction = example["instruction"]
        instruction_token_length = len(tokenizer.encode(instruction))
        if (
            2 <= instruction_token_length <= max_seq_length
            and instruction not in set_of_unique_instructions
        ):
            filtered_examples.append(example)
            set_of_unique_instructions.add(instruction)

    return filtered_examples


def load_json(share_gpt_data_path):
    with open(share_gpt_data_path, "r") as f:
        share_gpt_data = json.load(f)
    examples = []
    for data in share_gpt_data:
        examples.append(
            {
                "instruction": data["instruction"],
                "input": "",
                "output": "",
            }
        )
    print(f"{share_gpt_data_path} examples:", len(examples))
    return examples


def main(
    sharegpt_prompt_path: str = "/path/to/sharegpt_prompts.json",
    openorca_prompt_path: str = "/path/to/openorca_prompts.json",
    output_file: str = "/path/to/salmon_merged_prompts.json",
    general_length_bonus: float = 1.0,
    reasoning_length_bonus: float = -0.4,
    redteaming_length_bonus: float = 0.0,
):
    del redteaming_length_bonus  # unused

    dolly_examples = load_dolly_data(length_bonus=general_length_bonus)
    oasst_examples = load_oasst_data(length_bonus=general_length_bonus)

    sharegpt_examples = load_json(sharegpt_prompt_path)
    sharegpt_examples = [
        {
            "instruction": example["instruction"],
            "input": "",
            "output": "",
            "length_bonus": general_length_bonus,
            "question_type": 0,
        }
        for example in sharegpt_examples
    ]

    openorca_examples = load_json(openorca_prompt_path)
    openorca_examples = [
        {
            "instruction": example["instruction"],
            "input": "",
            "output": "",
            "length_bonus": reasoning_length_bonus,
            "question_type": 1,
        }
        for example in openorca_examples
    ]
    math_examples = load_math_data(length_bonus=reasoning_length_bonus)

    merged_examples = (
        dolly_examples
        + oasst_examples
        + sharegpt_examples
        + openorca_examples
        + math_examples
    )
    filtered_examples = filter_and_clean_examples(merged_examples)

    print("Total examples:", len(filtered_examples))

    with open(output_file, "w") as f:
        json.dump(filtered_examples, f, indent=2)


if __name__ == "__main__":
    fire.Fire(main)

training/step3_ppo_training/clean_and_merge_prompts.py
IBM/SALMON
import tqdm
import json
import random
import pandas as pd
import fire
from transformers import LlamaTokenizerFast

# https://huggingface.co/datasets/Open-Orca/OpenOrca/blob/main/1M-GPT4-Augmented.parquet


def main(
    train_data_path: str = "path/to/data/1M-GPT4-Augmented.parquet",
    output_path: str = "path/to/data/openorca_prompts.json",
    tokenizer_name: str = "TheBloke/dromedary-65b-lora-HF",  # a random llama-based model
    max_samples_per_dataset: int = 10000,
    max_prompt_len: int = 256,
):
    train_data = pd.read_parquet(train_data_path)
    tokenizer = LlamaTokenizerFast.from_pretrained(tokenizer_name)
    niv_data = []
    flan_data = []
    t0_data = []
    cot_data = []

    for i in tqdm.tqdm(range(len(train_data))):
        datapoint = train_data.iloc[i]

        d_id = datapoint["id"]

        if "niv" in d_id:
            niv_data.append(datapoint)
            continue
        if "flan" in d_id:
            flan_data.append(datapoint)
            continue
        if "t0" in d_id:
            t0_data.append(datapoint)
            continue
        if "cot" in d_id:
            cot_data.append(datapoint)
            continue

        raise ValueError("Unknown dataset")

    print("Stats:")
    print(f"niv: {len(niv_data)}")
    print(f"flan: {len(flan_data)}")
    print(f"t0: {len(t0_data)}")
    print(f"cot: {len(cot_data)}")

    total_data = []

    set_of_instructions = set()

    for data in tqdm.tqdm([niv_data, flan_data, t0_data, cot_data]):
        data_count = 0
        random.shuffle(data)
        for datapoint in tqdm.tqdm(data):
            ex_instruction = datapoint["question"].strip()
            ex_input = ""
            ex_output = ""

            if ex_instruction in set_of_instructions:
                continue
            else:
                set_of_instructions.add(ex_instruction)

            if (
                len(tokenizer.encode(ex_instruction + "\n\n" + ex_input))
                > max_prompt_len
            ):
                continue

            # We delete NLI tasks
            if "it is not possible to tell" in datapoint["question"]:
                continue

            continue_flag = False
            for keyword in ["Premise", "premise", "Hypothesis", "hypothesis"]:
                if keyword in datapoint["question"]:
                    continue_flag = True
                    break
            if continue_flag:
                continue

            total_data.append(
                {
                    "instruction": ex_instruction,
                    "input": ex_input,
                    "output": ex_output,
                }
            )

            data_count += 1
            if data_count >= max_samples_per_dataset:
                break
        print("data_count:", data_count, len(total_data))

    random.shuffle(total_data)

    print(f"Total data: {len(total_data)}")

    with open(output_path, "w") as f:
        json.dump(total_data, f, indent=2)


if __name__ == "__main__":
    fire.Fire(main)

training/step3_ppo_training/subsample_openorca_prompts.py
IBM/SALMON
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import argparse
from dataclasses import dataclass, field
import os
import sys
from typing import Optional, List
import logging

import torch

from accelerate import DistributedDataParallelKwargs

import transformers

try:
    from transformers import LlamaTokenizerFast as LlamaTokenizer

    print("Using fast tokenizer")
except:
    from transformers import LlamaTokenizer

    print("Using slow tokenizer")

from transformers import AutoTokenizer

from data_utils.data_utils_ppo import make_rl_data_module
from qlora_utils import (
    get_last_checkpoint,
    DEFAULT_PAD_TOKEN,
)
from models.ppo_trainer import (
    PPOTrainer,
    make_models,
)
from models.rl_trainer import AlpacaAccelerator


torch.backends.cuda.matmul.allow_tf32 = True

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:
    reward_model_name_or_path: Optional[str] = field(default="EleutherAI/pythia-12b")
    policy_model_name_or_path: Optional[str] = field(default="EleutherAI/pythia-1.4b")
    trust_remote_code: Optional[bool] = field(
        default=False,
        metadata={
            "help": "Enable unpickling of arbitrary code in AutoModelForCausalLM#from_pretrained."
        },
    )
    tokenizer_model_name: Optional[str] = field(
        default="TheBloke/dromedary-65b-lora-HF"
    )
    reward_model_tokenizer_model_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Tokenizer model name for the reward model. "
            "Defaults to the same as the policy model."
        },
    )
    base_model_mapping: Optional[dict] = field(
        default=None,
        metadata={
            "help": "Mapping from base model names to other base model names. "
            "This is useful for loading models trained on different clusters."
        },
    )


@dataclass
class DataArguments:
    dataset_path: str = field(default="tatsu-lab/alpaca_farm")
    dataset_name: str = field(default="alpaca_instructions")
    eval_dataset_path: str = field(default="tatsu-lab/alpaca_farm")
    eval_dataset_name: str = field(default="alpaca_instructions")
    train_splits: List[str] = field(default_factory=lambda: ["unlabeled"])
    eval_splits: List[str] = field(default_factory=lambda: ["val"])
    policy_meta_prompt_pattern: Optional[str] = field(
        default=None, metadata={"help": "Which meta prompt pattern to use."}
    )
    reward_meta_prompt_pattern: Optional[str] = field(
        default=None, metadata={"help": "Which meta prompt pattern to use."}
    )
    principle_collection_path: Optional[str] = field(
        default=None, metadata={"help": "Path to the principle collection."}
    )
    max_principles: Optional[int] = field(
        default=None, metadata={"help": "Maximum number of principles to use."}
    )
    stop_token: Optional[str] = field(
        default=None,
        metadata={"help": "Token to stop generation with."},
    )


@dataclass
class TrainingArguments(transformers.Seq2SeqTrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    # From AlpacaFarm
    truncate_tokens: Optional[List[str]] = field(
        default_factory=lambda: None,
        metadata={
            "help": "Tokens in strings to truncate at first occurrence. "
            "This was used in original OAI summarization paper to avoid models returning incomplete sentences. "
        },
    )
    truncate_after: Optional[int] = field(
        default=None,
        metadata={
            "help": "Truncate after this number of tokens. Prevents early truncation."
        },
    )
    penalty_reward_value: float = field(
        default=-1.0,
        metadata={
            "help": "Reward assigned to sequences that are truncated, "
            "e.g., due to outputting incomplete sentences for given context window."
        },
    )
    penalize_no_stop_token: bool = field(
        default=False,
        metadata={
            "help": "Whether to penalize sequences that do not contain stop_token."
        },
    )
    length_bonus_score: float = field(
        default=0.0,
        metadata={
            "help": "Add the reward for longer sequences by this amount. "
            "This is useful for encouraging longer sequences."
        },
    )
    length_bonus_upper_bound: float = field(
        default=1.0,
        metadata={"help": "Upper bound on the length bonus. "},
    )
    policy_model_bits: int = field(
        default=4, metadata={"help": "How many bits to use."}
    )
    reward_model_bits: int = field(
        default=4, metadata={"help": "How many bits to use."}
    )
    fully_initialize_policy: bool = field(
        default=False,
        metadata={
            "help": "Whether to fully initialize the policy model. "
            "This is useful for debugging."
        },
    )
    enable_reasoning_principles: bool = field(
        default=False,
        metadata={
            "help": "Whether to enable reasoning principles. "
            "This is useful for improving reasoning capabilities."
        },
    )
    enable_redteaming_principles: bool = field(
        default=False,
        metadata={
            "help": "Whether to enable redteaming principles. "
            "This is useful for improving harmlessness."
        },
    )
    relative_stop_token_penalty: bool = field(
        default=False,
        metadata={
            "help": "Whether to penalize sequences that do not contain stop_token "
            "with a relative penalty based on the original reward."
        },
    )
    clean_tokens_after_eos: bool = field(
        default=False,
        metadata={
            "help": "Whether to clean up tokens after the first occurrence of stop_token."
        },
    )
    total_epochs: int = field(default=10)
    rollout_batch_size: int = field(default=512)
    step_batch_size: int = field(default=256)
    rollout_per_device_batch_size: int = field(default=32)
    step_per_device_batch_size: int = field(default=2)
    reward_model_per_device_batch_size: int = field(default=None)
    noptepochs: int = field(default=2)
    vf_coef: float = field(default=0.1)
    cliprange: float = field(default=0.2)
    cliprange_value: float = field(default=0.2)
    gamma: float = field(default=1.0)
    lam: float = field(default=1.0)
    whiten_rewards: bool = field(default=True)
    temperature: float = field(default=1.0)
    kl_coef: float = field(default=0.2)
    kl_approximator: str = field(default="k1")
    target_kl: float = field(default=6.0)
    k_beta: float = field(default=0.1)
    adaptive_kl: bool = field(default=False)
    eval_batches: int = field(
        default=sys.maxsize,
        metadata={"help": "Maximum number of batches to evaluate on."},
    )
    init_value_with_reward: bool = field(
        default=True,
        metadata={"help": "Initialize the value model with the reward model."},
    )
    save_steps_extra: Optional[str] = field(
        default=None,
        metadata={
            "help": "A list of predetermined checkpoints to save, represented in the format 'no1__no2__no3'. "
            "Parse this with str.split('__')."
        },
    )
    query_len: int = field(default=128)
    min_token_limit: int = field(default=None)
    response_len: int = field(default=384)
    model_max_length: int = field(default=1024)
    whitening_async_stats: str = field(
        default="per_gpu",
        metadata={"help": "How to sync statistics for advantage whitening."},
    )
    # From QLoRA
    full_finetune: bool = field(
        default=False, metadata={"help": "Finetune the entire model without adapters."}
    )
    adam8bit: bool = field(default=False, metadata={"help": "Use 8-bit adam."})
    report_to: str = field(
        default="none",
        metadata={"help": "To use wandb or something else for reporting."},
    )
    resume_dir: str = field(
        default=None,
        metadata={
            "help": "The directory to resume from. If None, will start from scratch."
        },
    )
    output_dir: str = field(
        default="./output", metadata={"help": "The output dir for logs and checkpoints"}
    )
    optim: str = field(
        default="paged_adamw_32bit", metadata={"help": "The optimizer to be used"}
    )
    per_device_train_batch_size: int = field(
        default=1,
        metadata={
            "help": "The training batch size per GPU. Increase for better speed."
        },
    )
    gradient_accumulation_steps: int = field(
        default=16,
        metadata={
            "help": "How many gradients to accumulate before to perform an optimizer step"
        },
    )
    weight_decay: float = field(
        default=0.0, metadata={"help": "The L2 weight decay rate of AdamW"}
    )  # use lora dropout instead for regularization if needed
    learning_rate: float = field(default=0.0002, metadata={"help": "The learnign rate"})
    remove_unused_columns: bool = field(
        default=False,
        metadata={"help": "Removed unused columns. Needed to make this codebase work."},
    )
    max_grad_norm: float = field(
        default=0.3,
        metadata={
            "help": "Gradient clipping max norm. This is tuned and works well for all models tested."
        },
    )
    gradient_checkpointing: bool = field(
        default=True,
        metadata={"help": "Use gradient checkpointing. You want to use this."},
    )
    do_train: bool = field(
        default=True,
        metadata={"help": "To train or not to train, that is the question?"},
    )
    lr_scheduler_type: str = field(
        default="constant",
        metadata={
            "help": "Learning rate schedule. Constant a bit better than cosine, and has advantage for analysis"
        },
    )
    logging_steps: int = field(
        default=10,
        metadata={"help": "The frequency of update steps after which to log the loss"},
    )
    group_by_length: bool = field(
        default=True,
        metadata={
            "help": "Group sequences into batches with same length. Saves memory and speeds up training considerably."
        },
    )
    save_strategy: str = field(
        default="steps", metadata={"help": "When to save checkpoints"}
    )
    save_steps: int = field(default=250, metadata={"help": "How often to save a model"})
    save_total_limit: int = field(
        default=40,
        metadata={
            "help": "How many checkpoints to save before the oldest is overwritten"
        },
    )
    resume_from_training: bool = field(
        default=False, metadata={"help": "Resume from training"}
    )
    config_file_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "The name of the config file. If not specified, will use the command line arguments."
        },
    )

    def __post_init__(self):
        super().__post_init__()
        world_size = int(os.environ.get("WORLD_SIZE", 1))

        # Checks on rollout_batch_size only matter for PPO.
        assert (
            self.rollout_batch_size >= self.rollout_per_device_batch_size * world_size
        ), (
            "rollout_batch_size is smaller than rollout_per_device_batch_size * world_size. "
            "Increase the former or decrease the latter to fix this."
        )
        assert (
            self.rollout_batch_size % (self.rollout_per_device_batch_size * world_size)
            == 0
        ), "rollout_batch_size is not a multiple of rollout_per_device_batch_size * world_size. "

        assert self.step_batch_size >= self.step_per_device_batch_size * world_size, (
            "step_batch_size is smaller than step_per_device_batch_size * world_size. "
            "Increase the former or decrease the latter to fix this."
        )
        assert (
            self.step_batch_size % (self.step_per_device_batch_size * world_size) == 0
        ), "step_batch_size is not a multiple of step_per_device_batch_size * world_size. "

        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        if local_rank == 0:
            logger.warning(
                f"Rollout stats:\n"
                f"\trollout_batch_size: {self.rollout_batch_size}\n"
                f"\trollout_per_device_batch_size: {self.rollout_per_device_batch_size}\n"
                f"\tworld_size: {world_size}\n",
            )
        assert (
            self.rollout_batch_size // self.rollout_per_device_batch_size
        ) % world_size == 0
        self.rollout_accumulation_steps = (
            self.rollout_batch_size // self.rollout_per_device_batch_size // world_size
        )

        logger.warning(
            f"Step stats:\n"
            f"\tstep_batch_size: {self.step_batch_size}\n"
            f"\tstep_per_device_batch_size: {self.step_per_device_batch_size}\n"
            f"\tworld_size: {world_size}\n",
        )
        assert (
            self.step_batch_size // self.step_per_device_batch_size
        ) % world_size == 0
        self.gradient_accumulation_steps = (
            self.step_batch_size // self.step_per_device_batch_size // world_size
        )

        logger.warning(
            f"Accumulation steps:\n"
            f"\trollout_accumulation_steps: {self.rollout_accumulation_steps}\n"
            f"\tgradient_accumulation_steps: {self.gradient_accumulation_steps}\n"
        )

        if self.save_steps_extra is not None:
            self.save_steps_extra_list = [
                int(string) for string in self.save_steps_extra.split("__")
            ]
        else:
            self.save_steps_extra_list = []

    def set_truncate_token_ids(self, tokenizer: transformers.PreTrainedTokenizer):
        """Convert truncation token to token ids.

        This is called in RLTrainer.
        """
        truncate_tokens = self.truncate_tokens
        if truncate_tokens is None:
            truncate_token_ids = None
        else:
            truncate_token_ids = tokenizer.convert_tokens_to_ids(truncate_tokens)
        self.truncate_token_ids = truncate_token_ids


def rank0_print(*args):
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    if local_rank == 0:
        print(*args)


def train():
    hfparser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )

    try:
        # Assumes that the first .yaml file is the config file (if any)
        config_file = next(iter(arg for arg in sys.argv if arg.endswith(".yaml")))
    except StopIteration:
        config_file = None

    if config_file:
        model_args, data_args, training_args = hfparser.parse_yaml_file(config_file)
    else:
        (
            model_args,
            data_args,
            training_args,
            extra_args,
        ) = hfparser.parse_args_into_dataclasses(return_remaining_strings=True)

    args = argparse.Namespace(
        **vars(model_args), **vars(data_args), **vars(training_args)
    )
    training_args.data_config = data_args

    checkpoint_dir, completed_training = get_last_checkpoint(args.output_dir)

    if checkpoint_dir is None and args.resume_dir is not None:
        checkpoint_dir, _ = get_last_checkpoint(args.resume_dir)
        completed_training = False

    if completed_training:
        rank0_print("Detected that training was already completed!")

    if checkpoint_dir is None:
        rank0_print("Training from scratch.")
    else:
        rank0_print("Loading from checkpoint:", checkpoint_dir)

    accelerator = AlpacaAccelerator(
        log_with=args.report_to,
        project_dir=args.logging_dir,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        even_batches=True,  # Make sure the batch size on each device is the same.
        split_batches=False,  # Don't break a batch into smaller chunks.
        step_scheduler_with_optimizer=False,  # Untie optimizer and scheduler step.
        # Value model might not use all parameters (e.g., lm-head) in the forward pass.
        kwargs_handlers=[
            DistributedDataParallelKwargs(
                find_unused_parameters=args.ddp_find_unused_parameters,
            )
        ],
    )
    dict_args = vars(args)
    # value should be one of int, float, str, bool, or torch.Tensor
    for k in dict_args:
        if type(dict_args[k]) not in [int, float, str, bool, torch.Tensor]:
            dict_args[k] = str(dict_args[k])
    # print(dict_args)
    accelerator.init_trackers(
        "rl_trainer",
        config=dict_args,
        init_kwargs={
            "wandb": {
                "name": args.run_name or args.output_dir.split("/")[-1],
                "dir": args.output_dir,
                "resume": "allow",
            }
        },
    )
    logger.warning(
        accelerator.state,
        # main_process_only=False,
    )

    use_llama_base_model = (
        "dromedary" in args.tokenizer_model_name.lower()
        or "llama" in args.tokenizer_model_name.lower()
    )

    if use_llama_base_model:
        tokenizer_model_name = (
            "TheBloke/dromedary-65b-lora-HF"  # TODO(zhiqings): hacking
        )
        TokenizerClass = LlamaTokenizer
    else:
        tokenizer_model_name = args.tokenizer_model_name
        TokenizerClass = AutoTokenizer

    # Tokenizer
    tokenizer = TokenizerClass.from_pretrained(
        tokenizer_model_name,
        # cache_dir=args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="left",
        truncation_side="right",
    )

    if use_llama_base_model:
        num_new_tokens = 0
        if tokenizer._pad_token is None:
            tokenizer.pad_token_id = (
                0  # unk. we want this to be different from the eos token
            )
    else:
        num_new_tokens = tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN))

    reward_tokenizer = tokenizer
    reward_num_new_tokens = 0
    if (
        args.reward_model_tokenizer_model_name is not None
        and args.reward_model_tokenizer_model_name != args.tokenizer_model_name
    ):
        reward_tokenizer = AutoTokenizer.from_pretrained(
            args.reward_model_tokenizer_model_name,
            # cache_dir=args.cache_dir,
            model_max_length=training_args.model_max_length,
            padding_side="left",
            truncation_side="right",
        )
        reward_num_new_tokens = reward_tokenizer.add_special_tokens(
            dict(pad_token=DEFAULT_PAD_TOKEN)
        )
        print("Using different reward tokenizer:", tokenizer, reward_tokenizer)

    # Dataset
    data_module: dict = make_rl_data_module(
        tokenizer=tokenizer, data_args=data_args, training_args=training_args
    )

    if accelerator.is_main_process:
        for i in range(3):
            token_ids = data_module["train_dataset"][i]["queries"]
            print(tokenizer.decode(token_ids, skip_special_tokens=True))
            print("=" * 20)

    rank = int(os.environ.get("RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    node_id = rank // torch.cuda.device_count()

    print(f"Distributed info: rank={rank}, world_size={world_size}, node_id={node_id}")
    model_module = make_models(
        tokenizer=tokenizer,
        reward_tokenizer=reward_tokenizer,
        args=args,
        accelerator=accelerator,
        num_new_tokens=num_new_tokens,
        reward_num_new_tokens=reward_num_new_tokens,
        resume_from_checkpoint=(
            checkpoint_dir if training_args.resume_from_training else None
        ),
    )

    trainer = PPOTrainer(
        args=training_args,
        accelerator=accelerator,
        **data_module,
        **model_module,
        tokenizer=tokenizer,
        reward_tokenizer=reward_tokenizer,
    )

    trainer.train(
        resume_training_ckpt=checkpoint_dir
        if training_args.resume_from_training
        else None
    )


if __name__ == "__main__":
    train()

training/train_qlora_ppo.py
IBM/SALMON
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import json

import os
from dataclasses import dataclass, field
from typing import Optional, List, Literal
import logging

import torch
import transformers
import argparse
from transformers import set_seed

try:
    from transformers import LlamaTokenizerFast as LlamaTokenizer

    print("Using fast tokenizer")
except:
    from transformers import LlamaTokenizer

    print("Using slow tokenizer")

from transformers import AutoTokenizer, AutoModelForCausalLM

from qlora_utils import (
    SavePeftModelCallback,
    print_trainable_parameters,
    get_last_checkpoint,
    DEFAULT_PAD_TOKEN,
)
from data_utils.data_utils_rm import make_binary_reward_modeling_data_module
from models.reward_model import (
    RewardConfig,
    RewardModel,
    RewardModelTrainer as Trainer,
    compute_reward_modeling_metrics,
)

torch.backends.cuda.matmul.allow_tf32 = True

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="EleutherAI/pythia-12b")
    trust_remote_code: Optional[bool] = field(
        default=False,
        metadata={
            "help": "Enable unpickling of arbitrary code in AutoModelForCausalLM#from_pretrained."
        },
    )


@dataclass
class DataArguments:
    dataset_path: str = field(default="tatsu-lab/alpaca_farm")
    dataset_name: str = field(default=None, metadata={"help": "Dataset name"})
    eval_dataset_path: str = field(default="tatsu-lab/alpaca_farm")
    eval_dataset_name: str = field(default="alpaca_human_preference")
    eval_size: int = field(
        default=500,
        metadata={
            "help": "Number of examples to split out from training to use for evaluation."
        },
    )
    meta_prompt_pattern: Optional[str] = field(
        default=None, metadata={"help": "Which meta prompt pattern to use."}
    )
    principle_collection_path: Optional[str] = field(
        default=None, metadata={"help": "Path to the principle collection."}
    )


@dataclass
class TrainingArguments(transformers.Seq2SeqTrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    # From AlpacaFarm
    model_max_length: int = field(
        default=512,
        metadata={
            "help": "Maximum sequence length. Sequences will be left padded to this length always during training."
        },
    )
    query_len: int = field(default=None, metadata={"help": "Length of the query."})
    response_len: int = field(
        default=None, metadata={"help": "Length of the response."}
    )
    label_names: List[str] = field(
        default_factory=lambda: ["index_0", "index_1", "choice"],
        metadata={
            "help": "Names of the labels in the dataset. "
            "This is needed to get transformers.Trainer to not throw those tensors away before `compute_loss`."
            "By default, the trainer throws away columns it doesn't recognize when creating the "
            "`train_dataloader` (see `_remove_unused_columns`). "
        },
    )
    padding: Literal["max_length", "longest"] = field(
        default="longest",
        metadata={
            "help": "Padding strategy. If 'max_length', pads to `model_max_length` always; this might lead to some "
            "redundant compute. If 'longest', pads to the longest sequence in the batch, capped by `model_max_length`."
        },
    )
    end_sequence_with_eos: bool = field(
        default=False,
        metadata={
            "help": "Whether to end sequences with EOS. "
            "Ending with EOS might help the reward model realize it's time to predict."
        },
    )
    # From QLoRA
    full_finetune: bool = field(
        default=False, metadata={"help": "Finetune the entire model without adapters."}
    )
    adam8bit: bool = field(default=False, metadata={"help": "Use 8-bit adam."})
    double_quant: bool = field(
        default=True,
        metadata={
            "help": "Compress the quantization statistics through double quantization."
        },
    )
    quant_type: str = field(
        default="nf4",
        metadata={
            "help": "Quantization data type to use. Should be one of `fp4` or `nf4`."
        },
    )
    bits: int = field(default=4, metadata={"help": "How many bits to use."})
    lora_modules: Optional[List[str]] = field(
        default=None,
        metadata={
            "help": "Which modules to use LoRA on. If None, will use all linear layers."
        },
    )
    lora_r: int = field(default=64, metadata={"help": "Lora R dimension."})
    lora_alpha: float = field(default=16, metadata={"help": " Lora alpha."})
    lora_dropout: float = field(default=0.0, metadata={"help": "Lora dropout."})
    report_to: str = field(
        default="none",
        metadata={"help": "To use wandb or something else for reporting."},
    )
    resume_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Path to the directory containing the checkpoint to resume."},
    )
    output_dir: str = field(
        default="./output", metadata={"help": "The output dir for logs and checkpoints"}
    )
    optim: str = field(
        default="paged_adamw_32bit", metadata={"help": "The optimizer to be used"}
    )
    per_device_train_batch_size: int = field(
        default=1,
        metadata={
            "help": "The training batch size per GPU. Increase for better speed."
        },
    )
    gradient_accumulation_steps: int = field(
        default=16,
        metadata={
            "help": "How many gradients to accumulate before to perform an optimizer step"
        },
    )
    weight_decay: float = field(
        default=0.0, metadata={"help": "The L2 weight decay rate of AdamW"}
    )  # use lora dropout instead for regularization if needed
    learning_rate: float = field(default=0.0002, metadata={"help": "The learnign rate"})
    remove_unused_columns: bool = field(
        default=False,
        metadata={"help": "Removed unused columns. Needed to make this codebase work."},
    )
    max_grad_norm: float = field(
        default=0.3,
        metadata={
            "help": "Gradient clipping max norm. This is tuned and works well for all models tested."
        },
    )
    gradient_checkpointing: bool = field(
        default=True,
        metadata={"help": "Use gradient checkpointing. You want to use this."},
    )
    do_train: bool = field(
        default=True,
        metadata={"help": "To train or not to train, that is the question?"},
    )
    lr_scheduler_type: str = field(
        default="constant",
        metadata={
            "help": "Learning rate schedule. Constant a bit better than cosine, and has advantage for analysis"
        },
    )
    warmup_ratio: float = field(
        default=0.03, metadata={"help": "Fraction of steps to do a warmup for"}
    )
    logging_steps: int = field(
        default=10,
        metadata={"help": "The frequency of update steps after which to log the loss"},
    )
    group_by_length: bool = field(
        default=True,
        metadata={
            "help": "Group sequences into batches with same length. Saves memory and speeds up training considerably."
        },
    )
    save_strategy: str = field(
        default="steps", metadata={"help": "When to save checkpoints"}
    )
    save_steps: int = field(default=250, metadata={"help": "How often to save a model"})
    save_total_limit: int = field(
        default=40,
        metadata={
            "help": "How many checkpoints to save before the oldest is overwritten"
        },
    )
    resume_from_training: bool = field(
        default=False, metadata={"help": "Resume from training"}
    )


def rank0_print(*args):
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    if local_rank == 0:
        print(*args)


def train():
    hfparser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )
    (
        model_args,
        data_args,
        training_args,
        extra_args,
    ) = hfparser.parse_args_into_dataclasses(return_remaining_strings=True)
    args = argparse.Namespace(
        **vars(model_args), **vars(data_args), **vars(training_args)
    )

    if args.resume_dir is not None:
        checkpoint_dir, completed_training = args.resume_dir, False
    else:
        checkpoint_dir, completed_training = get_last_checkpoint(args.output_dir)

    if completed_training:
        rank0_print("Detected that training was already completed!")

    if checkpoint_dir is None:
        rank0_print("Training from scratch.")
    else:
        rank0_print("Loading from checkpoint:", checkpoint_dir)
        if args.resume_from_training:
            rank0_print("Resuming from training not supported yet. Exiting.")
            exit(1)

    use_llama_base_model = (
        "dromedary" in args.model_name_or_path.lower()
        or "llama" in args.model_name_or_path.lower()
        or "vicuna" in args.model_name_or_path.lower()
    )

    if use_llama_base_model:
        tokenizer_model_name = (
            "TheBloke/dromedary-65b-lora-HF"  # TODO(zhiqings): hacking
        )
        TokenizerClass = LlamaTokenizer
    else:
        tokenizer_model_name = args.model_name_or_path
        TokenizerClass = AutoTokenizer

    # Tokenizer
    tokenizer = TokenizerClass.from_pretrained(
        tokenizer_model_name,
        cache_dir=args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="left",
        truncation_side="right",
    )

    if use_llama_base_model:
        if tokenizer._pad_token is None:
            tokenizer.pad_token_id = (
                0  # unk. we want this to be different from the eos token
            )
    else:
        raise NotImplementedError

    data_module = make_binary_reward_modeling_data_module(
        tokenizer=tokenizer,
        data_args=data_args,
        training_args=training_args,
    )

    if args.do_train:
        training_data = data_module["train_dataset"]
        rank0_print("Training data size:", len(training_data))
        rank0_print("Training data example:")
        for i in range(min(3, len(training_data))):
            ex_input_ids_0 = training_data[i]["input_ids"][0]
            rank0_print(tokenizer.decode(ex_input_ids_0, skip_special_tokens=True))
            rank0_print("=" * 20)
            ex_input_ids_1 = training_data[i]["input_ids"][1]
            rank0_print(tokenizer.decode(ex_input_ids_1, skip_special_tokens=True))
            rank0_print("=" * 20)
            rank0_print("=" * 20)

    config = RewardConfig(backbone_model_name_or_path=model_args.model_name_or_path)

    model = RewardModel(
        args=args,
        config=config,
        qlora=True,
        checkpoint_dir=checkpoint_dir,
    )

    model.backbone_model.config.use_cache = False
    print_trainable_parameters(args, model)
    print("loaded model")
    set_seed(args.seed)

    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        args=training_args,
        compute_metrics=compute_reward_modeling_metrics,
        **{k: v for k, v in data_module.items() if k != "predict_dataset"},
    )

    # Callbacks
    if not args.full_finetune:
        trainer.add_callback(SavePeftModelCallback)

    # Verifying the datatypes.
    dtypes = {}
    for _, p in model.named_parameters():
        dtype = p.dtype
        if dtype not in dtypes:
            dtypes[dtype] = 0
        dtypes[dtype] += p.numel()
    total = 0
    for k, v in dtypes.items():
        total += v
    for k, v in dtypes.items():
        print(k, v, v / total)

    all_metrics = {"run_name": args.run_name}

    # Training
    if args.do_train:
        logger.info("*** Train ***")
        train_result = trainer.train()
        metrics = train_result.metrics
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()
        all_metrics.update(metrics)

    # Evaluation
    if args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate(metric_key_prefix="eval")
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
        all_metrics.update(metrics)

    if args.do_train or args.do_eval:
        with open(os.path.join(args.output_dir, "metrics.json"), "w") as fout:
            fout.write(json.dumps(all_metrics))


if __name__ == "__main__":
    train()

training/train_qlora_rm.py
IBM/SALMON
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import json

import os
from dataclasses import dataclass, field
from typing import Optional, List
import logging

import torch
import transformers
import argparse
from transformers import (
    set_seed,
    Trainer,
)

try:
    from transformers import LlamaTokenizerFast as LlamaTokenizer

    print("Using fast tokenizer")
except:
    from transformers import LlamaTokenizer

    print("Using slow tokenizer")

from transformers import AutoTokenizer

from qlora_utils import (
    SavePeftModelCallback,
    print_trainable_parameters,
    get_last_checkpoint,
    DEFAULT_PAD_TOKEN,
)
from data_utils.data_utils_sft import (
    make_sft_data_module,
    IGNORE_INDEX,
)
from models.qlora_model import get_accelerate_model


torch.backends.cuda.matmul.allow_tf32 = True

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="EleutherAI/pythia-12b")
    trust_remote_code: Optional[bool] = field(
        default=False,
        metadata={
            "help": "Enable unpickling of arbitrary code in AutoModelForCausalLM#from_pretrained."
        },
    )


@dataclass
class DataArguments:
    eval_dataset_size: int = field(
        default=1024, metadata={"help": "Size of validation dataset."}
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
            "value if set."
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
            "value if set."
        },
    )
    source_max_len: int = field(
        default=1024,
        metadata={
            "help": "Maximum source sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    target_max_len: int = field(
        default=256,
        metadata={
            "help": "Maximum target sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    dataset: str = field(
        default="alpaca",
        metadata={"help": "Which dataset to finetune on. See datamodule for options."},
    )
    dataset_format: Optional[str] = field(
        default=None,
        metadata={
            "help": "Which dataset format is used. [alpaca|chip2|self-instruct|hh-rlhf]"
        },
    )
    meta_prompt_pattern: Optional[str] = field(
        default=None, metadata={"help": "Which meta prompt pattern to use."}
    )
    add_eos_to_target: bool = field(
        default=True, metadata={"help": "Whether to add an EOS token to the target."}
    )


@dataclass
class TrainingArguments(transformers.Seq2SeqTrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    train_on_source: Optional[bool] = field(
        default=False,
        metadata={
            "help": "Whether to train on the input in addition to the target text."
        },
    )
    full_finetune: bool = field(
        default=False, metadata={"help": "Finetune the entire model without adapters."}
    )
    adam8bit: bool = field(default=False, metadata={"help": "Use 8-bit adam."})
    double_quant: bool = field(
        default=True,
        metadata={
            "help": "Compress the quantization statistics through double quantization."
        },
    )
    quant_type: str = field(
        default="nf4",
        metadata={
            "help": "Quantization data type to use. Should be one of `fp4` or `nf4`."
        },
    )
    bits: int = field(default=4, metadata={"help": "How many bits to use."})
    lora_modules: Optional[List[str]] = field(
        default=None,
        metadata={
            "help": "Which modules to use LoRA on. If None, will use all linear layers."
        },
    )
    lora_r: int = field(default=64, metadata={"help": "Lora R dimension."})
    lora_alpha: float = field(default=16, metadata={"help": " Lora alpha."})
    lora_dropout: float = field(default=0.0, metadata={"help": "Lora dropout."})
    max_memory_MB: int = field(default=80000, metadata={"help": "Free memory per gpu."})
    report_to: str = field(
        default="none",
        metadata={"help": "To use wandb or something else for reporting."},
    )
    resume_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Path to the directory containing the checkpoint to resume."},
    )
    output_dir: str = field(
        default="./output", metadata={"help": "The output dir for logs and checkpoints"}
    )
    optim: str = field(
        default="paged_adamw_32bit", metadata={"help": "The optimizer to be used"}
    )
    per_device_train_batch_size: int = field(
        default=1,
        metadata={
            "help": "The training batch size per GPU. Increase for better speed."
        },
    )
    gradient_accumulation_steps: int = field(
        default=16,
        metadata={
            "help": "How many gradients to accumulate before to perform an optimizer step"
        },
    )
    weight_decay: float = field(
        default=0.0, metadata={"help": "The L2 weight decay rate of AdamW"}
    )  # use lora dropout instead for regularization if needed
    learning_rate: float = field(default=0.0002, metadata={"help": "The learnign rate"})
    remove_unused_columns: bool = field(
        default=False,
        metadata={"help": "Removed unused columns. Needed to make this codebase work."},
    )
    max_grad_norm: float = field(
        default=0.3,
        metadata={
            "help": "Gradient clipping max norm. This is tuned and works well for all models tested."
        },
    )
    gradient_checkpointing: bool = field(
        default=True,
        metadata={"help": "Use gradient checkpointing. You want to use this."},
    )
    do_train: bool = field(
        default=True,
        metadata={"help": "To train or not to train, that is the question?"},
    )
    lr_scheduler_type: str = field(
        default="constant",
        metadata={
            "help": "Learning rate schedule. Constant a bit better than cosine, and has advantage for analysis"
        },
    )
    warmup_ratio: float = field(
        default=0.03, metadata={"help": "Fraction of steps to do a warmup for"}
    )
    logging_steps: int = field(
        default=10,
        metadata={"help": "The frequency of update steps after which to log the loss"},
    )
    group_by_length: bool = field(
        default=True,
        metadata={
            "help": "Group sequences into batches with same length. Saves memory and speeds up training considerably."
        },
    )
    save_strategy: str = field(
        default="steps", metadata={"help": "When to save checkpoints"}
    )
    save_steps: int = field(default=250, metadata={"help": "How often to save a model"})
    save_total_limit: int = field(
        default=40,
        metadata={
            "help": "How many checkpoints to save before the oldest is overwritten"
        },
    )
    resume_from_training: bool = field(
        default=False, metadata={"help": "Resume from training"}
    )


def rank0_print(*args):
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    if local_rank == 0:
        print(*args)


def train():
    hfparser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )
    (
        model_args,
        data_args,
        training_args,
        extra_args,
    ) = hfparser.parse_args_into_dataclasses(return_remaining_strings=True)
    args = argparse.Namespace(
        **vars(model_args), **vars(data_args), **vars(training_args)
    )

    if args.resume_dir is not None:
        checkpoint_dir, completed_training = args.resume_dir, False
    else:
        checkpoint_dir, completed_training = get_last_checkpoint(args.output_dir)

    if completed_training:
        rank0_print("Detected that training was already completed!")

    if checkpoint_dir is None:
        rank0_print("Training from scratch.")
    else:
        rank0_print("Loading from checkpoint:", checkpoint_dir)
        if args.resume_from_training:
            rank0_print("Resuming from training not supported yet. Exiting.")
            exit(1)

    use_llama_base_model = (
        "dromedary" in args.model_name_or_path.lower()
        or "llama" in args.model_name_or_path.lower()
    )

    if use_llama_base_model:
        tokenizer_model_name = (
            "TheBloke/dromedary-65b-lora-HF"  # a random llama-based model
        )
        TokenizerClass = LlamaTokenizer
    else:
        tokenizer_model_name = args.model_name_or_path
        TokenizerClass = AutoTokenizer

    left_truncated_tokenizer = TokenizerClass.from_pretrained(
        tokenizer_model_name,
        cache_dir=args.cache_dir,
        truncation_side="left",
        padding_side="right",
        # use_fast=False, # Fast tokenizer giving issues.
    )

    # Tokenizer
    tokenizer = TokenizerClass.from_pretrained(
        tokenizer_model_name,
        cache_dir=args.cache_dir,
        truncation_side="right",
        padding_side="right",
    )

    if use_llama_base_model:
        if tokenizer._pad_token is None:
            left_truncated_tokenizer.pad_token_id = (
                0  # unk. we want this to be different from the eos token
            )
            tokenizer.pad_token_id = (
                0  # unk. we want this to be different from the eos token
            )

    else:
        raise NotImplementedError

    data_module = make_sft_data_module(
        left_truncated_tokenizer=left_truncated_tokenizer,
        tokenizer=tokenizer,
        args=args,
    )

    model = get_accelerate_model(args, checkpoint_dir)

    model.config.use_cache = False
    print_trainable_parameters(args, model)
    print("loaded model")
    set_seed(args.seed)

    if args.do_train:
        training_data = data_module["train_dataset"]
        rank0_print("Training data size:", len(training_data))
        rank0_print("Training data example:")
        for i in range(min(3, len(training_data))):
            rank0_print(training_data[i])

    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        args=training_args,
        **{k: v for k, v in data_module.items() if k != "predict_dataset"},
    )

    # Callbacks
    if not args.full_finetune:
        trainer.add_callback(SavePeftModelCallback)

    # Verifying the datatypes.
    dtypes = {}
    for _, p in model.named_parameters():
        dtype = p.dtype
        if dtype not in dtypes:
            dtypes[dtype] = 0
        dtypes[dtype] += p.numel()
    total = 0
    for k, v in dtypes.items():
        total += v
    for k, v in dtypes.items():
        print(k, v, v / total)

    all_metrics = {"run_name": args.run_name}
    # Training
    if args.do_train:
        logger.info("*** Train ***")
        # Note: `resume_from_checkpoint` not supported for adapter checkpoints by HF.
        # Currently adapter checkpoint is reloaded as expected but optimizer/scheduler states are not.
        train_result = trainer.train()
        metrics = train_result.metrics
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()
        all_metrics.update(metrics)

    if args.do_train:
        with open(os.path.join(args.output_dir, "metrics.json"), "w") as fout:
            fout.write(json.dumps(all_metrics))


if __name__ == "__main__":
    train()

training/train_qlora_sft.py
IBM/SALMON