modity code structure

Former-commit-id: 0682ed357210897e0b67c4a6eb31a94b3eb929f1

src/llmtuner/tuner/__init__.py

@@ -0,0 +1,5 @@
+from llmtuner.tuner.core import get_train_args, get_infer_args, load_model_and_tokenizer
+from llmtuner.tuner.pt import run_pt
+from llmtuner.tuner.sft import run_sft
+from llmtuner.tuner.rm import run_rm
+from llmtuner.tuner.ppo import run_ppo

src/llmtuner/tuner/core/__init__.py

@@ -0,0 +1,2 @@
+from llmtuner.tuner.core.parser import get_train_args, get_infer_args
+from llmtuner.tuner.core.loader import load_model_and_tokenizer

src/llmtuner/tuner/core/adapter.py

@@ -0,0 +1,94 @@
+import os
+import torch
+
+from transformers.modeling_utils import PreTrainedModel
+from peft import (
+    PeftModel,
+    TaskType,
+    LoraConfig,
+    get_peft_model
+)
+from peft.utils import CONFIG_NAME, WEIGHTS_NAME
+
+from llmtuner.extras.logging import get_logger
+from llmtuner.extras.save_and_load import load_trainable_params
+from llmtuner.hparams import ModelArguments, FinetuningArguments
+
+
+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:
+            assert load_trainable_params(model, model_args.checkpoint_dir[0]), "Model checkpoint is not correctly loaded."
+
+    if finetuning_args.finetuning_type == "lora":
+        logger.info("Fine-tuning method: LoRA")
+        latest_checkpoint = None
+
+        if model_args.checkpoint_dir is not None:
+            assert os.path.exists(os.path.join(model_args.checkpoint_dir[0], WEIGHTS_NAME)), \
+                "Provided path ({}) does not contain a LoRA weight.".format(model_args.checkpoint_dir[0])
+            assert os.path.exists(os.path.join(model_args.checkpoint_dir[0], CONFIG_NAME)), \
+                "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, latest_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 latest_checkpoint is not None: # resume lora training or quantized inference
+                model = PeftModel.from_pretrained(model, latest_checkpoint, is_trainable=is_trainable)
+
+        if is_trainable and latest_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

src/llmtuner/tuner/core/loader.py

@@ -0,0 +1,151 @@
+import os
+import torch
+from typing import Literal, Optional, Tuple
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    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
+from trl import AutoModelForCausalLMWithValueHead
+
+from llmtuner.extras.logging import get_logger
+from llmtuner.extras.misc import prepare_model_for_training, print_trainable_params
+from llmtuner.extras.save_and_load import load_valuehead_params
+from llmtuner.hparams import ModelArguments, FinetuningArguments
+from llmtuner.tuner.core.adapter import init_adapter
+
+
+logger = get_logger(__name__)
+
+
+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")
+
+
+def load_model_and_tokenizer(
+    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=model_args.padding_side,
+        **config_kwargs
+    )
+    if tokenizer.pad_token_id is None or tokenizer.pad_token_id == 64000: # 64000 for baichuan model (older version)
+        tokenizer.pad_token_id = 0 # set as the <unk> token
+
+    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"
+
+    if model_args.checkpoint_dir is not None and finetuning_args.finetuning_type == "full":
+        model_to_load = model_args.checkpoint_dir[0]
+    else:
+        model_to_load = model_args.model_name_or_path
+
+    # Load and prepare pretrained models (without valuehead).
+    model = AutoModelForCausalLM.from_pretrained(
+        model_to_load,
+        config=config,
+        torch_dtype=torch.bfloat16 if model_args.compute_dtype == torch.bfloat16 else torch.float16,
+        low_cpu_mem_usage=True,
+        **config_kwargs
+    )
+
+    # Register auto class to save the custom code files.
+    if hasattr(config, "auto_map") and "AutoConfig" in config.auto_map:
+        config.__class__.register_for_auto_class()
+    if hasattr(config, "auto_map") and "AutoTokenizer" in config.auto_map:
+        tokenizer.__class__.register_for_auto_class()
+    if hasattr(config, "auto_map") and "AutoModelForCausalLM" in config.auto_map:
+        model.__class__.register_for_auto_class()
+
+    # Initialize adapters
+    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

src/llmtuner/tuner/core/parser.py

@@ -0,0 +1,134 @@
+import os
+import sys
+import torch
+import datasets
+import transformers
+from typing import Any, Dict, Optional, Tuple
+from transformers import HfArgumentParser, Seq2SeqTrainingArguments
+
+from llmtuner.extras.logging import get_logger
+from llmtuner.hparams import (
+    ModelArguments,
+    DataArguments,
+    FinetuningArguments,
+    GeneratingArguments,
+    GeneralArguments
+)
+
+
+logger = get_logger(__name__)
+
+
+def get_train_args(
+    args: Optional[Dict[str, Any]] = None
+) -> Tuple[ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments]:
+
+    parser = HfArgumentParser((ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments))
+
+    if args is not None:
+        model_args, data_args, training_args, finetuning_args, general_args = parser.parse_dict(args)
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+        model_args, data_args, training_args, finetuning_args, general_args = parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        model_args, data_args, training_args, finetuning_args, general_args = parser.parse_json_file(os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args, finetuning_args, general_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)
+    data_args.init_for_training()
+
+    assert general_args.stage == "sft" or (not training_args.predict_with_generate), \
+        "`predict_with_generate` cannot be set as True at PT, RM and PPO stages."
+
+    assert not (training_args.do_train and training_args.predict_with_generate), \
+        "`predict_with_generate` cannot be set as True while training."
+
+    assert (not training_args.do_predict) or training_args.predict_with_generate, \
+        "Please enable `predict_with_generate` to save model predictions."
+
+    assert model_args.quantization_bit is None or finetuning_args.finetuning_type == "lora", \
+        "Quantization is only compatible with the LoRA method."
+
+    if model_args.checkpoint_dir is not None:
+        if finetuning_args.finetuning_type != "lora":
+            assert len(model_args.checkpoint_dir) == 1, "Only LoRA tuning accepts multiple checkpoints."
+        else:
+            assert model_args.quantization_bit is None or len(model_args.checkpoint_dir) == 1, \
+                "Quantized model only accepts a single checkpoint."
+
+    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 == "default":
+        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, general_args
+
+
+def get_infer_args(
+    args: Optional[Dict[str, Any]] = None
+) -> Tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]:
+
+    parser = HfArgumentParser((ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments))
+
+    if args is not None:
+        model_args, data_args, finetuning_args, generating_args = parser.parse_dict(args)
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+        model_args, data_args, finetuning_args, generating_args = parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        model_args, data_args, finetuning_args, generating_args = parser.parse_json_file(os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, finetuning_args, generating_args = parser.parse_args_into_dataclasses()
+
+    assert model_args.quantization_bit is None or finetuning_args.finetuning_type == "lora", \
+        "Quantization is only compatible with the LoRA method."
+
+    if model_args.checkpoint_dir is not None:
+        if finetuning_args.finetuning_type != "lora":
+            assert len(model_args.checkpoint_dir) == 1, "Only LoRA tuning accepts multiple checkpoints."
+        else:
+            assert model_args.quantization_bit is None or len(model_args.checkpoint_dir) == 1, \
+                "Quantized model only accepts a single checkpoint."
+
+    if data_args.prompt_template == "default":
+        logger.warning("Please specify `prompt_template` if you are using other pre-trained models.")
+
+    return model_args, data_args, finetuning_args, generating_args

src/llmtuner/tuner/core/trainer.py

@@ -0,0 +1,85 @@
+import os
+import torch
+from typing import Dict, Optional
+
+from transformers import Seq2SeqTrainer
+from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.modeling_utils import unwrap_model
+
+from llmtuner.extras.constants import FINETUNING_ARGS_NAME, VALUE_HEAD_FILE_NAME
+from llmtuner.extras.logging import get_logger
+from llmtuner.extras.save_and_load import get_state_dict, load_trainable_params, load_valuehead_params
+from llmtuner.hparams import FinetuningArguments
+
+
+logger = get_logger(__name__)
+
+
+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 (currently using LoRA only)
+            backbone_model = getattr(model, "pretrained_model")
+            torch.save(get_state_dict(getattr(model, "v_head")), os.path.join(output_dir, VALUE_HEAD_FILE_NAME))
+        else:
+            backbone_model = model
+
+        if self.finetuning_args.finetuning_type == "lora":
+            backbone_model.save_pretrained(output_dir, state_dict=get_state_dict(backbone_model))
+        else: # freeze/full tuning
+            backbone_model.config.use_cache = True
+            backbone_model.save_pretrained(
+                output_dir,
+                state_dict=get_state_dict(backbone_model),
+                safe_serialization=self.args.save_safetensors
+            )
+            backbone_model.config.use_cache = False
+            if self.tokenizer is not None:
+                self.tokenizer.save_pretrained(output_dir)
+
+        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)
+        backbone_model = getattr(model, "pretrained_model") if hasattr(model, "pretrained_model") else model
+
+        if self.finetuning_args.finetuning_type == "lora":
+            backbone_model.load_adapter(self.state.best_model_checkpoint, getattr(backbone_model, "active_adapter"))
+            if hasattr(model, "v_head") and load_valuehead_params(model, self.state.best_model_checkpoint):
+                model.v_head.load_state_dict({
+                    "summary.weight": getattr(model, "reward_head_weight"),
+                    "summary.bias": getattr(model, "reward_head_bias")
+                })
+        else: # freeze/full-tuning
+            load_trainable_params(backbone_model, self.state.best_model_checkpoint)

src/llmtuner/tuner/ppo/__init__.py

@@ -0,0 +1 @@
+from llmtuner.tuner.ppo.workflow import run_ppo

src/llmtuner/tuner/ppo/trainer.py

@@ -0,0 +1,195 @@
+import os
+import math
+import torch
+from tqdm import tqdm
+from typing import Callable, Dict, List, Optional
+
+from transformers import Seq2SeqTrainingArguments, TrainerState, TrainerControl
+from transformers.modeling_utils import PreTrainedModel
+
+from trl import PPOTrainer
+from trl.core import LengthSampler
+
+from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.logging import get_logger
+from llmtuner.extras.misc import AverageMeter, get_logits_processor
+from llmtuner.hparams import FinetuningArguments
+from llmtuner.tuner.core.trainer import PeftTrainer
+from llmtuner.tuner.ppo.utils import cast_layernorm_dtype, replace_model
+
+
+logger = get_logger(__name__)
+
+
+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.control = TrainerControl()
+        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()
+        self.log_callback.on_train_begin(self.args, self.state, self.control)
+
+        for step in tqdm(range(max_steps), disable=not self.is_world_process_zero(), leave=False):
+
+            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 self.control.should_epoch_stop or self.control.should_training_stop:
+                    break
+
+                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, self.control)
+                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}"))
+
+            if self.control.should_training_stop:
+                break
+
+    @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/llmtuner/tuner/ppo/utils.py

@@ -0,0 +1,37 @@
+import torch
+from typing import Dict, List, Literal, Optional, Tuple
+from trl import AutoModelForCausalLMWithValueHead
+
+from llmtuner.extras.constants import LAYERNORM_NAMES
+
+
+def replace_model(model: AutoModelForCausalLMWithValueHead, target: Literal["default", "reward"]) -> None:
+    if target == "reward": # save default head temporarily
+        valuehead_state_dict = model.v_head.state_dict()
+        setattr(model, "default_head_weight", valuehead_state_dict["summary.weight"])
+        setattr(model, "default_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] = LAYERNORM_NAMES,
+    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

src/llmtuner/tuner/ppo/workflow.py

@@ -0,0 +1,68 @@
+# 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 trl import PPOConfig
+from torch.optim import AdamW
+from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments
+from transformers.optimization import get_scheduler
+
+from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.ploting import plot_loss
+from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+from llmtuner.tuner.core import load_model_and_tokenizer
+from llmtuner.tuner.ppo.trainer import PPOPeftTrainer
+
+
+def run_ppo(
+    model_args: ModelArguments,
+    data_args: DataArguments,
+    training_args: Seq2SeqTrainingArguments,
+    finetuning_args: FinetuningArguments
+):
+    dataset = get_dataset(model_args, data_args)
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="ppo")
+    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="ppo")
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, label_pad_token_id=tokenizer.pad_token_id)
+
+    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"])

src/llmtuner/tuner/pt/__init__.py

@@ -0,0 +1 @@
+from llmtuner.tuner.pt.workflow import run_pt

src/llmtuner/tuner/pt/workflow.py

@@ -0,0 +1,73 @@
+# Inspired by: https://github.com/huggingface/transformers/blob/v4.29.2/examples/pytorch/language-modeling/run_clm.py
+
+import math
+from typing import Optional, List
+from transformers import Seq2SeqTrainingArguments, DataCollatorForSeq2Seq, TrainerCallback
+
+from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.extras.ploting import plot_loss
+from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+from llmtuner.tuner.core import load_model_and_tokenizer
+from llmtuner.tuner.core.trainer import PeftTrainer
+
+
+def run_pt(
+    model_args: ModelArguments,
+    data_args: DataArguments,
+    training_args: Seq2SeqTrainingArguments,
+    finetuning_args: FinetuningArguments,
+    callbacks: Optional[List[TrainerCallback]] = [LogCallback()]
+):
+    dataset = get_dataset(model_args, data_args)
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="pt")
+    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="pt")
+    data_collator = DataCollatorForSeq2Seq(
+        tokenizer=tokenizer,
+        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+    )
+
+    # 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=callbacks,
+        **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)

src/llmtuner/tuner/rm/__init__.py

@@ -0,0 +1 @@
+from llmtuner.tuner.rm.workflow import run_rm

src/llmtuner/tuner/rm/collator.py

@@ -0,0 +1,19 @@
+import torch
+from typing import Any, Dict, Sequence
+from transformers import DataCollatorWithPadding
+
+
+class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
+    r"""
+    Data collator for pairwise data.
+    """
+
+    def __call__(self, features: Sequence[Dict[str, Any]]) -> 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)

src/llmtuner/tuner/rm/metric.py

@@ -0,0 +1,7 @@
+import numpy as np
+from typing import Dict, Sequence, Tuple, Union
+
+
+def compute_accuracy(eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
+    preds, _ = eval_preds
+    return {"accuracy": (preds[0] > preds[1]).sum() / len(preds[0])}

src/llmtuner/tuner/rm/trainer.py

@@ -0,0 +1,38 @@
+import torch
+from typing import Dict, List, Optional, Tuple, Union
+from transformers.modeling_utils import PreTrainedModel
+
+from llmtuner.tuner.core.trainer import PeftTrainer
+
+
+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: PreTrainedModel,
+        inputs: Dict[str, torch.Tensor],
+        return_outputs: Optional[bool] = False
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
+        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.
+
+        Note that the first element will be removed from the output tuple.
+
+        See: https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/trainer.py#L3509
+        """
+        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, [loss, r_accept, r_reject]) if return_outputs else loss

src/llmtuner/tuner/rm/workflow.py

@@ -0,0 +1,66 @@
+# 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 transformers import Seq2SeqTrainingArguments
+
+from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.ploting import plot_loss
+from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+from llmtuner.tuner.core import load_model_and_tokenizer
+from llmtuner.tuner.rm.metric import compute_accuracy
+from llmtuner.tuner.rm.collator import PairwiseDataCollatorWithPadding
+from llmtuner.tuner.rm.trainer import PairwisePeftTrainer
+
+
+def run_rm(
+    model_args: ModelArguments,
+    data_args: DataArguments,
+    training_args: Seq2SeqTrainingArguments,
+    finetuning_args: FinetuningArguments
+):
+    dataset = get_dataset(model_args, data_args)
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="rm")
+    dataset = preprocess_dataset(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)

src/llmtuner/tuner/sft/__init__.py

@@ -0,0 +1 @@
+from llmtuner.tuner.sft.workflow import run_sft

src/llmtuner/tuner/sft/metric.py

@@ -0,0 +1,51 @@
+import numpy as np
+from dataclasses import dataclass
+from typing import Dict, Sequence, Tuple, Union
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+import jieba
+from rouge_chinese import Rouge
+from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+
+from llmtuner.extras.constants import IGNORE_INDEX
+
+
+@dataclass
+class ComputeMetrics:
+    r"""
+    Wraps the tokenizer into metric functions, used in Seq2SeqPeftTrainer.
+    """
+
+    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
+        score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
+
+        preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
+        labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)
+
+        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True)
+        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        for pred, label in zip(decoded_preds, decoded_labels):
+            hypothesis = list(jieba.cut(pred))
+            reference = list(jieba.cut(label))
+
+            if len(" ".join(hypothesis).split()) == 0 or len(" ".join(reference).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()}

src/llmtuner/tuner/sft/trainer.py

@@ -0,0 +1,71 @@
+import os
+import json
+import torch
+import numpy as np
+import torch.nn as nn
+from typing import Any, Dict, List, Optional, Tuple, Union
+from transformers.trainer import PredictionOutput
+
+from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.extras.logging import get_logger
+from llmtuner.tuner.core.trainer import PeftTrainer
+
+
+logger = get_logger(__name__)
+
+
+class Seq2SeqPeftTrainer(PeftTrainer):
+    r"""
+    Inherits PeftTrainer to compute generative metrics such as BLEU and ROUGE.
+    """
+
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        r"""
+        Removes the prompt part in the generated tokens.
+
+        Subclass and override to inject custom behavior.
+        """
+        prompt_len, label_len = inputs["input_ids"].size(-1), inputs["labels"].size(-1)
+        if self.tokenizer.padding_side == "right": # pads the labels to the same length as the inputs
+            inputs["labels"] = torch.cat((inputs["labels"], torch.zeros_like(inputs["input_ids"])[:, label_len:]), dim=-1)
+        else:
+            inputs["labels"] = torch.cat((torch.zeros_like(inputs["input_ids"])[:, label_len:], inputs["labels"]), dim=-1)
+        loss, generated_tokens, labels = super().prediction_step(
+            model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
+        )
+        generated_tokens = generated_tokens[:, prompt_len:] if generated_tokens is not None else None
+
+        return (loss, generated_tokens, labels)
+
+    def save_predictions(
+        self,
+        predict_results: PredictionOutput
+    ) -> None:
+        r"""
+        Saves model predictions to `output_dir`.
+
+        A custom behavior that not contained in Seq2SeqTrainer.
+        """
+        if not self.is_world_process_zero():
+            return
+
+        output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
+        logger.info(f"Saving prediction results to {output_prediction_file}")
+
+        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)
+
+        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+
+        with open(output_prediction_file, "w", encoding="utf-8") as writer:
+            res: List[str] = []
+            for pred, label in zip(decoded_preds, decoded_labels):
+                res.append(json.dumps({"label": label, "predict": pred}, ensure_ascii=False))
+            writer.write("\n".join(res))

src/llmtuner/tuner/sft/workflow.py

@@ -0,0 +1,94 @@
+# Inspired by: https://github.com/huggingface/transformers/blob/v4.29.2/examples/pytorch/summarization/run_summarization.py
+
+from typing import Optional, List
+from transformers import Seq2SeqTrainingArguments, DataCollatorForSeq2Seq, TrainerCallback
+
+from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.extras.misc import get_logits_processor
+from llmtuner.extras.ploting import plot_loss
+from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+from llmtuner.tuner.core import load_model_and_tokenizer
+from llmtuner.tuner.sft.metric import ComputeMetrics
+from llmtuner.tuner.sft.trainer import Seq2SeqPeftTrainer
+
+
+def run_sft(
+    model_args: ModelArguments,
+    data_args: DataArguments,
+    training_args: Seq2SeqTrainingArguments,
+    finetuning_args: FinetuningArguments,
+    callbacks: Optional[List[TrainerCallback]] = [LogCallback()]
+):
+    dataset = get_dataset(model_args, data_args)
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="sft")
+    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="sft")
+    data_collator = DataCollatorForSeq2Seq(
+        tokenizer=tokenizer,
+        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+    )
+
+    # 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=callbacks,
+        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)
+        if training_args.predict_with_generate: # eval_loss will be wrong if predict_with_generate is enabled
+            metrics.pop("eval_loss", None)
+        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)
+        if training_args.predict_with_generate: # predict_loss will be wrong if predict_with_generate is enabled
+            predict_results.metrics.pop("predict_loss", None)
+        trainer.log_metrics("predict", predict_results.metrics)
+        trainer.save_metrics("predict", predict_results.metrics)
+        trainer.save_predictions(predict_results)