Add src folder

src/api_demo.py

@@ -0,0 +1,101 @@
+# coding=utf-8
+# Implements API for fine-tuned models.
+# Usage: python api_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint
+
+# Request:
+# curl http://127.0.0.1:8000 --header 'Content-Type: application/json' --data '{"prompt": "Hello there!", "history": []}'
+
+# Response:
+# {
+#   "response": "'Hi there!'",
+#   "history": "[('Hello there!', 'Hi there!')]",
+#   "status": 200,
+#   "time": "2000-00-00 00:00:00"
+# }
+
+
+import json
+import torch
+import uvicorn
+import datetime
+from fastapi import FastAPI, Request
+
+from utils import (
+    Template,
+    load_pretrained,
+    prepare_infer_args,
+    get_logits_processor
+)
+
+
+def torch_gc():
+    if torch.cuda.is_available():
+        num_gpus = torch.cuda.device_count()
+        for device_id in range(num_gpus):
+            with torch.cuda.device(device_id):
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()
+
+
+app = FastAPI()
+
+
+@app.post("/")
+async def create_item(request: Request):
+    global model, tokenizer, prompt_template, generating_args
+
+    # Parse the request JSON
+    json_post_raw = await request.json()
+    json_post = json.dumps(json_post_raw)
+    json_post_list = json.loads(json_post)
+    prompt = json_post_list.get("prompt")
+    history = json_post_list.get("history")
+    max_new_tokens = json_post_list.get("max_new_tokens", None)
+    top_p = json_post_list.get("top_p", None)
+    temperature = json_post_list.get("temperature", None)
+
+    # Tokenize the input prompt
+    input_ids = tokenizer([prompt_template.get_prompt(prompt, history)], return_tensors="pt")["input_ids"]
+    input_ids = input_ids.to(model.device)
+
+    # Generation arguments
+    gen_kwargs = generating_args.to_dict()
+    gen_kwargs["input_ids"] = input_ids
+    gen_kwargs["logits_processor"] = get_logits_processor()
+    gen_kwargs["max_new_tokens"] = max_new_tokens if max_new_tokens else gen_kwargs["max_new_tokens"]
+    gen_kwargs["top_p"] = top_p if top_p else gen_kwargs["top_p"]
+    gen_kwargs["temperature"] = temperature if temperature else gen_kwargs["temperature"]
+
+    # Generate response
+    with torch.no_grad():
+        generation_output = model.generate(**gen_kwargs)
+    outputs = generation_output.tolist()[0][len(input_ids[0]):]
+    response = tokenizer.decode(outputs, skip_special_tokens=True)
+
+    # Update history
+    history = history + [(prompt, response)]
+
+    # Prepare response
+    now = datetime.datetime.now()
+    time = now.strftime("%Y-%m-%d %H:%M:%S")
+    answer = {
+        "response": repr(response),
+        "history": repr(history),
+        "status": 200,
+        "time": time
+    }
+
+    # Log and clean up
+    log = "[" + time + "] " + "\", prompt:\"" + prompt + "\", response:\"" + repr(response) + "\""
+    print(log)
+    torch_gc()
+
+    return answer
+
+
+if __name__ == "__main__":
+    model_args, data_args, finetuning_args, generating_args = prepare_infer_args()
+    model, tokenizer = load_pretrained(model_args, finetuning_args)
+    prompt_template = Template(data_args.prompt_template)
+
+    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)

src/cli_demo.py

@@ -0,0 +1,70 @@
+# coding=utf-8
+# Implements stream chat in command line for fine-tuned models.
+# Usage: python cli_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint
+
+
+from utils import (
+    Template,
+    load_pretrained,
+    prepare_infer_args,
+    get_logits_processor
+)
+from threading import Thread
+from transformers import TextIteratorStreamer
+
+
+def main():
+
+    model_args, data_args, finetuning_args, generating_args = prepare_infer_args()
+    model, tokenizer = load_pretrained(model_args, finetuning_args)
+
+    model_name = "BLOOM" if "bloom" in model_args.model_name_or_path else "LLaMA"
+    prompt_template = Template(data_args.prompt_template)
+
+    def predict_and_print(query, history: list) -> list:
+        input_ids = tokenizer([prompt_template.get_prompt(query, history)], return_tensors="pt")["input_ids"]
+        input_ids = input_ids.to(model.device)
+
+        streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+
+        gen_kwargs = generating_args.to_dict()
+        gen_kwargs["input_ids"] = input_ids
+        gen_kwargs["logits_processor"] = get_logits_processor()
+        gen_kwargs["streamer"] = streamer
+
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+
+        print("{}: ".format(model_name), end="", flush=True)
+        response = ""
+        for new_text in streamer:
+            print(new_text, end="", flush=True)
+            response += new_text
+        print()
+        history = history + [(query, response)]
+        return history
+
+    history = []
+    print("欢迎使用 {} 模型，输入内容即可对话，clear清空对话历史，stop终止程序".format(model_name))
+    while True:
+        try:
+            query = input("\nInput: ")
+        except UnicodeDecodeError:
+            print("Detected decoding error at the inputs, please set the terminal encoding to utf-8.")
+            continue
+        except Exception:
+            raise
+
+        if query.strip() == "stop":
+            break
+
+        if query.strip() == "clear":
+            history = []
+            print("History has been removed.")
+            continue
+
+        history = predict_and_print(query, history)
+
+
+if __name__ == "__main__":
+    main()

src/export_model.py

@@ -0,0 +1,19 @@
+# coding=utf-8
+# Exports the fine-tuned model.
+# Usage: python export_model.py --checkpoint_dir path_to_checkpoint --output_dir path_to_save_model
+
+
+from utils import load_pretrained, prepare_args
+
+
+def main():
+
+    model_args, _, training_args, finetuning_args = prepare_args(stage="sft")
+    model, tokenizer = load_pretrained(model_args, finetuning_args)
+    model.save_pretrained(training_args.output_dir, max_shard_size="10GB")
+    tokenizer.save_pretrained(training_args.output_dir)
+    print("model and tokenizer have been saved at:", training_args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

src/train_ppo.py

@@ -0,0 +1,81 @@
+# coding=utf-8
+# Implements parameter-efficient PPO training of fine-tuned models.
+# This code is inspired by:
+# https://github.com/lvwerra/trl/blob/main/examples/sentiment/scripts/gpt-neox-20b_peft/gpt-neo-20b_sentiment_peft.py
+
+import math
+
+from torch.optim import AdamW
+from transformers.optimization import get_scheduler
+from trl import PPOConfig
+
+from utils import (
+    DynamicDataCollatorWithPadding,
+    PPOPeftTrainer,
+    LogCallback,
+    load_pretrained,
+    prepare_args,
+    prepare_data,
+    preprocess_data,
+    plot_loss
+)
+
+
+def main():
+
+    # Prepare pretrained model and dataset
+    model_args, data_args, training_args, finetuning_args = prepare_args(stage="ppo")
+    dataset = prepare_data(model_args, data_args)
+    model, tokenizer = load_pretrained(model_args, finetuning_args, training_args.do_train, stage="ppo")
+    dataset = preprocess_data(dataset, tokenizer, data_args, training_args, stage="ppo")
+    data_collator = DynamicDataCollatorWithPadding(tokenizer)
+
+    ppo_config = PPOConfig(
+        model_name=model_args.model_name_or_path,
+        learning_rate=training_args.learning_rate,
+        mini_batch_size=training_args.per_device_train_batch_size,
+        batch_size=training_args.per_device_train_batch_size,
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        ppo_epochs=1,
+        max_grad_norm=training_args.max_grad_norm
+    )
+
+    optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=ppo_config.learning_rate)
+    total_train_batch_size = \
+        training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
+    lr_scheduler = get_scheduler(
+        training_args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=training_args.warmup_steps,
+        num_training_steps=(training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size))
+    )
+
+    # Initialize our Trainer
+    ppo_trainer = PPOPeftTrainer(
+        training_args=training_args,
+        finetuning_args=finetuning_args,
+        callbacks=[LogCallback()],
+        config=ppo_config,
+        model=model,
+        ref_model=None,
+        tokenizer=tokenizer,
+        dataset=dataset,
+        data_collator=data_collator,
+        optimizer=optimizer,
+        lr_scheduler=lr_scheduler
+    )
+
+    ppo_trainer.ppo_train(max_target_length=data_args.max_target_length)
+    ppo_trainer.save_model()
+    ppo_trainer.save_state() # must be after save_model
+    if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
+        plot_loss(training_args.output_dir, keys=["loss", "reward"])
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

src/train_pt.py

@@ -0,0 +1,81 @@
+# coding=utf-8
+# Implements several parameter-efficient pre-training method.
+# This code is inspired by
+# https://github.com/huggingface/transformers/blob/v4.29.2/examples/pytorch/language-modeling/run_clm.py
+
+
+import math
+
+from utils import (
+    DynamicDataCollatorWithPadding,
+    PeftTrainer,
+    LogCallback,
+    load_pretrained,
+    prepare_args,
+    prepare_data,
+    preprocess_data,
+    plot_loss
+)
+
+
+def main():
+
+    # Prepare pretrained model and dataset
+    model_args, data_args, training_args, finetuning_args = prepare_args(stage="pt")
+    dataset = prepare_data(model_args, data_args)
+    model, tokenizer = load_pretrained(model_args, finetuning_args, training_args.do_train, stage="pt")
+    dataset = preprocess_data(dataset, tokenizer, data_args, training_args, stage="pt")
+    data_collator = DynamicDataCollatorWithPadding(tokenizer, data_args.ignore_pad_token_for_loss)
+
+    # Split the dataset
+    if training_args.do_train:
+        if data_args.dev_ratio > 1e-6:
+            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
+            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
+        else:
+            trainer_kwargs = {"train_dataset": dataset}
+    else: # do_eval or do_predict
+        trainer_kwargs = {"eval_dataset": dataset}
+
+    # Initialize our Trainer
+    trainer = PeftTrainer(
+        finetuning_args=finetuning_args,
+        model=model,
+        args=training_args,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        callbacks=[LogCallback()],
+        **trainer_kwargs
+    )
+
+    # Training
+    if training_args.do_train:
+        train_result = trainer.train()
+        trainer.log_metrics("train", train_result.metrics)
+        trainer.save_metrics("train", train_result.metrics)
+        trainer.save_state()
+        trainer.save_model()
+        if trainer.is_world_process_zero() and model_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
+
+    # Evaluation
+    if training_args.do_eval:
+        metrics = trainer.evaluate(metric_key_prefix="eval")
+
+        try:
+            perplexity = math.exp(metrics["eval_loss"])
+        except OverflowError:
+            perplexity = float("inf")
+        metrics["perplexity"] = perplexity
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

src/train_rm.py

@@ -0,0 +1,76 @@
+# coding=utf-8
+# Implements parameter-efficient training of reward models.
+# This code is inspired by:
+# https://github.com/lvwerra/trl/blob/main/examples/summarization/scripts/reward_summarization.py
+# https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/train_reward_model_gptj.py
+
+
+from utils import (
+    PairwiseDataCollatorWithPadding,
+    PairwisePeftTrainer,
+    LogCallback,
+    load_pretrained,
+    prepare_args,
+    prepare_data,
+    preprocess_data,
+    compute_accuracy,
+    plot_loss
+)
+
+def main():
+
+    # Prepare pretrained model and dataset
+    model_args, data_args, training_args, finetuning_args = prepare_args(stage="rm")
+    dataset = prepare_data(model_args, data_args)
+    model, tokenizer = load_pretrained(model_args, finetuning_args, training_args.do_train, stage="rm")
+    dataset = preprocess_data(dataset, tokenizer, data_args, training_args, stage="rm")
+    data_collator = PairwiseDataCollatorWithPadding(tokenizer)
+
+    training_args.remove_unused_columns = False # important for pairwise dataset
+
+    # Split the dataset
+    if training_args.do_train:
+        if data_args.dev_ratio > 1e-6:
+            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
+            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
+        else:
+            trainer_kwargs = {"train_dataset": dataset}
+    else: # do_eval or do_predict
+        trainer_kwargs = {"eval_dataset": dataset}
+
+    # Initialize our Trainer
+    trainer = PairwisePeftTrainer(
+        finetuning_args=finetuning_args,
+        model=model,
+        args=training_args,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        callbacks=[LogCallback()],
+        compute_metrics=compute_accuracy,
+        **trainer_kwargs
+    )
+
+    # Training
+    if training_args.do_train:
+        train_result = trainer.train()
+        trainer.log_metrics("train", train_result.metrics)
+        trainer.save_metrics("train", train_result.metrics)
+        trainer.save_state()
+        trainer.save_model()
+        if trainer.is_world_process_zero() and model_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
+
+    # Evaluation
+    if training_args.do_eval:
+        metrics = trainer.evaluate(metric_key_prefix="eval")
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

src/train_sft.py

@@ -0,0 +1,97 @@
+# coding=utf-8
+# Implements several parameter-efficient supervised fine-tuning method.
+# This code is inspired by
+# https://github.com/huggingface/transformers/blob/v4.29.2/examples/pytorch/summarization/run_summarization.py
+
+
+from utils import (
+    DynamicDataCollatorWithPadding,
+    Seq2SeqPeftTrainer,
+    ComputeMetrics,
+    LogCallback,
+    load_pretrained,
+    prepare_args,
+    prepare_data,
+    preprocess_data,
+    get_logits_processor,
+    plot_loss
+)
+
+
+def main():
+
+    # Prepare pretrained model and dataset
+    model_args, data_args, training_args, finetuning_args = prepare_args(stage="sft")
+    dataset = prepare_data(model_args, data_args)
+    model, tokenizer = load_pretrained(model_args, finetuning_args, training_args.do_train, stage="sft")
+    dataset = preprocess_data(dataset, tokenizer, data_args, training_args, stage="sft")
+    data_collator = DynamicDataCollatorWithPadding(tokenizer, data_args.ignore_pad_token_for_loss)
+
+    # Override the decoding parameters of Seq2SeqTrainer
+    training_args.generation_max_length = training_args.generation_max_length if \
+                training_args.generation_max_length is not None else data_args.max_target_length
+    training_args.generation_num_beams = data_args.eval_num_beams if \
+                data_args.eval_num_beams is not None else training_args.generation_num_beams
+
+    # Split the dataset
+    if training_args.do_train:
+        if data_args.dev_ratio > 1e-6:
+            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
+            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
+        else:
+            trainer_kwargs = {"train_dataset": dataset}
+    else: # do_eval or do_predict
+        trainer_kwargs = {"eval_dataset": dataset}
+
+    # Initialize our Trainer
+    trainer = Seq2SeqPeftTrainer(
+        finetuning_args=finetuning_args,
+        model=model,
+        args=training_args,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        callbacks=[LogCallback()],
+        compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
+        **trainer_kwargs
+    )
+
+    # Keyword arguments for `model.generate`
+    gen_kwargs = {
+        "do_sample": True,
+        "top_p": 0.7,
+        "max_new_tokens": data_args.max_target_length + 1,
+        "temperature": 0.95,
+        "logits_processor": get_logits_processor()
+    }
+
+    # Training
+    if training_args.do_train:
+        train_result = trainer.train()
+        trainer.log_metrics("train", train_result.metrics)
+        trainer.save_metrics("train", train_result.metrics)
+        trainer.save_state()
+        trainer.save_model()
+        if trainer.is_world_process_zero() and model_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
+
+    # Evaluation
+    if training_args.do_eval:
+        metrics = trainer.evaluate(metric_key_prefix="eval", **gen_kwargs)
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    # Predict
+    if training_args.do_predict:
+        predict_results = trainer.predict(dataset, metric_key_prefix="predict", **gen_kwargs)
+        trainer.log_metrics("predict", predict_results.metrics)
+        trainer.save_metrics("predict", predict_results.metrics)
+        trainer.save_predictions(predict_results, tokenizer)
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

src/utils/__init__.py

@@ -0,0 +1,19 @@
+from .common import (
+    load_pretrained,
+    prepare_args,
+    prepare_infer_args,
+    prepare_data,
+    preprocess_data
+)
+
+from .data_collator import DynamicDataCollatorWithPadding
+
+from .peft_trainer import PeftTrainer, LogCallback
+
+from .seq2seq import ComputeMetrics, Seq2SeqPeftTrainer
+from .pairwise import PairwiseDataCollatorWithPadding, PairwisePeftTrainer, compute_accuracy
+from .ppo import PPOPeftTrainer
+
+from .template import Template
+
+from .other import get_logits_processor, plot_loss

src/utils/common.py

@@ -0,0 +1,561 @@
+import os
+import sys
+import torch
+import hashlib
+from itertools import chain
+from typing import List, Literal, Optional, Tuple
+
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    HfArgumentParser,
+    Seq2SeqTrainingArguments,
+    BitsAndBytesConfig
+)
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+from transformers.modeling_utils import PreTrainedModel
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+import datasets
+from datasets import Dataset, concatenate_datasets, load_dataset
+
+from peft import (
+    PeftModel,
+    TaskType,
+    LoraConfig,
+    get_peft_model
+)
+
+from peft.utils import CONFIG_NAME, WEIGHTS_NAME
+
+from trl import AutoModelForCausalLMWithValueHead
+
+from .config import (
+    ModelArguments,
+    DataTrainingArguments,
+    FinetuningArguments,
+    GeneratingArguments
+)
+
+from .template import Template
+
+from .other import (
+    get_logger,
+    load_trainable_params,
+    load_valuehead_params,
+    print_trainable_params,
+    prepare_model_for_training,
+    IGNORE_INDEX
+)
+
+check_min_version("4.29.1")
+require_version("datasets>=2.12.0", "To fix: pip install datasets>=2.12.0")
+require_version("accelerate>=0.19.0", "To fix: pip install accelerate>=0.19.0")
+require_version("peft>=0.3.0", "To fix: pip install peft>=0.3.0")
+require_version("trl>=0.4.4", "To fix: pip install trl>=0.4.4")
+
+
+logger = get_logger(__name__)
+
+
+def _init_adapter(
+        model: PreTrainedModel,
+        model_args: ModelArguments,
+        finetuning_args: FinetuningArguments,
+        is_trainable: bool,
+        is_mergeable: bool
+) -> PreTrainedModel:
+    r"""
+    Initializes the adapters.
+
+    Support full-parameter, freeze and LoRA training.
+
+    Note that the trainable parameters must be cast to float32.
+    """
+
+    if finetuning_args.finetuning_type == "none" and is_trainable:
+        raise ValueError("You cannot use finetuning_type=none while training.")
+
+    if finetuning_args.finetuning_type == "full":
+        logger.info("Fine-tuning method: Full")
+        model = model.float()
+
+    if finetuning_args.finetuning_type == "freeze":
+        logger.info("Fine-tuning method: Freeze")
+        for name, param in model.named_parameters():
+            if not any(trainable_layer in name for trainable_layer in finetuning_args.trainable_layers):
+                param.requires_grad_(False)
+            else:
+                param.data = param.data.to(torch.float32)
+
+    if model_args.checkpoint_dir is not None:
+        if finetuning_args.finetuning_type != "lora":
+            assert is_mergeable and len(model_args.checkpoint_dir) == 1, "Only LoRA tuning accepts multiple checkpoints."
+            assert load_trainable_params(model, model_args.checkpoint_dir[0]), "Model checkpoint is not correctly loaded."
+        else:
+            assert is_mergeable or len(model_args.checkpoint_dir) == 1, "Quantized model only accepts a single checkpoint."
+
+    if finetuning_args.finetuning_type == "lora":
+        logger.info("Fine-tuning method: LoRA")
+        lastest_checkpoint = None
+
+        if model_args.checkpoint_dir is not None:
+            if os.path.exists(os.path.join(model_args.checkpoint_dir[0], WEIGHTS_NAME)) and \
+                not os.path.exists(os.path.join(model_args.checkpoint_dir[0], CONFIG_NAME)):
+                raise ValueError("The given checkpoint may be not a LoRA checkpoint, \
+                                  please specify `--finetuning_type full/freeze` instead.")
+
+            if (is_trainable and model_args.resume_lora_training) or (not is_mergeable): # continually train on the lora weights
+                checkpoints_to_merge, lastest_checkpoint = model_args.checkpoint_dir[:-1], model_args.checkpoint_dir[-1]
+            else:
+                checkpoints_to_merge = model_args.checkpoint_dir
+
+            for checkpoint in checkpoints_to_merge:
+                model = PeftModel.from_pretrained(model, checkpoint)
+                model = model.merge_and_unload()
+
+            if len(checkpoints_to_merge) > 0:
+                logger.info("Merged {} model checkpoint(s).".format(len(checkpoints_to_merge)))
+
+            if lastest_checkpoint is not None: # resume lora training or quantized inference
+                model = PeftModel.from_pretrained(model, lastest_checkpoint, is_trainable=is_trainable)
+
+        if is_trainable and lastest_checkpoint is None: # create new lora weights while training
+            lora_config = LoraConfig(
+                task_type=TaskType.CAUSAL_LM,
+                inference_mode=False,
+                r=finetuning_args.lora_rank,
+                lora_alpha=finetuning_args.lora_alpha,
+                lora_dropout=finetuning_args.lora_dropout,
+                target_modules=finetuning_args.lora_target
+            )
+            model = get_peft_model(model, lora_config)
+
+    if model_args.checkpoint_dir is not None:
+        logger.info("Loaded fine-tuned model from checkpoint(s): {}".format(",".join(model_args.checkpoint_dir)))
+
+    return model
+
+
+def load_pretrained(
+        model_args: ModelArguments,
+        finetuning_args: FinetuningArguments,
+        is_trainable: Optional[bool] = False,
+        stage: Optional[Literal["pt", "sft", "rm", "ppo"]] = "sft"
+) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+    r"""
+    Loads pretrained model and tokenizer.
+
+    Support both training and inference.
+    """
+    if (not is_trainable) and model_args.checkpoint_dir is None:
+        logger.warning("Checkpoint is not found at evaluation, load the original model.")
+        finetuning_args = FinetuningArguments(finetuning_type="none")
+
+    assert stage in ["pt", "sft"] or finetuning_args.finetuning_type == "lora", \
+        "RM and PPO training can only be performed with the LoRA method."
+
+    config_kwargs = {
+        "trust_remote_code": True,
+        "cache_dir": model_args.cache_dir,
+        "revision": model_args.model_revision,
+        "use_auth_token": True if model_args.use_auth_token else None,
+    }
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        use_fast=model_args.use_fast_tokenizer,
+        padding_side="left",
+        **config_kwargs
+    )
+    tokenizer.pad_token_id = 0 if tokenizer.pad_token_id is None else tokenizer.pad_token_id # set as the <unk> token
+    tokenizer.pad_token_id = 0 if tokenizer.pad_token_id == 64000 else tokenizer.pad_token_id # for baichuan model (older version)
+
+    config = AutoConfig.from_pretrained(model_args.model_name_or_path, **config_kwargs)
+    is_mergeable = True
+
+    # Quantization configurations (using bitsandbytes library).
+    if model_args.quantization_bit is not None:
+        if model_args.quantization_bit == 8:
+            require_version("bitsandbytes>=0.37.0", "To fix: pip install bitsandbytes>=0.37.0")
+            config_kwargs["load_in_8bit"] = True
+            config_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_8bit=True,
+                llm_int8_threshold=6.0
+            )
+        elif model_args.quantization_bit == 4:
+            require_version("bitsandbytes>=0.39.0", "To fix: pip install bitsandbytes>=0.39.0")
+            require_version("transformers>=4.30.1", "To fix: pip install transformers>=4.30.1")
+            require_version("accelerate>=0.20.3", "To fix: pip install accelerate>=0.20.3")
+            require_version("peft>=0.4.0.dev0", "To fix: pip install git+https://github.com/huggingface/peft.git")
+            config_kwargs["load_in_4bit"] = True
+            config_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=model_args.compute_dtype,
+                bnb_4bit_use_double_quant=model_args.double_quantization,
+                bnb_4bit_quant_type=model_args.quantization_type
+            )
+        is_mergeable = False
+        config_kwargs["device_map"] = {"": int(os.environ.get("LOCAL_RANK", "0"))}
+        logger.info("Quantizing model to {} bit.".format(model_args.quantization_bit))
+
+    if not is_trainable: # `device_map=auto` should be used for inference only
+        config_kwargs["device_map"] = "auto"
+
+    # Load and prepare pretrained models (without valuehead).
+    model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        config=config,
+        torch_dtype=torch.bfloat16 if model_args.compute_dtype == torch.bfloat16 else torch.float16,
+        low_cpu_mem_usage=True,
+        **config_kwargs
+    )
+    model = prepare_model_for_training(model, finetuning_args.finetuning_type) if is_trainable else model
+    model = _init_adapter(model, model_args, finetuning_args, is_trainable, is_mergeable)
+
+    if stage == "rm" or stage == "ppo": # add value head
+        model = AutoModelForCausalLMWithValueHead.from_pretrained(model)
+
+        if stage == "rm" and model_args.checkpoint_dir is not None: # load valuehead weights to evaluate reward model
+            logger.warning("Only the last checkpoint containing valuehead will be loaded as the valuehead.")
+            if load_valuehead_params(model, model_args.checkpoint_dir[-1]):
+                model.v_head.load_state_dict({
+                    "summary.weight": getattr(model, "reward_head_weight"),
+                    "summary.bias": getattr(model, "reward_head_bias")
+                })
+
+        if stage == "ppo": # load reward model
+            assert is_trainable, "PPO stage cannot be performed at evaluation."
+            assert model_args.reward_model is not None, "Reward model is necessary for PPO training."
+            logger.info("Load reward model from {}".format(model_args.reward_model))
+            model.pretrained_model.load_adapter(model_args.reward_model, "reward", is_trainable=False)
+            assert load_valuehead_params(model, model_args.reward_model), "Reward model is not correctly loaded."
+
+    if not is_trainable:
+        model.requires_grad_(False) # fix all model params
+        model = model.half() if model_args.quantization_bit is None else model # cast from fp32 to fp16
+
+    print_trainable_params(model)
+
+    return model, tokenizer
+
+
+def prepare_args(
+        stage: Literal["pt", "sft", "rm", "ppo"]
+) -> Tuple[ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, FinetuningArguments]:
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, FinetuningArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): # Provide arguments with a json file.
+        model_args, data_args, training_args, finetuning_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args, finetuning_args = parser.parse_args_into_dataclasses()
+
+    # Setup logging
+    if training_args.should_log:
+        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
+        transformers.utils.logging.set_verbosity_info()
+
+    log_level = training_args.get_process_log_level()
+    datasets.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.enable_default_handler()
+    transformers.utils.logging.enable_explicit_format()
+
+    # Check arguments (do not check finetuning_args since it may be loaded from checkpoints)
+    if stage != "sft" and training_args.predict_with_generate:
+        raise ValueError("`predict_with_generate` cannot be set as True at PT, RM and PPO stages.")
+
+    if training_args.do_train and training_args.predict_with_generate:
+        raise ValueError("`predict_with_generate` cannot be set as True while training.")
+
+    if training_args.do_predict and (not training_args.predict_with_generate):
+        raise ValueError("Please enable `predict_with_generate` to save model predictions.")
+
+    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
+        raise ValueError("Quantization is only compatible with the LoRA method.")
+
+    if model_args.quantization_bit is not None and (not training_args.do_train):
+        logger.warning("Evaluating model in 4/8-bit mode may cause lower scores.")
+
+    if training_args.do_train and (not training_args.fp16):
+        logger.warning("We recommend enable fp16 mixed precision training.")
+
+    if data_args.prompt_template == "alpaca":
+        logger.warning("Please specify `prompt_template` if you are using other pre-trained models.")
+
+    if training_args.local_rank != -1 and training_args.ddp_find_unused_parameters is None:
+        logger.warning("`ddp_find_unused_parameters` needs to be set as False in DDP training.")
+        training_args.ddp_find_unused_parameters = False
+
+    training_args.optim = "adamw_torch" if training_args.optim == "adamw_hf" else training_args.optim # suppress warning
+
+    if model_args.quantization_bit is not None:
+        if training_args.fp16:
+            model_args.compute_dtype = torch.float16
+        elif training_args.bf16:
+            model_args.compute_dtype = torch.bfloat16
+        else:
+            model_args.compute_dtype = torch.float32
+
+    # Log on each process the small summary:
+    logger.info(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}\n"
+        + f"  distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+    logger.info(f"Training/evaluation parameters {training_args}")
+
+    # Set seed before initializing model.
+    transformers.set_seed(training_args.seed)
+
+    return model_args, data_args, training_args, finetuning_args
+
+
+def prepare_infer_args() -> Tuple[ModelArguments, DataTrainingArguments, FinetuningArguments, GeneratingArguments]:
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, FinetuningArguments, GeneratingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): # Provide arguments with a json file.
+        model_args, data_args, finetuning_args, generating_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, finetuning_args, generating_args = parser.parse_args_into_dataclasses()
+
+    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
+        raise ValueError("Quantization is only compatible with the LoRA method.")
+
+    if data_args.prompt_template == "alpaca":
+        logger.warning("Please specify `prompt_template` if you are using other pre-trained models.")
+
+    return model_args, data_args, finetuning_args, generating_args
+
+
+def prepare_data(
+        model_args: ModelArguments,
+        data_args: DataTrainingArguments
+) -> Dataset:
+
+    def checksum(file_path, hash):
+        with open(file_path, "rb") as datafile:
+            binary_data = datafile.read()
+        sha1 = hashlib.sha1(binary_data).hexdigest()
+        if sha1 != hash:
+            logger.warning("Checksum failed for {}. It may vary depending on the platform.".format(file_path))
+
+    max_samples = data_args.max_samples
+    all_datasets: List[Dataset] = [] # support multiple datasets
+
+    for dataset_attr in data_args.dataset_list:
+
+        logger.info("Loading dataset {}...".format(dataset_attr))
+
+        if dataset_attr.load_from == "hf_hub":
+            raw_datasets = load_dataset(dataset_attr.dataset_name, cache_dir=model_args.cache_dir)
+        elif dataset_attr.load_from == "script":
+            raw_datasets = load_dataset(
+                os.path.join(data_args.dataset_dir, dataset_attr.dataset_name),
+                cache_dir=model_args.cache_dir
+            )
+        elif dataset_attr.load_from == "file":
+            data_file = os.path.join(data_args.dataset_dir, dataset_attr.file_name)
+
+            extension = dataset_attr.file_name.split(".")[-1]
+            if extension == "csv":
+                file_type = "csv"
+            elif extension == "json" or extension == "jsonl":
+                file_type = "json"
+            else:
+                file_type = "text"
+
+            if dataset_attr.file_sha1 is not None:
+                checksum(data_file, dataset_attr.file_sha1)
+            else:
+                logger.warning("Checksum failed: missing SHA-1 hash value in dataset_info.json.")
+
+            raw_datasets = load_dataset(
+                file_type,
+                data_files=data_file,
+                cache_dir=model_args.cache_dir,
+                use_auth_token=True if model_args.use_auth_token else None
+            )
+        else:
+            raise NotImplementedError
+
+        dataset = raw_datasets[data_args.split]
+
+        if max_samples is not None:
+            max_samples_temp = min(len(dataset), max_samples)
+            dataset = dataset.select(range(max_samples_temp))
+
+        dummy_data = [None] * len(dataset)
+        for column_name, target_name in [
+            ("prompt_column", "prompt"),
+            ("query_column", "query"),
+            ("response_column", "response"),
+            ("history_column", "history")
+        ]: # every dataset will have 4 columns same as each other
+            if getattr(dataset_attr, column_name) != target_name:
+                if getattr(dataset_attr, column_name):
+                    dataset = dataset.rename_column(getattr(dataset_attr, column_name), target_name)
+                else: # None or empty string
+                    dataset = dataset.add_column(target_name, dummy_data)
+        all_datasets.append(dataset)
+
+    if len(data_args.dataset_list) == 1:
+        all_datasets = all_datasets[0]
+    else:
+        all_datasets = concatenate_datasets(all_datasets)
+
+    return all_datasets
+
+
+def preprocess_data(
+        dataset: Dataset,
+        tokenizer: PreTrainedTokenizer,
+        data_args: DataTrainingArguments,
+        training_args: Seq2SeqTrainingArguments,
+        stage: Literal["pt", "sft", "rm", "ppo"]
+) -> Dataset:
+
+    column_names = list(dataset.column_names)
+    prefix = data_args.source_prefix if data_args.source_prefix is not None else ""
+    prompt_template = Template(data_args.prompt_template)
+
+    # support question with a single answer or multiple answers
+    def get_dialog(examples):
+        for i in range(len(examples["prompt"])):
+            if examples["prompt"][i] and examples["response"][i]:
+                query, answer = examples["prompt"][i], examples["response"][i]
+                query = query + "\n" + examples["query"][i] if examples["query"][i] else query
+                dialog = prompt_template.get_dialog(query, answer, examples["history"][i], prefix)
+                yield dialog
+
+    def preprocess_pretrain_dataset(examples):
+        # build grouped texts with format `[BOS] X1 X2 X3 ...` (without [EOS])
+        text_ids = tokenizer(examples["prompt"], add_special_tokens=False)["input_ids"]
+        concatenated_ids = list(chain(*text_ids))
+        total_length = len(concatenated_ids)
+        block_size = data_args.max_source_length - 1
+        # we drop the small remainder, and if the total_length < block_size, we exclude this batch
+        total_length = (total_length // block_size) * block_size
+        # split by chunks of max_source_length
+        result = [[tokenizer.bos_token_id] + concatenated_ids[i: i + block_size]
+                  for i in range(0, total_length, block_size)]
+        return {
+            "input_ids": result,
+            "labels": result.copy()
+        }
+
+    def preprocess_supervised_dataset(examples):
+        # build inputs with format `X [BOS] Y [EOS]` and labels with format `[IGNORE] ... [IGNORE] Y [EOS]`
+        # for input with history, we build multiple input-label pairs just like:
+        # https://github.com/lm-sys/FastChat/blob/f17c092f64840fa6354ed52789dccb2daa793d0b/fastchat/train/train.py#L112
+        model_inputs = {"input_ids": [], "labels": []}
+        for dialog in get_dialog(examples):
+            input_ids, labels = [], []
+
+            for i in range(len(dialog) // 2):
+                source_ids = tokenizer.encode(text=dialog[2*i], add_special_tokens=False)
+                target_ids = tokenizer.encode(text=dialog[2*i+1], add_special_tokens=False)
+                input_ids += source_ids + [tokenizer.bos_token_id] + target_ids + [tokenizer.eos_token_id]
+                labels += [IGNORE_INDEX] * (len(source_ids) + 1) + target_ids + [tokenizer.eos_token_id]
+
+            model_inputs["input_ids"].append(input_ids[:data_args.max_source_length + data_args.max_target_length])
+            model_inputs["labels"].append(labels[:data_args.max_source_length + data_args.max_target_length])
+        return model_inputs
+
+    def preprocess_unsupervised_dataset(examples):
+        # build inputs with format `X [BOS]` and labels with format `Y [BOS]`
+        model_inputs = {"input_ids": [], "labels": []}
+        for dialog in get_dialog(examples):
+            prompt, answer = "".join(dialog[:-1]), dialog[-1]
+
+            source_ids = tokenizer.encode(text=prompt, add_special_tokens=False)
+            target_ids = tokenizer.encode(text=answer, add_special_tokens=False)
+
+            if len(source_ids) > data_args.max_source_length - 1: # bos token
+                source_ids = source_ids[:data_args.max_source_length - 1]
+            if len(target_ids) > data_args.max_target_length - 1: # bos token
+                target_ids = target_ids[:data_args.max_target_length - 1]
+
+            input_ids = source_ids + [tokenizer.bos_token_id]
+            labels = target_ids + [tokenizer.bos_token_id]
+
+            model_inputs["input_ids"].append(input_ids)
+            model_inputs["labels"].append(labels)
+        return model_inputs
+
+    def preprocess_pairwise_dataset(examples):
+        # build input pairs with format `X [BOS] Y1 [EOS]` and `X [BOS] Y2 [EOS]`
+        model_inputs = {"accept_ids": [], "reject_ids": []}
+        for dialog in get_dialog(examples):
+            prompt, answer = "".join(dialog[:-1]), dialog[-1]
+
+            source_ids = tokenizer.encode(text=prompt, add_special_tokens=False)
+            accept_ids = tokenizer.encode(text=answer[0], add_special_tokens=False)
+            reject_ids = tokenizer.encode(text=answer[1], add_special_tokens=False)
+
+            if len(source_ids) > data_args.max_source_length - 1: # bos token
+                source_ids = source_ids[:data_args.max_source_length - 1]
+            if len(accept_ids) > data_args.max_target_length - 1: # eos token
+                accept_ids = accept_ids[:data_args.max_target_length - 1]
+            if len(reject_ids) > data_args.max_target_length - 1: # eos token
+                reject_ids = reject_ids[:data_args.max_target_length - 1]
+
+            accept_ids = source_ids + [tokenizer.bos_token_id] + accept_ids + [tokenizer.eos_token_id]
+            reject_ids = source_ids + [tokenizer.bos_token_id] + reject_ids + [tokenizer.eos_token_id]
+
+            model_inputs["accept_ids"].append(accept_ids)
+            model_inputs["reject_ids"].append(reject_ids)
+        return model_inputs
+
+    def print_supervised_dataset_example(example):
+        print("input_ids:\n{}".format(example["input_ids"]))
+        print("inputs:\n{}".format(tokenizer.decode(example["input_ids"])))
+        print("label_ids:\n{}".format(example["labels"]))
+        print("labels:\n{}".format(
+            tokenizer.decode([d if d != IGNORE_INDEX else tokenizer.pad_token_id for d in example["labels"]]))
+        )
+
+    def print_pairwise_dataset_example(example):
+        print("accept_ids:\n{}".format(example["accept_ids"]))
+        print("accepts:\n{}".format(tokenizer.decode(example["accept_ids"])))
+        print("reject_ids:\n{}".format(example["reject_ids"]))
+        print("rejects:\n{}".format(tokenizer.decode(example["reject_ids"])))
+
+    def print_unsupervised_dataset_example(example):
+        print("input_ids:\n{}".format(example["input_ids"]))
+        print("inputs:\n{}".format(tokenizer.decode(example["input_ids"])))
+
+    if stage == "pt":
+        preprocess_function = preprocess_pretrain_dataset
+    elif stage == "sft":
+        preprocess_function = preprocess_unsupervised_dataset \
+            if training_args.predict_with_generate else preprocess_supervised_dataset
+    elif stage == "rm":
+        preprocess_function = preprocess_pairwise_dataset
+    elif stage == "ppo":
+        preprocess_function = preprocess_unsupervised_dataset
+
+    with training_args.main_process_first(desc="dataset map pre-processing"):
+        dataset = dataset.map(
+            preprocess_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not data_args.overwrite_cache,
+            desc="Running tokenizer on dataset"
+        )
+
+        if stage == "pt":
+            print_unsupervised_dataset_example(dataset[0])
+        elif stage == "sft":
+            print_supervised_dataset_example(dataset[0])
+        elif stage == "rm":
+            print_pairwise_dataset_example(dataset[0])
+        elif stage == "ppo":
+            print_unsupervised_dataset_example(dataset[0])
+
+        return dataset

src/utils/config.py

@@ -0,0 +1,283 @@
+import os
+import json
+import torch
+from typing import Any, Dict, List, Literal, Optional
+from dataclasses import asdict, dataclass, field
+
+
+@dataclass
+class DatasetAttr:
+
+    load_from: str
+    dataset_name: Optional[str] = None
+    file_name: Optional[str] = None
+    file_sha1: Optional[str] = None
+
+    def __repr__(self) -> str:
+        if self.dataset_name is not None:
+            return self.dataset_name
+        else:
+            return self.file_name
+
+    def __post_init__(self):
+        self.prompt_column = "instruction"
+        self.query_column = "input"
+        self.response_column = "output"
+        self.history_column = None
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune.
+    """
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models."}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co."}
+    )
+    use_fast_tokenizer: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}
+    )
+    use_auth_token: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Will use the token generated when running `huggingface-cli login`."}
+    )
+    model_revision: Optional[str] = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}
+    )
+    quantization_bit: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of bits to quantize the model."}
+    )
+    quantization_type: Optional[Literal["fp4", "nf4"]] = field(
+        default="nf4",
+        metadata={"help": "Quantization data type to use in int4 training."}
+    )
+    double_quantization: Optional[bool] = field(
+        default=True,
+        metadata={"help": "Whether to use double quantization in int4 training or not."}
+    )
+    compute_dtype: Optional[torch.dtype] = field(
+        default=None,
+        metadata={"help": "Used in quantization configs. Do not specify this argument manually."}
+    )
+    checkpoint_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to the directory(s) containing the delta model checkpoints as well as the configurations."}
+    )
+    reward_model: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to the directory containing the checkpoints of the reward model."}
+    )
+    resume_lora_training: Optional[bool] = field(
+        default=True,
+        metadata={"help": "Whether to resume training from the last LoRA weights or create new weights after merging them."}
+    )
+    plot_loss: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to plot the training loss after fine-tuning or not."}
+    )
+
+    def __post_init__(self):
+        if self.checkpoint_dir is not None: # support merging multiple lora weights
+            self.checkpoint_dir = [cd.strip() for cd in self.checkpoint_dir.split(",")]
+
+        if self.quantization_bit is not None:
+            assert self.quantization_bit in [4, 8], "We only accept 4-bit or 8-bit quantization."
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and evaluation.
+    """
+    dataset: Optional[str] = field(
+        default="alpaca_zh",
+        metadata={"help": "The name of provided dataset(s) to use. Use comma to separate multiple datasets."}
+    )
+    dataset_dir: Optional[str] = field(
+        default="data",
+        metadata={"help": "The name of the folder containing datasets."}
+    )
+    split: Optional[str] = field(
+        default="train",
+        metadata={"help": "Which dataset split to use for training and evaluation."}
+    )
+    overwrite_cache: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets."}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."}
+    )
+    max_source_length: Optional[int] = field(
+        default=512,
+        metadata={"help": "The maximum total input sequence length after tokenization."}
+    )
+    max_target_length: Optional[int] = field(
+        default=512,
+        metadata={"help": "The maximum total output sequence length after tokenization."}
+    )
+    max_samples: Optional[int] = field(
+        default=None,
+        metadata={"help": "For debugging purposes, truncate the number of examples for each dataset."}
+    )
+    eval_num_beams: Optional[int] = field(
+        default=None,
+        metadata={"help": "Number of beams to use for evaluation. This argument will be passed to `model.generate`"}
+    )
+    ignore_pad_token_for_loss: Optional[bool] = field(
+        default=True,
+        metadata={"help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."}
+    )
+    source_prefix: Optional[str] = field(
+        default=None,
+        metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
+    )
+    dev_ratio: Optional[float] = field(
+        default=0,
+        metadata={"help": "Proportion of the dataset to include in the development set, should be between 0.0 and 1.0."}
+    )
+    prompt_template: Optional[str] = field(
+        default="alpaca",
+        metadata={"help": "Which template to use for constructing prompts in training and inference."}
+    )
+
+    def __post_init__(self): # support mixing multiple datasets
+        dataset_names = [ds.strip() for ds in self.dataset.split(",")]
+        with open(os.path.join(self.dataset_dir, "dataset_info.json"), "r") as f:
+            dataset_info = json.load(f)
+
+        self.dataset_list: List[DatasetAttr] = []
+        for name in dataset_names:
+            if name not in dataset_info:
+                raise ValueError("Undefined dataset {} in dataset_info.json.".format(name))
+
+            if "hf_hub_url" in dataset_info[name]:
+                dataset_attr = DatasetAttr("hf_hub", dataset_name=dataset_info[name]["hf_hub_url"])
+            elif "script_url" in dataset_info[name]:
+                dataset_attr = DatasetAttr("script", dataset_name=dataset_info[name]["script_url"])
+            else:
+                dataset_attr = DatasetAttr(
+                    "file",
+                    file_name=dataset_info[name]["file_name"],
+                    file_sha1=dataset_info[name].get("file_sha1", None)
+                )
+
+            if "columns" in dataset_info[name]:
+                dataset_attr.prompt_column = dataset_info[name]["columns"].get("prompt", None)
+                dataset_attr.query_column = dataset_info[name]["columns"].get("query", None)
+                dataset_attr.response_column = dataset_info[name]["columns"].get("response", None)
+                dataset_attr.history_column = dataset_info[name]["columns"].get("history", None)
+
+            self.dataset_list.append(dataset_attr)
+
+
+@dataclass
+class FinetuningArguments:
+    """
+    Arguments pertaining to which techniques we are going to fine-tuning with.
+    """
+    finetuning_type: Optional[Literal["none", "freeze", "lora", "full"]] = field(
+        default="lora",
+        metadata={"help": "Which fine-tuning method to use."}
+    )
+    num_layer_trainable: Optional[int] = field(
+        default=3,
+        metadata={"help": "Number of trainable layers for Freeze fine-tuning."}
+    )
+    name_module_trainable: Optional[Literal["mlp", "self_attn", "self_attention"]] = field(
+        default="mlp",
+        metadata={"help": "Name of trainable modules for Freeze fine-tuning. \
+                  LLaMA choices: [\"mlp\", \"self_attn\"], \
+                  BLOOM choices: [\"mlp\", \"self_attention\"], \
+                  Baichuan choices: [\"mlp\", \"self_attn\"]"}
+    )
+    lora_rank: Optional[int] = field(
+        default=8,
+        metadata={"help": "The intrinsic dimension for LoRA fine-tuning."}
+    )
+    lora_alpha: Optional[float] = field(
+        default=32.0,
+        metadata={"help": "The scale factor for LoRA fine-tuning (similar with the learning rate)."}
+    )
+    lora_dropout: Optional[float] = field(
+        default=0.1,
+        metadata={"help": "Dropout rate for the LoRA fine-tuning."}
+    )
+    lora_target: Optional[str] = field(
+        default="q_proj,v_proj",
+        metadata={"help": "Name(s) of target modules to apply LoRA. Use comma to separate multiple modules. \
+                  LLaMA choices: [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
+                  BLOOM choices: [\"query_key_value\", \"self_attention.dense\", \"mlp.dense\"], \
+                  Baichuan choices: [\"W_pack\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"]"}
+    )
+
+    def __post_init__(self):
+        if isinstance(self.lora_target, str): # support custom target modules/layers of LoRA
+            self.lora_target = [target.strip() for target in self.lora_target.split(",")]
+
+        if self.num_layer_trainable > 0: # fine-tuning the last n layers if num_layer_trainable > 0
+            trainable_layer_ids = [27 - k for k in range(self.num_layer_trainable)]
+        else: # fine-tuning the first n layers if num_layer_trainable < 0
+            trainable_layer_ids = [k for k in range(-self.num_layer_trainable)]
+
+        self.trainable_layers = ["layers.{:d}.{}".format(idx, self.name_module_trainable) for idx in trainable_layer_ids]
+
+        assert self.finetuning_type in ["none", "freeze", "lora", "full"], "Invalid fine-tuning method."
+
+    def save_to_json(self, json_path: str):
+        """Saves the content of this instance in JSON format inside `json_path`."""
+        json_string = json.dumps(asdict(self), indent=2, sort_keys=True) + "\n"
+        with open(json_path, "w", encoding="utf-8") as f:
+            f.write(json_string)
+
+    @classmethod
+    def load_from_json(cls, json_path: str):
+        """Creates an instance from the content of `json_path`."""
+        with open(json_path, "r", encoding="utf-8") as f:
+            text = f.read()
+        return cls(**json.loads(text))
+
+
+@dataclass
+class GeneratingArguments:
+    """
+    Arguments pertaining to specify the decoding parameters.
+    """
+    do_sample: Optional[bool] = field(
+        default=True,
+        metadata={"help": "Whether or not to use sampling, use greedy decoding otherwise."}
+    )
+    temperature: Optional[float] = field(
+        default=0.95,
+        metadata={"help": "The value used to modulate the next token probabilities."}
+    )
+    top_p: Optional[float] = field(
+        default=0.7,
+        metadata={"help": "The smallest set of most probable tokens with probabilities that add up to top_p or higher are kept."}
+    )
+    top_k: Optional[int] = field(
+        default=50,
+        metadata={"help": "The number of highest probability vocabulary tokens to keep for top-k filtering."}
+    )
+    num_beams: Optional[int] = field(
+        default=1,
+        metadata={"help": "Number of beams for beam search. 1 means no beam search."}
+    )
+    max_new_tokens: Optional[int] = field(
+        default=512,
+        metadata={"help": "The maximum numbers of tokens to generate, ignoring the number of tokens in the prompt."}
+    )
+    repetition_penalty: Optional[float] = field(
+        default=1.0,
+        metadata={"help": "The parameter for repetition penalty. 1.0 means no penalty."}
+    )
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)

src/utils/data_collator.py

@@ -0,0 +1,64 @@
+import torch
+
+from typing import Dict, Optional, Sequence, Union
+
+from transformers import DataCollatorWithPadding, BatchEncoding
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+from .other import IGNORE_INDEX
+
+
+class DynamicDataCollatorWithPadding(DataCollatorWithPadding):
+    r"""
+    Inherits DataCollatorWithPadding. It is capable of dynamically padding for batched data.
+    """
+    def __init__(
+            self,
+            tokenizer: PreTrainedTokenizer,
+            ignore_pad_token_for_loss: Optional[bool] = False
+    ):
+        super().__init__(tokenizer, padding=True)
+        self.label_pad_token_id = IGNORE_INDEX if ignore_pad_token_for_loss else tokenizer.pad_token_id
+
+    def get_attention_masks(self, input_ids: torch.Tensor, device: torch.device) -> torch.Tensor:
+        r"""
+        Generates attention masks for left-padded sequences.
+        """
+        batch_size, seq_length = input_ids.size()
+        attention_mask = torch.ones((batch_size, seq_length), device=device)
+        for i, seq in enumerate(input_ids):
+            attention_mask[i, :(seq != self.tokenizer.pad_token_id).nonzero()[0].item()] = 0 # padding
+        attention_mask = attention_mask.bool()
+        return attention_mask
+
+    def __call__(self, features: Sequence[Dict[str, Union[torch.Tensor, Sequence[int]]]]) -> BatchEncoding:
+        r"""
+        Pads batched data to the longest sequence in the batch.
+
+        We adopt left-padding in both training and evaluation.
+        """
+        if isinstance(features[0]["input_ids"], torch.Tensor):
+            input_ids = [feature["input_ids"].clone().detach().flip(0) for feature in features]
+        else:
+            input_ids = [torch.tensor(feature["input_ids"]).flip(0) for feature in features]
+
+        if "labels" in features[0]:
+            if isinstance(features[0]["labels"], torch.Tensor):
+                labels = [feature["labels"].clone().detach().flip(0) for feature in features]
+            else:
+                labels = [torch.tensor(feature["labels"]).flip(0) for feature in features]
+            input_ids = input_ids + labels # pad them to the same length
+
+        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id).flip(-1)
+
+        batch = {}
+
+        if "labels" in features[0]:
+            input_ids, labels = input_ids.split(len(features), dim=0)
+            labels = torch.where(labels != self.tokenizer.pad_token_id, labels, self.label_pad_token_id)
+            batch["labels"] = labels
+
+        batch["input_ids"] = input_ids
+        batch["attention_mask"] = self.get_attention_masks(input_ids, device=input_ids.device)
+
+        return BatchEncoding(batch)

src/utils/other.py

@@ -0,0 +1,196 @@
+import os
+import sys
+import json
+import torch
+import logging
+from typing import Dict, List, Optional
+
+from transformers.trainer import TRAINER_STATE_NAME
+from transformers.modeling_utils import PreTrainedModel
+from transformers.generation.utils import LogitsProcessorList
+from transformers.generation.logits_process import LogitsProcessor
+
+from peft.utils import WEIGHTS_NAME
+
+
+IGNORE_INDEX = -100
+VALUE_HEAD_FILE_NAME = "value_head.bin"
+FINETUNING_ARGS_NAME = "finetuning_args.json"
+
+
+def get_logger(name: str) -> logging.Logger:
+    return logging.getLogger(name)
+
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.INFO,
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+
+
+logger = get_logger(__name__)
+
+
+class AverageMeter:
+    r"""
+    Computes and stores the average and current value.
+    """
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+# Avoid runtime error in model.generate(do_sample=True).
+class InvalidScoreLogitsProcessor(LogitsProcessor):
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        if torch.isnan(scores).any() or torch.isinf(scores).any():
+            scores.zero_()
+            scores[..., 0] = 1.0
+        return scores
+
+
+def get_logits_processor() -> LogitsProcessorList:
+    logits_processor = LogitsProcessorList()
+    logits_processor.append(InvalidScoreLogitsProcessor())
+    return logits_processor
+
+
+# Includes: (1) cast the layernorm in fp32 (2) make output embedding layer require grads (3) upcast the lm_head to fp32
+# Inspired by: https://github.com/huggingface/peft/blob/c0209c35abbf88c63aa267800d98a8e212ed0a42/src/peft/utils/other.py#L35
+def prepare_model_for_training(
+        model: PreTrainedModel,
+        finetuning_type: str,
+        output_embedding_layer_name: Optional[str] = "lm_head",
+        use_gradient_checkpointing: Optional[bool] = True,
+        layer_norm_names: Optional[List[str]] = ["norm", "ln_f"] # for LLaMA and BLOOM setting
+) -> PreTrainedModel:
+
+    for name, param in model.named_parameters():
+        if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
+            param.data = param.data.to(torch.float32)
+
+    if use_gradient_checkpointing:
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False # turn off when gradient checkpointing is enabled
+
+    if finetuning_type != "full" and hasattr(model, output_embedding_layer_name):
+        output_embedding_layer: torch.nn.Linear = getattr(model, output_embedding_layer_name)
+        input_dtype = output_embedding_layer.weight.dtype
+
+        class CastOutputToFloat(torch.nn.Sequential):
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return super().forward(x.to(input_dtype)).to(torch.float32)
+
+        setattr(model, output_embedding_layer_name, CastOutputToFloat(output_embedding_layer))
+
+    return model
+
+
+def print_trainable_params(model: torch.nn.Module) -> None:
+    trainable_params, all_param = 0, 0
+    for param in model.parameters():
+        num_params = param.numel()
+        # if using DS Zero 3 and the weights are initialized empty
+        if num_params == 0 and hasattr(param, "ds_numel"):
+            num_params = param.ds_numel
+        all_param += num_params
+        if param.requires_grad:
+            trainable_params += num_params
+    print("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
+                trainable_params, all_param, 100 * trainable_params / all_param))
+
+
+def get_state_dict(model: torch.nn.Module) -> Dict[str, torch.Tensor]: # get state dict containing trainable parameters
+    state_dict = model.state_dict()
+    filtered_state_dict = {}
+
+    for k, v in model.named_parameters():
+        if v.requires_grad:
+            filtered_state_dict[k] = state_dict[k].cpu().clone().detach()
+
+    return filtered_state_dict
+
+
+def load_trainable_params(model: torch.nn.Module, checkpoint_dir: os.PathLike) -> bool:
+    weights_file = os.path.join(checkpoint_dir, WEIGHTS_NAME)
+    if not os.path.exists(weights_file):
+        logger.warning("Provided path ({}) does not contain pre-trained weights.".format(checkpoint_dir))
+        return False
+    model_state_dict = torch.load(weights_file, map_location="cpu")
+    model.load_state_dict(model_state_dict, strict=False) # skip missing keys
+    return True
+
+
+def load_valuehead_params(model: torch.nn.Module, checkpoint_dir: os.PathLike) -> bool:
+    valuehead_file = os.path.join(checkpoint_dir, VALUE_HEAD_FILE_NAME)
+    if not os.path.exists(valuehead_file):
+        logger.warning("Provided path ({}) does not contain valuehead weights.".format(checkpoint_dir))
+        return False
+    valuehead_state_dict = torch.load(valuehead_file, map_location="cpu")
+    model.register_buffer("reward_head_weight", valuehead_state_dict["summary.weight"])
+    model.register_buffer("reward_head_bias", valuehead_state_dict["summary.bias"])
+    model.register_buffer("default_head_weight", torch.zeros_like(valuehead_state_dict["summary.weight"]))
+    model.register_buffer("default_head_bias", torch.zeros_like(valuehead_state_dict["summary.bias"]))
+    return True
+
+
+def smooth(scalars: List[float], weight: Optional[float] = 0.9) -> List[float]:
+    r"""
+    EMA implementation according to TensorBoard.
+    """
+    last = scalars[0]
+    smoothed = list()
+    for next_val in scalars:
+        smoothed_val = last * weight + (1 - weight) * next_val
+        smoothed.append(smoothed_val)
+        last = smoothed_val
+    return smoothed
+
+
+def plot_loss(save_dictionary: os.PathLike, keys: Optional[List[str]] = ["loss"]) -> None:
+    import matplotlib.pyplot as plt
+    with open(os.path.join(save_dictionary, TRAINER_STATE_NAME), "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    for key in keys:
+        steps, metrics = [], []
+        for i in range(len(data["log_history"])):
+            if key in data["log_history"][i]:
+                steps.append(data["log_history"][i]["step"])
+                metrics.append(data["log_history"][i][key])
+
+        if len(metrics) == 0:
+            logger.warning(f"No metric {key} to plot.")
+            continue
+
+        plt.figure()
+        plt.plot(steps, metrics, alpha=0.4, label="original")
+        plt.plot(steps, smooth(metrics), label="smoothed")
+        plt.title("training {} of {}".format(key, save_dictionary))
+        plt.xlabel("step")
+        plt.ylabel(key)
+        plt.legend()
+        plt.savefig(os.path.join(save_dictionary, "training_{}.png".format(key)), format="png", dpi=100)
+        print("Figure saved:", os.path.join(save_dictionary, "training_{}.png".format(key)))

src/utils/pairwise.py

@@ -0,0 +1,57 @@
+import torch
+import numpy as np
+from typing import Dict, Sequence, Tuple, Union
+
+from .data_collator import DynamicDataCollatorWithPadding
+
+from .peft_trainer import PeftTrainer
+
+from .other import get_logger
+
+logger = get_logger(__name__)
+
+
+def compute_accuracy(eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
+    preds, _ = eval_preds
+    preds = np.array(preds)
+    return {"accuracy": (preds[:, 0] > preds[:, 1]).sum() / len(preds)}
+
+
+class PairwiseDataCollatorWithPadding(DynamicDataCollatorWithPadding):
+    r"""
+    Data collator for pairwise data.
+    """
+
+    def __call__(self, features: Sequence[Dict[str, Union[torch.Tensor, Sequence[int]]]]) -> Dict[str, torch.Tensor]:
+        r"""
+        Pads batched data to the longest sequence in the batch.
+
+        We generate 2 * n examples where the first n examples represent chosen examples and
+        the last n examples represent rejected examples.
+        """
+        features = [{"input_ids": feature[key]} for key in ("accept_ids", "reject_ids") for feature in features]
+        return super().__call__(features)
+
+
+class PairwisePeftTrainer(PeftTrainer):
+    r"""
+    Inherits PeftTrainer to compute pairwise loss.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.can_return_loss = True # override property to return eval_loss
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        r"""
+        Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.
+
+        We use score on the EOS token to represent reward of the whole sentence.
+
+        Subclass and override to inject custom behavior. It should not be directly used by external scripts.
+        """
+        batch_size = inputs["input_ids"].size(0) // 2
+        _, _, values = model(**inputs)
+        r_accept, r_reject = values[:, -1].split(batch_size, dim=0)
+        loss = -torch.log(torch.sigmoid(r_accept - r_reject)).mean()
+        return (loss, torch.stack((r_accept, r_reject), dim=-1)) if return_outputs else loss

src/utils/peft_trainer.py

@@ -0,0 +1,132 @@
+import os
+import json
+import time
+import torch
+from typing import Dict, Optional
+from datetime import timedelta
+
+from transformers import (
+    Seq2SeqTrainer,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments
+)
+
+from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.modeling_utils import unwrap_model
+
+from peft.utils.other import WEIGHTS_NAME
+
+from .config import FinetuningArguments
+
+from .other import (
+    get_logger,
+    get_state_dict,
+    load_trainable_params,
+    load_valuehead_params,
+    FINETUNING_ARGS_NAME,
+    VALUE_HEAD_FILE_NAME
+)
+
+
+logger = get_logger(__name__)
+
+
+class LogCallback(TrainerCallback):
+    r"""
+    TrainerCallback includes the state function during training, for more details refer to the TrainerCallback class.
+    The on_log function primarily collects process parameters during training, such as training loss, learning rate,
+    and training epochs, as well as progress parameters like the current percentage progress and estimated remaining
+    time. Every time a log is triggered, a new record is appended to the file "messages.log" for dynamic visualization
+    purposes.
+    """
+
+    def __init__(self):
+        self.start_time = time.time()
+
+    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs) -> None:
+        r"""
+        Event called after logging the last logs.
+        """
+        if "loss" not in state.log_history[-1]:
+            return
+        cur_time = time.time()
+        cur_steps = state.log_history[-1].get("step")
+        elapsed_time = cur_time - self.start_time
+        avg_time_per_step = elapsed_time / cur_steps if cur_steps != 0 else 0
+        remaining_steps = state.max_steps - cur_steps
+        remaining_time = remaining_steps * avg_time_per_step
+        log_dict = {
+            "current_steps": cur_steps,
+            "total_steps": state.max_steps,
+            "loss": state.log_history[-1].get("loss", None),
+            "reward": state.log_history[-1].get("reward", None),
+            "learning_rate": state.log_history[-1].get("learning_rate", None),
+            "epoch": state.log_history[-1].get("epoch", None),
+            "percentage": round(cur_steps / state.max_steps * 100, 2) if state.max_steps != 0 else 100,
+            "elapsed_time": str(timedelta(seconds=int(elapsed_time))),
+            "remaining_time": str(timedelta(seconds=int(remaining_time)))
+        }
+        os.makedirs(args.output_dir, exist_ok=True)
+        with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a") as f:
+            f.write(json.dumps(log_dict) + "\n")
+
+
+class PeftTrainer(Seq2SeqTrainer):
+    r"""
+    Inherits Seq2SeqTrainer to support parameter-efficient checkpoints.
+    """
+
+    def __init__(self, finetuning_args: FinetuningArguments, **kwargs):
+        super().__init__(**kwargs)
+        self.finetuning_args = finetuning_args
+        if self.is_world_process_zero() and os.path.exists(os.path.join(self.args.output_dir, "trainer_log.jsonl")):
+            logger.warning("Previous log file in this folder will be deleted.")
+            os.remove(os.path.join(self.args.output_dir, "trainer_log.jsonl"))
+
+    def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, torch.Tensor]] = None) -> None:
+        r"""
+        Saves trainable parameters as model checkpoint.
+
+        This function will only be executed at the process zero.
+
+        Subclass and override to inject custom behavior. It should not be directly used by external scripts.
+        """
+        output_dir = output_dir if output_dir is not None else self.args.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+        logger.info(f"Saving model checkpoint to {output_dir}")
+        model = unwrap_model(self.model)
+
+        if hasattr(model, "pretrained_model"): # for models with valuehead
+            backbone_model = getattr(model, "pretrained_model")
+        else:
+            backbone_model = model
+
+        if hasattr(backbone_model, "peft_config"): # peft methods
+            backbone_model.save_pretrained(output_dir, state_dict=get_state_dict(backbone_model)) # save lora weights
+        else:
+            torch.save(get_state_dict(backbone_model), os.path.join(output_dir, WEIGHTS_NAME)) # save trainable weights
+
+        if hasattr(model, "v_head"): # save valuehead weights
+            torch.save(get_state_dict(getattr(model, "v_head")), os.path.join(output_dir, VALUE_HEAD_FILE_NAME))
+
+        with open(os.path.join(output_dir, TRAINING_ARGS_NAME), "w", encoding="utf-8") as f:
+            f.write(self.args.to_json_string() + "\n")
+        self.finetuning_args.save_to_json(os.path.join(output_dir, FINETUNING_ARGS_NAME))
+
+    def _load_best_model(self):
+        r"""
+        Loads trainable parameters from model checkpoint.
+
+        Subclass and override to inject custom behavior. It should not be directly used by external scripts.
+        """
+        logger.info(f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric}).")
+        model = unwrap_model(self.model)
+        if hasattr(model, "peft_config"): # peft methods
+            model.load_adapter(self.state.best_model_checkpoint, getattr(model, "active_adapter"))
+        else:
+            load_trainable_params(model, self.state.best_model_checkpoint)
+
+        if hasattr(model, "v_head"):
+            load_valuehead_params(model, self.state.best_model_checkpoint)

src/utils/ppo.py

@@ -0,0 +1,223 @@
+import os
+import math
+import torch
+from tqdm import tqdm
+from typing import Callable, Dict, List, Literal, Optional, Tuple
+
+from transformers import Seq2SeqTrainingArguments, TrainerState
+from transformers.modeling_utils import PreTrainedModel
+
+from trl import PPOTrainer, AutoModelForCausalLMWithValueHead
+from trl.core import LengthSampler
+
+from .peft_trainer import PeftTrainer, LogCallback
+
+from .config import FinetuningArguments
+
+from .other import (
+    AverageMeter,
+    get_logger,
+    get_logits_processor
+)
+
+
+logger = get_logger(__name__)
+
+
+def replace_model(model: AutoModelForCausalLMWithValueHead, target: Literal["default", "reward"]) -> None:
+    if target == "reward": # save original head temporarily
+        valuehead_state_dict = model.v_head.state_dict()
+
+        setattr(model, "origin_head_weight", valuehead_state_dict["summary.weight"])
+        setattr(model, "origin_head_bias", valuehead_state_dict["summary.bias"])
+
+    model.pretrained_model.set_adapter(target) # set the LoRA adapter to be active
+    model.v_head.load_state_dict({
+        "summary.weight": getattr(model, "{}_head_weight".format(target)),
+        "summary.bias": getattr(model, "{}_head_bias".format(target))
+    })
+
+
+def cast_layernorm_dtype(
+        model: AutoModelForCausalLMWithValueHead,
+        layer_norm_names: List[str] = ["norm", "ln_f"], # for LLaMA and BLOOM setting
+        layer_norm_params: Optional[Dict[str, torch.Tensor]] = None
+) -> Tuple[AutoModelForCausalLMWithValueHead, Dict[str, torch.Tensor]]:
+
+    layer_norm_state_dict = {}
+
+    for name, param in model.named_parameters():
+        if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
+            if layer_norm_params is not None:
+                param.data = layer_norm_params[name] # restore float32 weights
+            else:
+                layer_norm_state_dict[name] = param.data.detach().clone() # store float32 weights for stability
+                param.data = param.data.to(torch.float16)
+
+    return model, layer_norm_state_dict
+
+
+class PPOPeftTrainer(PPOTrainer, PeftTrainer):
+    r"""
+    Inherits PPOTrainer.
+    """
+
+    def __init__(
+            self,
+            training_args: Seq2SeqTrainingArguments,
+            finetuning_args: FinetuningArguments,
+            callbacks: List[LogCallback],
+            **kwargs
+    ):
+        PPOTrainer.__init__(self, **kwargs)
+        self.args = training_args
+        self.finetuning_args = finetuning_args
+        self.log_callback = callbacks[0]
+        self.state = TrainerState()
+        self.data_collator = self.accelerator.prepare(kwargs["data_collator"]) # override the data collator of PPOTrainer
+
+    def ppo_train(self, max_target_length: int) -> None:
+        r"""
+        Implements training loop for the PPO stage, like _inner_training_loop() in Huggingface's Trainer.
+        """
+        total_train_batch_size = self.config.batch_size * self.config.gradient_accumulation_steps * self.args.world_size
+        len_dataloader = len(self.dataloader)
+        num_steps_per_epoch = max(len_dataloader // self.config.gradient_accumulation_steps, 1)
+        num_examples = len(self.dataset)
+        num_train_epochs = self.args.num_train_epochs
+        max_steps = math.ceil(num_train_epochs * num_steps_per_epoch)
+
+        self.state.max_steps = max_steps
+        self.state.num_train_epochs = num_train_epochs
+        self.state.is_local_process_zero = self.is_local_process_zero()
+        self.state.is_world_process_zero = self.is_world_process_zero()
+
+        if self.is_world_process_zero():
+            logger.info("***** Running training *****")
+            logger.info(f"  Num examples = {num_examples}")
+            logger.info(f"  Num Epochs = {num_train_epochs}")
+            logger.info(f"  Instantaneous batch size per device = {self.config.batch_size}")
+            logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
+            logger.info(f"  Gradient Accumulation steps = {self.config.gradient_accumulation_steps}")
+            logger.info(f"  Total optimization steps = {max_steps}")
+            logger.info(f"  Number of trainable parameters = {sum(p.numel() for p in self.model.parameters() if p.requires_grad)}")
+
+        # Keyword arguments for `model.generate`
+        gen_kwargs = {
+            "top_k": 0.0,
+            "top_p": 1.0,
+            "do_sample": True,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+            "logits_processor": get_logits_processor()
+        }
+        output_length_sampler = LengthSampler(max_target_length // 2, max_target_length)
+        unwrapped_model: PreTrainedModel = self.accelerator.unwrap_model(self.model)
+
+        dataiter = iter(self.dataloader)
+        steps_trained = 0
+        loss_meter = AverageMeter()
+        reward_meter = AverageMeter()
+
+        for step in tqdm(range(max_steps), disable=not self.is_world_process_zero()):
+
+            for _ in range(self.config.gradient_accumulation_steps):
+
+                batch = next(dataiter)
+                steps_trained += 1
+
+                unwrapped_model.gradient_checkpointing_disable()
+                unwrapped_model.config.use_cache = True
+
+                # Get response from model
+                query_tensors: torch.Tensor = batch["input_ids"]
+                response_tensors = self.generate(batch, length_sampler=output_length_sampler, return_prompt=False, **gen_kwargs)
+
+                queries: List[torch.Tensor] = []
+                responses: List[torch.Tensor] = []
+                for i in range(len(query_tensors)):
+                    query_length = (query_tensors[i] != self.tokenizer.pad_token_id).nonzero()[0]
+                    response_length = (response_tensors[i] != self.tokenizer.pad_token_id).nonzero()[-1] + 1
+                    queries.append(query_tensors[i, query_length:]) # remove padding from left
+                    if response_length < 2: # make response have at least 2 tokens
+                        responses.append(response_tensors.new_empty(2).fill_(self.tokenizer.eos_token_id))
+                    else:
+                        responses.append(response_tensors[i, :response_length]) # remove padding from right
+
+                # Compute rewards
+                replace_model(unwrapped_model, target="reward")
+                _, _, values = self.model(**self.prepare_model_inputs(queries, responses))
+                rewards = [reward for reward in values[:, -1].to(torch.float32)] # use float32 type
+                replace_model(unwrapped_model, target="default") # make sure the model is default at the end
+
+                # Run PPO step
+                unwrapped_model.gradient_checkpointing_enable()
+                unwrapped_model.config.use_cache = False
+
+                stats = self.step(queries, responses, rewards)
+
+                loss_meter.update(stats["ppo/loss/total"], n=len(rewards))
+                reward_meter.update(torch.stack(rewards).mean().item(), n=len(rewards))
+
+                if steps_trained == len_dataloader:
+                    dataiter = iter(self.dataloader)
+                    steps_trained = 0
+
+            if self.is_world_process_zero() and (step+1) % self.args.logging_steps == 0:
+                logs = {
+                    "loss": round(loss_meter.avg, 4),
+                    "reward": round(reward_meter.avg, 4),
+                    "learning_rate": stats["ppo/learning_rate"],
+                    "epoch": round(step / num_steps_per_epoch, 2)
+                }
+                print(logs)
+                logs["step"] = step
+                self.state.log_history.append(logs)
+                self.log_callback.on_log(self.args, self.state, None)
+                loss_meter.reset()
+                reward_meter.reset()
+
+            if (step+1) % self.args.save_steps == 0: # save checkpoint
+                self.save_model(os.path.join(self.args.output_dir, f"checkpoint-{step+1}"))
+
+    @torch.no_grad()
+    def generate(
+            self,
+            inputs: Dict[str, torch.Tensor],
+            length_sampler: Optional[Callable] = None,
+            return_prompt: Optional[bool] = True,
+            **generation_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Generates model's responses given queries.
+
+        Subclass and override to inject custom behavior.
+        """
+        self.model, layer_norm_params = cast_layernorm_dtype(self.model)
+
+        if length_sampler is not None:
+            generation_kwargs["max_new_tokens"] = length_sampler()
+
+        unwrapped_model = self.accelerator.unwrap_model(self.model)
+
+        response = unwrapped_model.generate(**inputs, **generation_kwargs)
+
+        # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
+        # Inspired by: https://github.com/huggingface/transformers/blob/v4.28.1/src/transformers/trainer_seq2seq.py#L273
+        if unwrapped_model.pretrained_model.generation_config._from_model_config:
+            unwrapped_model.pretrained_model.generation_config._from_model_config = False
+
+        self.model, _ = cast_layernorm_dtype(self.model, layer_norm_params)
+
+        if not return_prompt and not self.is_encoder_decoder:
+            return response[:, inputs["input_ids"].size(1):]
+        return response
+
+    def save_model(self, output_dir: Optional[str] = None) -> None:
+        r"""
+        Saves model checkpoint.
+
+        Subclass and override to inject custom behavior.
+        """
+        if self.args.should_save:
+            self._save(output_dir)

src/utils/seq2seq.py

@@ -0,0 +1,96 @@
+import os
+import json
+import numpy as np
+from dataclasses import dataclass
+from typing import Dict, List, Sequence, Tuple, Union
+
+from transformers.trainer import PredictionOutput
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+import jieba
+from rouge_chinese import Rouge
+from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+
+from .peft_trainer import PeftTrainer
+
+from .other import get_logger, IGNORE_INDEX
+
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ComputeMetrics:
+    r"""
+    Wraps the tokenizer into metric functions, used in Seq2SeqPeftTrainer.
+
+    Borrowed from: https://github.com/THUDM/ChatGLM-6B/blob/0c2806fea82683349194e21996dd6b3acc3c265b/ptuning/main.py#L307
+    """
+
+    tokenizer: PreTrainedTokenizer
+
+    def __call__(self, eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
+        r"""
+        Uses the model predictions to compute metrics.
+        """
+        preds, labels = eval_preds
+        if isinstance(preds, tuple):
+            preds = preds[0]
+        # Replace IGNORE_INDEX in the labels with pad_token_id as we cannot decode them if ignore_pad_token_for_loss=True.
+        preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
+        labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)
+
+        score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
+        for pred, label in zip(preds, labels):
+            pred = pred[(pred == self.tokenizer.bos_token_id).nonzero()[0][0]:] # remove the query
+            hypothesis = list(jieba.cut(self.tokenizer.decode(pred, skip_special_tokens=True)))
+            reference = list(jieba.cut(self.tokenizer.decode(label, skip_special_tokens=True)))
+
+            if len(" ".join(hypothesis).split()) == 0:
+                result = {"rouge-1": {"f": 0.0}, "rouge-2": {"f": 0.0}, "rouge-l": {"f": 0.0}}
+            else:
+                rouge = Rouge()
+                scores = rouge.get_scores(" ".join(hypothesis), " ".join(reference))
+                result = scores[0]
+
+            for k, v in result.items():
+                score_dict[k].append(round(v["f"] * 100, 4))
+
+            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
+            score_dict["bleu-4"].append(round(bleu_score * 100, 4))
+
+        return {k: float(np.mean(v)) for k, v in score_dict.items()}
+
+
+class Seq2SeqPeftTrainer(PeftTrainer):
+    r"""
+    Inherits PeftTrainer to compute generative metrics such as BLEU and ROUGE.
+    """
+
+    def save_predictions(
+            self,
+            predict_results: PredictionOutput,
+            tokenizer: PreTrainedTokenizer
+    ) -> None:
+        r"""
+        Saves model predictions to `output_dir`.
+
+        A custom behavior that not contained in Seq2SeqTrainer.
+        """
+        if not self.is_world_process_zero():
+            return
+
+        preds = np.where(predict_results.predictions != IGNORE_INDEX, predict_results.predictions, self.tokenizer.pad_token_id)
+        labels = np.where(predict_results.label_ids != IGNORE_INDEX, predict_results.label_ids, self.tokenizer.pad_token_id)
+
+        preds = [pred[(pred == self.tokenizer.bos_token_id).nonzero()[0][0]:] for pred in preds] # remove the queries
+        preds = [tokenizer.decode(pred, skip_special_tokens=True).strip() for pred in preds]
+        labels = [tokenizer.decode(label, skip_special_tokens=True).strip() for label in labels]
+
+        output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
+        logger.info(f"Saving prediction results to {output_prediction_file}")
+        with open(output_prediction_file, "w", encoding="utf-8") as writer:
+            res: List[str] = []
+            for pred, label in zip(preds, labels):
+                res.append(json.dumps({"label": label, "predict": pred}, ensure_ascii=False))
+            writer.write("\n".join(res))

src/utils/template.py

@@ -0,0 +1,138 @@
+from typing import List, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class Template:
+
+    name: str
+
+    def __post_init__(self):
+
+        if self.name == "vanilla":
+            r"""
+            Supports language model inference without histories.
+            """
+            self._register_template(
+                prefix="",
+                prompt="{query}",
+                sep="",
+                use_history=False
+            )
+
+        elif self.name == "alpaca":
+            r"""
+            Supports: https://huggingface.co/tatsu-lab/alpaca-7b-wdiff
+                      https://github.com/ymcui/Chinese-LLaMA-Alpaca
+            """
+            self._register_template(
+                prefix="Below is an instruction that describes a task. "
+                       "Write a response that appropriately completes the request.\n\n",
+                prompt="### Instruction:\n{query}\n\n### Response:\n",
+                sep="\n\n",
+                use_history=True
+            )
+
+        elif self.name == "vicuna":
+            r"""
+            Supports: https://huggingface.co/lmsys/vicuna-7b-delta-v1.1
+                      https://huggingface.co/lmsys/vicuna-13b-delta-v1.1
+            """
+            self._register_template(
+                prefix="A chat between a curious user and an artificial intelligence assistant. "
+                       "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+                prompt="USER: {query} ASSISTANT: ",
+                sep="</s>",
+                use_history=True
+            )
+
+        elif self.name == "belle":
+            r"""
+            Supports: https://huggingface.co/BelleGroup/BELLE-LLaMA-EXT-13B
+            """
+            self._register_template(
+                prefix="",
+                prompt="Human: {query}\n\nBelle: ",
+                sep="\n\n",
+                use_history=True
+            )
+
+        elif self.name == "linly":
+            r"""
+            Supports: https://github.com/CVI-SZU/Linly
+            """
+            self._register_template(
+                prefix="",
+                prompt="User: {query}\nBot: ",
+                sep="\n",
+                use_history=True
+            )
+
+        elif self.name == "billa":
+            r"""
+            Supports: https://github.com/Neutralzz/BiLLa
+            """
+            self._register_template(
+                prefix="",
+                prompt="Human: {query}\nAssistant: ",
+                sep="\n",
+                use_history=True
+            )
+
+        elif self.name == "ziya":
+            r"""
+            Supports: https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1
+            """
+            self._register_template(
+                prefix="",
+                prompt="<human>:{query}\n<bot>:",
+                sep="\n",
+                use_history=True
+            )
+
+        elif self.name == "aquila":
+            r"""
+            Supports: https://huggingface.co/qhduan/aquilachat-7b
+            """
+            self._register_template(
+                prefix="A chat between a curious human and an artificial intelligence assistant. "
+                       "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+                prompt="Human: {query}\nAssistant: ",
+                sep="###",
+                use_history=True
+            )
+
+        else:
+            raise ValueError("Template {} does not exist.".format(self.name))
+
+    def get_prompt(self, query: str, history: Optional[list] = None, prefix: Optional[str] = "") -> str:
+        r"""
+        Returns a string containing prompt without response.
+        """
+        return "".join(self._format_example(query, history, prefix))
+
+    def get_dialog(self, query: str, resp: str, history: Optional[list] = None, prefix: Optional[str] = "") -> List[str]:
+        r"""
+        Returns a list containing 2 * n elements where the 2k-th is a query and the (2k+1)-th is a response.
+        """
+        return self._format_example(query, history, prefix) + [resp]
+
+    def _register_template(self, prefix: str, prompt: str, sep: str, use_history: Optional[bool] = True) -> None:
+        self.prefix = prefix
+        self.prompt = prompt
+        self.sep = sep
+        self.use_history = use_history
+
+    def _format_example(self, query: str, history: Optional[list] = None, prefix: Optional[str] = "") -> List[str]:
+        prefix = prefix if prefix else self.prefix
+        history = history if (history and self.use_history) else []
+        history = history + [(query, "<dummy>")]
+        convs = []
+        for turn_idx, (user_query, bot_resp) in enumerate(history):
+            if turn_idx == 0:
+                convs.append(prefix + self.prompt.format(query=user_query))
+                convs.append(bot_resp)
+            else:
+                convs.append(self.sep + self.prompt.format(query=user_query))
+                convs.append(bot_resp)
+        return convs[:-1] # drop last

src/web_demo.py

@@ -0,0 +1,150 @@
+# coding=utf-8
+# Implements user interface in browser for fine-tuned models.
+# Usage: python web_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint
+
+
+import mdtex2html
+import gradio as gr
+
+from threading import Thread
+from utils import (
+    Template,
+    load_pretrained,
+    prepare_infer_args,
+    get_logits_processor
+)
+
+from transformers import TextIteratorStreamer
+from transformers.utils.versions import require_version
+
+
+require_version("gradio>=3.30.0", "To fix: pip install gradio>=3.30.0")
+
+
+model_args, data_args, finetuning_args, generating_args = prepare_infer_args()
+model, tokenizer = load_pretrained(model_args, finetuning_args)
+
+prompt_template = Template(data_args.prompt_template)
+
+
+def postprocess(self, y):
+    r"""
+    Overrides Chatbot.postprocess
+    """
+    if y is None:
+        return []
+    for i, (message, response) in enumerate(y):
+        y[i] = (
+            None if message is None else mdtex2html.convert((message)),
+            None if response is None else mdtex2html.convert(response),
+        )
+    return y
+
+
+gr.Chatbot.postprocess = postprocess
+
+
+def parse_text(text): # copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT
+    lines = text.split("\n")
+    lines = [line for line in lines if line != ""]
+    count = 0
+    for i, line in enumerate(lines):
+        if "```" in line:
+            count += 1
+            items = line.split("`")
+            if count % 2 == 1:
+                lines[i] = "<pre><code class=\"language-{}\">".format(items[-1])
+            else:
+                lines[i] = "<br /></code></pre>"
+        else:
+            if i > 0:
+                if count % 2 == 1:
+                    line = line.replace("`", "\`")
+                    line = line.replace("<", "&lt;")
+                    line = line.replace(">", "&gt;")
+                    line = line.replace(" ", "&nbsp;")
+                    line = line.replace("*", "&ast;")
+                    line = line.replace("_", "&lowbar;")
+                    line = line.replace("-", "&#45;")
+                    line = line.replace(".", "&#46;")
+                    line = line.replace("!", "&#33;")
+                    line = line.replace("(", "&#40;")
+                    line = line.replace(")", "&#41;")
+                    line = line.replace("$", "&#36;")
+                lines[i] = "<br />" + line
+    text = "".join(lines)
+    return text
+
+
+def predict(query, chatbot, max_length, top_p, temperature, history):
+    chatbot.append((parse_text(query), ""))
+
+    input_ids = tokenizer([prompt_template.get_prompt(query, history)], return_tensors="pt")["input_ids"]
+    input_ids = input_ids.to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+
+    gen_kwargs = {
+        "input_ids": input_ids,
+        "do_sample": generating_args.do_sample,
+        "top_p": top_p,
+        "temperature": temperature,
+        "num_beams": generating_args.num_beams,
+        "max_length": max_length,
+        "repetition_penalty": generating_args.repetition_penalty,
+        "logits_processor": get_logits_processor(),
+        "streamer": streamer
+    }
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        new_history = history + [(query, response)]
+        chatbot[-1] = (parse_text(query), parse_text(response))
+        yield chatbot, new_history
+
+
+def reset_user_input():
+    return gr.update(value="")
+
+
+def reset_state():
+    return [], []
+
+
+with gr.Blocks() as demo:
+
+    gr.HTML("""
+    <h1 align="center">
+        <a href="https://chato.cn/" target="_blank">
+            		百姓AI助手
+        </a>
+    </h1>
+    """)
+
+    chatbot = gr.Chatbot()
+
+    with gr.Row():
+        with gr.Column(scale=4):
+            with gr.Column(scale=12):
+                user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style(container=False)
+            with gr.Column(min_width=32, scale=1):
+                submitBtn = gr.Button("Submit", variant="primary")
+
+        with gr.Column(scale=1):
+            emptyBtn = gr.Button("Clear History")
+            max_length = gr.Slider(0, 2048, value=1024, step=1.0, label="Maximum length", interactive=True)
+            top_p = gr.Slider(0, 1, value=generating_args.top_p, step=0.01, label="Top P", interactive=True)
+            temperature = gr.Slider(0, 1.5, value=generating_args.temperature, step=0.01, label="Temperature", interactive=True)
+
+    history = gr.State([])
+
+    submitBtn.click(predict, [user_input, chatbot, max_length, top_p, temperature, history], [chatbot, history], show_progress=True)
+    submitBtn.click(reset_user_input, [], [user_input])
+
+    emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)
+
+demo.queue().launch(server_name="0.0.0.0", share=True, inbrowser=True)