Merge branch 'main' into main

Former-commit-id: 870d2c7bf74d0da5a927bef4b8b01d15cc66a3e9
2023-08-18 01:37:23 +08:00
parent 458955d0fb ed10486cad
commit 3126164aa6
55 changed files with 1656 additions and 670 deletions
--- a/src/llmtuner/tuner/core/adapter.py
+++ b/src/llmtuner/tuner/core/adapter.py
@@ -39,7 +39,7 @@ def init_adapter(
    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":
+    if finetuning_args.finetuning_type == "full" and is_trainable:
        logger.info("Fine-tuning method: Full")
        model = model.float()

@@ -65,7 +65,7 @@ def init_adapter(
            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
+            if (is_trainable and finetuning_args.resume_lora_training) or (not is_mergeable): # continually fine-tuning
                checkpoints_to_merge, latest_checkpoint = model_args.checkpoint_dir[:-1], model_args.checkpoint_dir[-1]
            else:
                checkpoints_to_merge = model_args.checkpoint_dir
--- a/src/llmtuner/tuner/core/loader.py
+++ b/src/llmtuner/tuner/core/loader.py
@@ -1,5 +1,7 @@
 import os
+import math
 import torch
+from types import MethodType
 from typing import TYPE_CHECKING, Literal, Optional, Tuple

 from transformers import (
@@ -34,7 +36,7 @@ check_min_version("4.29.1")
 require_version("datasets>=2.12.0", "To fix: pip install datasets>=2.12.0")
 require_version("accelerate>=0.21.0", "To fix: pip install accelerate>=0.21.0")
 require_version("peft>=0.4.0", "To fix: pip install peft>=0.4.0")
-require_version("trl>=0.4.7", "To fix: pip install trl>=0.4.7")
+require_version("trl>=0.5.0", "To fix: pip install trl>=0.5.0")


 def load_model_and_tokenizer(
@@ -52,9 +54,6 @@ def load_model_and_tokenizer(
        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,
@@ -69,15 +68,58 @@ def load_model_and_tokenizer(
        **config_kwargs
    )

-    if model_args.checkpoint_dir is not None and finetuning_args.finetuning_type == "full":
+    if finetuning_args.finetuning_type == "full" and model_args.checkpoint_dir is not None:
        model_to_load = model_args.checkpoint_dir[0]
    else:
        model_to_load = model_args.model_name_or_path

    config = AutoConfig.from_pretrained(model_to_load, **config_kwargs)
-    is_mergeable = True
+
+    if hasattr(config, "fp16") and hasattr(config, "bf16"): # fix Qwen config
+        if model_args.compute_dtype == torch.bfloat16:
+            setattr(config, "bf16", True)
+        else:
+            setattr(config, "fp16", True)
+
+    # Set RoPE scaling
+    if model_args.rope_scaling is not None:
+        if hasattr(config, "use_dynamic_ntk"): # for Qwen models
+            if is_trainable:
+                logger.warning("Qwen model does not support RoPE scaling in training.")
+            else:
+                setattr(config, "use_dynamic_ntk", True)
+                setattr(config, "use_logn_attn", True)
+                logger.info("Using dynamic NTK scaling.")
+
+        elif hasattr(config, "rope_scaling"): # for LLaMA models
+            require_version("transformers>=4.31.0", "RoPE scaling requires transformers>=4.31.0")
+
+            if is_trainable:
+                if model_args.rope_scaling == "dynamic":
+                    logger.warning(
+                        "Dynamic NTK may not work well with fine-tuning. "
+                        "See: https://github.com/huggingface/transformers/pull/24653"
+                    )
+
+                current_max_length = getattr(config, "max_position_embeddings", None)
+                if current_max_length and model_args.model_max_length > current_max_length:
+                    scaling_factor = float(math.ceil(model_args.model_max_length / current_max_length))
+                else:
+                    logger.warning("Input length is smaller than max length. Consider increase input length.")
+                    scaling_factor = 1.0
+            else:
+                scaling_factor = 2.0
+
+            setattr(config, "rope_scaling", {"type": model_args.rope_scaling, "factor": scaling_factor})
+            logger.info("Using {} scaling strategy and setting scaling factor to {}".format(
+                model_args.rope_scaling, scaling_factor
+            ))
+
+        else:
+            logger.warning("Current model does not support RoPE scaling.")

    # Quantization configurations (using bitsandbytes library).
+    is_mergeable = True
    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")
@@ -95,10 +137,10 @@ def load_model_and_tokenizer(
            )

        is_mergeable = False
-        config_kwargs["device_map"] = {"": int(os.environ.get("LOCAL_RANK", "0"))}
+        config_kwargs["device_map"] = {"": int(os.environ.get("LOCAL_RANK", "0"))} if is_trainable else "auto"
        logger.info("Quantizing model to {} bit.".format(model_args.quantization_bit))

-    # Load and prepare pretrained models (without valuehead).
+    # Load and prepare pre-trained models (without valuehead).
    model = AutoModelForCausalLM.from_pretrained(
        model_to_load,
        config=config,
@@ -107,6 +149,14 @@ def load_model_and_tokenizer(
        **config_kwargs
    )

+    # Disable custom generate method (for Qwen)
+    if "GenerationMixin" not in str(model.generate.__func__):
+        model.generate = MethodType(PreTrainedModel.generate, model)
+
+    # Fix LM head (for ChatGLM2)
+    if not hasattr(model, "lm_head") and hasattr(model, "transformer"):
+        setattr(model, "lm_head", model.transformer.output_layer)
+
    # Register auto class to save the custom code files.
    if isinstance(config, PretrainedConfig) and "AutoConfig" in getattr(config, "auto_map", {}):
        config.__class__.register_for_auto_class()
@@ -119,10 +169,10 @@ def load_model_and_tokenizer(
    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" = AutoModelForCausalLMWithValueHead.from_pretrained(model)
+    # Prepare model with valuehead for RLHF
+    if stage == "rm" or stage == "ppo":
+        model: AutoModelForCausalLMWithValueHead = AutoModelForCausalLMWithValueHead.from_pretrained(model)
        reset_logging()
-
        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]):
@@ -132,15 +182,15 @@ def load_model_and_tokenizer(
                })

        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."

+    # Prepare model for inference
    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
+        infer_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 # detect cuda capability
+        model = model.to(infer_dtype) if model_args.quantization_bit is None else model

    trainable_params, all_param = count_parameters(model)
    logger.info("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
--- a/src/llmtuner/tuner/core/parser.py
+++ b/src/llmtuner/tuner/core/parser.py
@@ -19,7 +19,7 @@ from llmtuner.hparams import (
 logger = get_logger(__name__)


-def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None):
+def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
    if args is not None:
        return parser.parse_dict(args)
    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
@@ -32,26 +32,53 @@ def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None)

 def parse_train_args(
    args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments]:
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    Seq2SeqTrainingArguments,
+    FinetuningArguments,
+    GeneratingArguments,
+    GeneralArguments
+]:
    parser = HfArgumentParser((
-        ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments
+        ModelArguments,
+        DataArguments,
+        Seq2SeqTrainingArguments,
+        FinetuningArguments,
+        GeneratingArguments,
+        GeneralArguments
    ))
    return _parse_args(parser, args)


 def parse_infer_args(
    args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]:
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    FinetuningArguments,
+    GeneratingArguments
+]:
    parser = HfArgumentParser((
-        ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+        ModelArguments,
+        DataArguments,
+        FinetuningArguments,
+        GeneratingArguments
    ))
    return _parse_args(parser, args)


 def get_train_args(
    args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments]:
-    model_args, data_args, training_args, finetuning_args, general_args = parse_train_args(args)
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    Seq2SeqTrainingArguments,
+    FinetuningArguments,
+    GeneratingArguments,
+    GeneralArguments
+]:
+    model_args, data_args, training_args, finetuning_args, generating_args, general_args = parse_train_args(args)

    # Setup logging
    if training_args.should_log:
@@ -67,33 +94,42 @@ def get_train_args(
    # 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."
+    if general_args.stage != "sft" and training_args.predict_with_generate:
+        raise ValueError("`predict_with_generate` cannot be set as True except SFT.")

-    assert not (training_args.do_train and training_args.predict_with_generate), \
-        "`predict_with_generate` cannot be set as True while training."
+    if training_args.do_train and training_args.predict_with_generate:
+        raise ValueError("`predict_with_generate` cannot be set as True while training.")

-    assert general_args.stage != "sft" or (not training_args.do_predict) or training_args.predict_with_generate, \
-        "Please enable `predict_with_generate` to save model predictions."
+    if general_args.stage == "sft" and training_args.do_predict and not training_args.predict_with_generate:
+        raise ValueError("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 general_args.stage in ["rm", "ppo"] and finetuning_args.finetuning_type != "lora":
+        raise ValueError("RM and PPO training can only be performed with the LoRA method.")

-    assert not (training_args.max_steps == -1 and data_args.streaming), \
-        "Please specify `max_steps` in streaming mode."
+    if general_args.stage in ["ppo", "dpo"] and not training_args.do_train:
+        raise ValueError("PPO and DPO stage can only be performed at training.")

-    assert training_args.evaluation_strategy == "no" or (not data_args.streaming), \
-        "Streaming mode does not support evaluation currently."
+    if general_args.stage == "ppo" and model_args.reward_model is None:
+        raise ValueError("Reward model is necessary for PPO training.")

-    assert not (general_args.stage == "ppo" and data_args.streaming), \
-        "Streaming mode does not suppport PPO training currently."
+    if training_args.max_steps == -1 and data_args.streaming:
+        raise ValueError("Please specify `max_steps` in streaming mode.")
+
+    if general_args.stage == "ppo" and data_args.streaming:
+        raise ValueError("Streaming mode does not suppport PPO training currently.")
+
+    if data_args.val_size > 1e-6 and data_args.val_size < 1 and data_args.streaming:
+        raise ValueError("Streaming mode should have an integer val size.")
+
+    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.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 len(model_args.checkpoint_dir) != 1:
+                raise ValueError("Only LoRA tuning accepts multiple checkpoints.")
+        elif model_args.quantization_bit is not None and len(model_args.checkpoint_dir) != 1:
+                raise ValueError("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.")
@@ -113,46 +149,48 @@ def get_train_args(
        logger.warning("`max_samples` is incompatible with `streaming`. Disabling max_samples.")
        data_args.max_samples = None

-    if data_args.dev_ratio > 1e-6 and data_args.streaming:
-        logger.warning("`dev_ratio` is incompatible with `streaming`. Disabling development set.")
-        data_args.dev_ratio = 0
-
    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
+    if training_args.bf16:
+        if not torch.cuda.is_bf16_supported():
+            raise ValueError("Current device does not support bf16 training.")
+        model_args.compute_dtype = torch.bfloat16
+    else:
+        model_args.compute_dtype = torch.float16
+
+    model_args.model_max_length = data_args.max_source_length + data_args.max_target_length

    # Log on each process the small summary:
-    logger.info("Process rank: {}, device: {}, n_gpu: {}\n  distributed training: {}, 16-bits training: {}".format(
+    logger.info("Process rank: {}, device: {}, n_gpu: {}\n  distributed training: {}, compute dtype: {}".format(
        training_args.local_rank, training_args.device, training_args.n_gpu,
-        bool(training_args.local_rank != -1), training_args.fp16
+        bool(training_args.local_rank != -1), str(model_args.compute_dtype)
    ))
    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
+    return model_args, data_args, training_args, finetuning_args, generating_args, general_args


 def get_infer_args(
    args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]:
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    FinetuningArguments,
+    GeneratingArguments
+]:
    model_args, data_args, finetuning_args, generating_args = parse_infer_args(args)

-    assert model_args.quantization_bit is None or finetuning_args.finetuning_type == "lora", \
-        "Quantization is only compatible with the LoRA method."
+    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.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 len(model_args.checkpoint_dir) != 1:
+                raise ValueError("Only LoRA tuning accepts multiple checkpoints.")
+        elif model_args.quantization_bit is not None and len(model_args.checkpoint_dir) != 1:
+                raise ValueError("Quantized model only accepts a single checkpoint.")

    return model_args, data_args, finetuning_args, generating_args
--- a/src/llmtuner/tuner/core/trainer.py
+++ b/src/llmtuner/tuner/core/trainer.py
@@ -13,26 +13,25 @@ from llmtuner.extras.logging import get_logger
 from llmtuner.extras.save_and_load import get_state_dict, load_trainable_params

 if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer, Seq2SeqTrainingArguments, TrainerState
    from llmtuner.hparams import FinetuningArguments


 logger = get_logger(__name__)


-class PeftTrainer(Seq2SeqTrainer):
+class PeftModelMixin:
    r"""
-    Inherits Seq2SeqTrainer to support parameter-efficient checkpoints.
+    Patches the save and load methods in Hugging Face Trainer for PeftModel and ModelWithValueHead.
    """

-    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs):
-        super().__init__(**kwargs)
-        self.finetuning_args = finetuning_args
-        self._remove_log()
-
-    def _remove_log(self):
-        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 __init__(self) -> None: # for type checking
+        self.model: PreTrainedModel = None
+        self.tokenizer: "PreTrainedTokenizer" = None
+        self.args: "Seq2SeqTrainingArguments" = None
+        self.finetuning_args: "FinetuningArguments" = None
+        self.state: "TrainerState" = None
+        raise AssertionError("Mixin should not be initialized.")

    def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, torch.Tensor]] = None) -> None:
        r"""
@@ -96,3 +95,13 @@ class PeftTrainer(Seq2SeqTrainer):
            model.load_adapter(self.state.best_model_checkpoint, model.active_adapter)
        else: # freeze/full-tuning
            load_trainable_params(model, self.state.best_model_checkpoint)
+
+
+class PeftTrainer(PeftModelMixin, Seq2SeqTrainer):
+    r"""
+    Inherits Seq2SeqTrainer to support parameter-efficient checkpoints.
+    """
+
+    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs):
+        Seq2SeqTrainer.__init__(self, **kwargs)
+        self.finetuning_args = finetuning_args
--- a/src/llmtuner/tuner/dpo/init.py
+++ b/src/llmtuner/tuner/dpo/init.py
@@ -0,0 +1 @@
+from llmtuner.tuner.dpo.workflow import run_dpo
--- a/src/llmtuner/tuner/dpo/collator.py
+++ b/src/llmtuner/tuner/dpo/collator.py
@@ -0,0 +1,51 @@
+import torch
+from dataclasses import dataclass
+from typing import Any, Dict, List, Sequence, Tuple
+from transformers import DataCollatorForSeq2Seq
+
+
+@dataclass
+class DPODataCollatorWithPadding(DataCollatorForSeq2Seq):
+    r"""
+    Data collator for pairwise data.
+    """
+
+    def _pad_labels(self, batch: torch.Tensor, positions: List[Tuple[int, int]]) -> torch.Tensor:
+        padded_labels = []
+        for feature, (prompt_len, answer_len) in zip(batch, positions):
+            if self.tokenizer.padding_side == "left":
+                start, end = feature.size(0) - answer_len, feature.size(0)
+            else:
+                start, end = prompt_len, answer_len
+            padded_tensor = self.label_pad_token_id * torch.ones_like(feature)
+            padded_tensor[start:end] = feature[start:end]
+            padded_labels.append(padded_tensor)
+        return torch.stack(padded_labels, dim=0).contiguous() # in contiguous memory
+
+    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.
+        """
+        concatenated_features = []
+        label_positions = []
+        for key in ("chosen_ids", "rejected_ids"):
+            for feature in features:
+                prompt_len, answer_len = len(feature["prompt_ids"]), len(feature[key])
+                concatenated_features.append({
+                    "input_ids": feature["prompt_ids"] + feature[key],
+                    "attention_mask": [1] * (prompt_len + answer_len)
+                })
+                label_positions.append((prompt_len, answer_len))
+
+        batch = self.tokenizer.pad(
+            concatenated_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors,
+        )
+        batch["labels"] = self._pad_labels(batch["input_ids"], label_positions)
+        return batch
--- a/src/llmtuner/tuner/dpo/trainer.py
+++ b/src/llmtuner/tuner/dpo/trainer.py
@@ -0,0 +1,77 @@
+import torch
+from collections import defaultdict
+from peft import PeftModel
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+from transformers import BatchEncoding, Trainer
+from trl import DPOTrainer
+
+from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.tuner.core.trainer import PeftModelMixin
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+    from llmtuner.hparams import FinetuningArguments, GeneratingArguments
+
+
+class DPOPeftTrainer(PeftModelMixin, DPOTrainer):
+
+    def __init__(
+        self,
+        finetuning_args: "FinetuningArguments",
+        generating_args: "GeneratingArguments",
+        ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]] = None,
+        **kwargs
+    ):
+        self.finetuning_args = finetuning_args
+        self.generating_args = generating_args
+        self.ref_model = ref_model
+        self.use_dpo_data_collator = True # hack to avoid warning
+        self.label_pad_token_id = IGNORE_INDEX
+        self.padding_value = 0
+        self.beta = finetuning_args.dpo_beta
+        self._stored_metrics = defaultdict(lambda: defaultdict(list))
+
+        Trainer.__init__(self, **kwargs)
+        if not hasattr(self, "accelerator"):
+            raise AttributeError("Please update `transformers`.")
+
+        if ref_model is not None:
+            self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+
+    def concatenated_forward(
+        self,
+        model: Optional[torch.nn.Module] = None,
+        batch: Optional[Dict[str, torch.Tensor]] = None
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+        batch_copied = BatchEncoding({k: v.detach().clone() for k, v in batch.items()}) # avoid error
+        unwrapped_model: "PreTrainedModel" = self.accelerator.unwrap_model(self.model)
+
+        if not torch.is_grad_enabled():
+            unwrapped_model.gradient_checkpointing_disable()
+
+        if model is None and isinstance(unwrapped_model, PeftModel): # peft model has no ref_model
+            with unwrapped_model.disable_adapter():
+                all_logits = self.model(
+                    input_ids=batch_copied["input_ids"],
+                    attention_mask=batch_copied["attention_mask"],
+                    return_dict=True
+                ).logits.to(torch.float32)
+        else:
+            all_logits = model(
+                input_ids=batch_copied["input_ids"],
+                attention_mask=batch_copied["attention_mask"],
+                return_dict=True
+            ).logits.to(torch.float32)
+
+        if not torch.is_grad_enabled():
+            unwrapped_model.gradient_checkpointing_enable()
+
+        all_logps = self._get_batch_logps(
+            all_logits,
+            batch["labels"],
+            average_log_prob=False
+        )
+        batch_size = batch["input_ids"].size(0) // 2
+        chosen_logps, rejected_logps = all_logps.split(batch_size, dim=0)
+        chosen_logits, rejected_logits = all_logits.split(batch_size, dim=0)
+        return chosen_logps, rejected_logps, chosen_logits, rejected_logits
--- a/src/llmtuner/tuner/dpo/workflow.py
+++ b/src/llmtuner/tuner/dpo/workflow.py
@@ -0,0 +1,59 @@
+# Inspired by: https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py
+
+from copy import deepcopy
+from peft import PeftModel
+from typing import TYPE_CHECKING, Optional, List
+
+from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
+from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.extras.ploting import plot_loss
+from llmtuner.tuner.core import load_model_and_tokenizer
+from llmtuner.tuner.dpo.collator import DPODataCollatorWithPadding
+from llmtuner.tuner.dpo.trainer import DPOPeftTrainer
+
+if TYPE_CHECKING:
+    from transformers import Seq2SeqTrainingArguments, TrainerCallback
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+
+
+def run_dpo(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    generating_args: "GeneratingArguments",
+    callbacks: Optional[List["TrainerCallback"]] = None
+):
+    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="rm")
+    data_collator = DPODataCollatorWithPadding(
+        tokenizer=tokenizer,
+        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+    )
+
+    training_args.remove_unused_columns = False # important for pairwise dataset
+    ref_model = deepcopy(model) if not isinstance(model, PeftModel) else None
+
+    # Initialize our Trainer
+    trainer = DPOPeftTrainer(
+        finetuning_args=finetuning_args,
+        generating_args=generating_args,
+        ref_model=ref_model,
+        model=model,
+        args=training_args,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        callbacks=callbacks,
+        **split_dataset(dataset, data_args, training_args)
+    )
+
+    # 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"])
--- a/src/llmtuner/tuner/ppo/trainer.py
+++ b/src/llmtuner/tuner/ppo/trainer.py
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
 from transformers import TrainerState, TrainerControl

 from trl import PPOTrainer
-from trl.core import LengthSampler
+from trl.core import LengthSampler, PPODecorators, logprobs_from_logits

 from llmtuner.extras.logging import get_logger
 from llmtuner.extras.misc import AverageMeter, count_parameters, get_logits_processor
@@ -18,7 +18,7 @@ if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments
    from trl import AutoModelForCausalLMWithValueHead
    from llmtuner.extras.callbacks import LogCallback
-    from llmtuner.hparams import FinetuningArguments
+    from llmtuner.hparams import FinetuningArguments, GeneratingArguments


 logger = get_logger(__name__)
@@ -33,16 +33,19 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
        self,
        training_args: "Seq2SeqTrainingArguments",
        finetuning_args: "FinetuningArguments",
+        generating_args: "GeneratingArguments",
        callbacks: List["LogCallback"],
+        compute_dtype: torch.dtype,
        **kwargs
    ):
        PPOTrainer.__init__(self, **kwargs)
        self.args = training_args
        self.finetuning_args = finetuning_args
+        self.generating_args = generating_args
        self.log_callback = callbacks[0]
+        self.compute_dtype = compute_dtype
        self.state = TrainerState()
        self.control = TrainerControl()
-        self._remove_log()

    def ppo_train(self, max_target_length: int) -> None:
        r"""
@@ -72,14 +75,11 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
            logger.info(f"  Number of trainable parameters = {count_parameters(self.model)[0]}")

        # 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()
-        }
+        gen_kwargs = self.generating_args.to_dict()
+        gen_kwargs["eos_token_id"] = list(set([self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids))
+        gen_kwargs["pad_token_id"] = self.tokenizer.pad_token_id
+        gen_kwargs["logits_processor"] = get_logits_processor()
+
        length_sampler = LengthSampler(max_target_length // 2, max_target_length)
        unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)

@@ -185,10 +185,74 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
        replace_model(unwrapped_model, target="reward")
        batch = self.prepare_model_inputs(queries, responses)
        _, _, values = self.model(**batch, output_hidden_states=True, return_dict=True)
+        if values.size(0) != batch["input_ids"].size(0): # adapt chatglm2
+            values = torch.transpose(values, 0, 1)
        rewards = [reward for reward in values[:, -1].float().detach().cpu()] # use fp32 type
        replace_model(unwrapped_model, target="default")
        return rewards

+    @PPODecorators.empty_cuda_cache()
+    def batched_forward_pass(
+        self,
+        model: "AutoModelForCausalLMWithValueHead",
+        queries: torch.Tensor,
+        responses: torch.Tensor,
+        model_inputs: dict,
+        return_logits: Optional[bool] = False
+    ):
+        r"""
+        Calculates model outputs in multiple batches.
+
+        Subclass and override to inject custom behavior.
+        """
+        bs = len(queries)
+        fbs = self.config.mini_batch_size
+        all_logprobs = []
+        all_logits = []
+        all_masks = []
+        all_values = []
+
+        for i in range(math.ceil(bs / fbs)):
+            input_kwargs = {key: value[i * fbs : (i + 1) * fbs] for key, value in model_inputs.items()}
+            query_batch = queries[i * fbs : (i + 1) * fbs]
+            response_batch = responses[i * fbs : (i + 1) * fbs]
+            input_ids = input_kwargs["input_ids"]
+            attention_mask = input_kwargs["attention_mask"]
+
+            with torch.cuda.amp.autocast(dtype=self.compute_dtype): # support bf16
+                logits, _, values = model(**input_kwargs)
+
+            if values.size(0) != input_ids.size(0): # adapt chatglm2
+                values = torch.transpose(values, 0, 1)
+
+            logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
+            masks = torch.zeros_like(attention_mask)
+            masks[:, :-1] = attention_mask[:, 1:]
+
+            for j in range(len(query_batch)):
+                start = len(query_batch[j]) - 1
+                if attention_mask[j, 0] == 0:  # offset left padding
+                    start += attention_mask[j, :].nonzero()[0]
+                end = start + len(response_batch[j])
+
+                masks[j, :start] = 0
+                masks[j, end:] = 0
+
+            if return_logits:
+                all_logits.append(logits)
+            else:
+                del logits
+            all_values.append(values)
+            all_logprobs.append(logprobs)
+            all_masks.append(masks)
+
+        return (
+            torch.cat(all_logprobs),
+            torch.cat(all_logits)[:, :-1] if return_logits else None,
+            torch.cat(all_values)[:, :-1],
+            torch.cat(all_masks)[:, :-1],
+        )
+
    def save_model(self, output_dir: Optional[str] = None) -> None:
        r"""
        Saves model checkpoint.
--- a/src/llmtuner/tuner/ppo/workflow.py
+++ b/src/llmtuner/tuner/ppo/workflow.py
@@ -1,11 +1,9 @@
-# Inspired by:
-# https://github.com/lvwerra/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py
+# Inspired by: https://github.com/lvwerra/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py

 import math
-from typing import TYPE_CHECKING
 from trl import PPOConfig
 from torch.optim import AdamW
-from typing import Optional, List
+from typing import TYPE_CHECKING, Optional, List
 from transformers import DataCollatorForSeq2Seq
 from transformers.optimization import get_scheduler

@@ -16,7 +14,7 @@ from llmtuner.tuner.ppo.trainer import PPOPeftTrainer

 if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments


 def run_ppo(
@@ -24,6 +22,7 @@ def run_ppo(
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
+    generating_args: "GeneratingArguments",
    callbacks: Optional[List["TrainerCallback"]] = None
 ):
    dataset = get_dataset(model_args, data_args)
@@ -38,24 +37,30 @@ def run_ppo(
        batch_size=training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps,
        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
        ppo_epochs=1,
-        max_grad_norm=training_args.max_grad_norm
+        max_grad_norm=training_args.max_grad_norm,
+        seed=training_args.seed,
+        optimize_cuda_cache=True
    )

-    optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=ppo_config.learning_rate)
-    total_train_batch_size = \
+    optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=training_args.learning_rate)
+    total_train_batch_size = (
        training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
+    )
+    num_training_steps = training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_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))
+        num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
+        num_training_steps=num_training_steps
    )

    # Initialize our Trainer
    ppo_trainer = PPOPeftTrainer(
        training_args=training_args,
        finetuning_args=finetuning_args,
+        generating_args=generating_args,
        callbacks=callbacks,
+        compute_dtype=model_args.compute_dtype,
        config=ppo_config,
        model=model,
        ref_model=None,
@@ -66,8 +71,10 @@ def run_ppo(
        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"])
+    # Training
+    if training_args.do_train:
+        ppo_trainer.ppo_train(max_target_length=data_args.max_target_length)
+        ppo_trainer.save_model()
+        ppo_trainer.save_state() # must be called after save_model to have a folder
+        if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "reward"])
--- a/src/llmtuner/tuner/pt/workflow.py
+++ b/src/llmtuner/tuner/pt/workflow.py
@@ -2,10 +2,9 @@

 import math
 from typing import TYPE_CHECKING, Optional, List
-from transformers import DataCollatorForSeq2Seq
+from transformers import DataCollatorForLanguageModeling

 from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
-from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.tuner.core import load_model_and_tokenizer
 from llmtuner.tuner.core.trainer import PeftTrainer
@@ -25,10 +24,7 @@ def run_pt(
    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
-    )
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Initialize our Trainer
    trainer = PeftTrainer(
@@ -38,7 +34,7 @@ def run_pt(
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,
-        **split_dataset(dataset, data_args.dev_ratio, training_args.do_train)
+        **split_dataset(dataset, data_args, training_args)
    )

    # Training
@@ -60,6 +56,5 @@ def run_pt(
            perplexity = float("inf")

        metrics["perplexity"] = perplexity
-
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
--- a/src/llmtuner/tuner/rm/collator.py
+++ b/src/llmtuner/tuner/rm/collator.py
@@ -1,8 +1,10 @@
 import torch
+from dataclasses import dataclass
 from typing import Any, Dict, Sequence
 from transformers import DataCollatorWithPadding


+@dataclass
 class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
    r"""
    Data collator for pairwise data.
@@ -16,7 +18,10 @@ class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
        the last n examples represent rejected examples.
        """
        features = [
-            {"input_ids": feature[key], "attention_mask": [1] * len(feature[key])}
-            for key in ("accept_ids", "reject_ids") for feature in features
+            {
+                "input_ids": feature["prompt_ids"] + feature[key],
+                "attention_mask": [1] * (len(feature["prompt_ids"]) + len(feature[key]))
+            }
+            for key in ("chosen_ids", "rejected_ids") for feature in features
        ]
        return super().__call__(features)
--- a/src/llmtuner/tuner/rm/trainer.py
+++ b/src/llmtuner/tuner/rm/trainer.py
@@ -42,6 +42,8 @@ class PairwisePeftTrainer(PeftTrainer):
        """
        batch_size = inputs["input_ids"].size(0) // 2
        _, _, values = model(**inputs, output_hidden_states=True, return_dict=True)
+        if values.size(0) != inputs["input_ids"].size(0): # adapt chatglm2
+            values = torch.transpose(values, 0, 1)
        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
--- a/src/llmtuner/tuner/rm/workflow.py
+++ b/src/llmtuner/tuner/rm/workflow.py
@@ -39,7 +39,7 @@ def run_rm(
        data_collator=data_collator,
        callbacks=callbacks,
        compute_metrics=compute_accuracy,
-        **split_dataset(dataset, data_args.dev_ratio, training_args.do_train)
+        **split_dataset(dataset, data_args, training_args)
    )

    # Training
--- a/src/llmtuner/tuner/sft/metric.py
+++ b/src/llmtuner/tuner/sft/metric.py
@@ -25,7 +25,7 @@ class ComputeMetrics:
        Uses the model predictions to compute metrics.
        """
        preds, labels = eval_preds
-        score_dict = {"accuracy": [], "rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
+        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)
@@ -49,6 +49,5 @@ class ComputeMetrics:

            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
            score_dict["bleu-4"].append(round(bleu_score * 100, 4))
-            score_dict["accuracy"].append(float(len(label) != 0 and pred[:len(label)] == label))

        return {k: float(np.mean(v)) for k, v in score_dict.items()}
--- a/src/llmtuner/tuner/sft/trainer.py
+++ b/src/llmtuner/tuner/sft/trainer.py
@@ -50,9 +50,10 @@ class Seq2SeqPeftTrainer(PeftTrainer):
        loss, generated_tokens, labels = super().prediction_step(
            model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
        )
-        generated_tokens = (
-            generated_tokens[:, max(prompt_len, label_len):] if generated_tokens is not None else None
-        )
+        if generated_tokens is not None:
+            generated_tokens[:, :max(prompt_len, label_len)] = (
+                self.tokenizer.pad_token_id * torch.ones_like(generated_tokens[:, :max(prompt_len, label_len)])
+            )

        return (loss, generated_tokens, labels)

@@ -72,14 +73,11 @@ class Seq2SeqPeftTrainer(PeftTrainer):
                assert self.tokenizer.padding_side == "left", "This method only accepts left-padded tensor."
                pad_token_id = self.tokenizer.pad_token_id
            else:
-                if self.model.config.pad_token_id is not None:
-                    pad_token_id = self.model.config.pad_token_id
-                else:
-                    raise ValueError("Pad_token_id must be set in the configuration of the model.")
+                raise ValueError("PAD token is required.")

        padded_tensor = pad_token_id * torch.ones_like(tgt_tensor)
        padded_tensor[:, -src_tensor.shape[-1]:] = src_tensor # adopt left-padding
-        return padded_tensor.contiguous()
+        return padded_tensor.contiguous() # in contiguous memory

    def save_predictions(
        self,
--- a/src/llmtuner/tuner/sft/workflow.py
+++ b/src/llmtuner/tuner/sft/workflow.py
@@ -16,7 +16,7 @@ from llmtuner.extras.logging import reset_logging, get_logger

 if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments

 logger = get_logger(__name__)

@@ -25,6 +25,7 @@ def run_sft(
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
+    generating_args: "GeneratingArguments",
    callbacks: Optional[List["TrainerCallback"]] = None
 ):
    dataset = get_dataset(model_args, data_args)
@@ -50,31 +51,15 @@ def run_sft(
        data_collator=data_collator,
        callbacks=callbacks,
        compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
-        **split_dataset(dataset, data_args.dev_ratio, training_args.do_train)
+        **split_dataset(dataset, data_args, training_args)
    )

    # 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()
-    }
-    # Detecting last checkpoint.
-    last_checkpoint = None
-    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
-        last_checkpoint = get_last_checkpoint(training_args.output_dir)
-        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
-            raise ValueError(
-                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
-                "Use --overwrite_output_dir to overcome."
-            )
-        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
-            logger.info(
-                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
-                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
-            )
+    gen_kwargs = generating_args.to_dict()
+    gen_kwargs["eos_token_id"] = list(set([tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids))
+    gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
+    gen_kwargs["logits_processor"] = get_logits_processor()
+
    # Training
    if training_args.do_train:
        checkpoint = None
--- a/src/llmtuner/tuner/tune.py
+++ b/src/llmtuner/tuner/tune.py
@@ -1,35 +1,47 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional

 from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.logging import get_logger
 from llmtuner.tuner.core import get_train_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
+from llmtuner.tuner.dpo import run_dpo

 if TYPE_CHECKING:
    from transformers import TrainerCallback


+logger = get_logger(__name__)
+
+
 def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["TrainerCallback"]] = None):
-    model_args, data_args, training_args, finetuning_args, general_args = get_train_args(args)
+    model_args, data_args, training_args, finetuning_args, generating_args, general_args = get_train_args(args)
    callbacks = [LogCallback()] if callbacks is None else callbacks

    if general_args.stage == "pt":
        run_pt(model_args, data_args, training_args, finetuning_args, callbacks)
    elif general_args.stage == "sft":
-        run_sft(model_args, data_args, training_args, finetuning_args, callbacks)
+        run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
    elif general_args.stage == "rm":
        run_rm(model_args, data_args, training_args, finetuning_args, callbacks)
    elif general_args.stage == "ppo":
-        run_ppo(model_args, data_args, training_args, finetuning_args, callbacks)
+        run_ppo(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
+    elif general_args.stage == "dpo":
+        run_dpo(model_args, data_args, training_args, finetuning_args, callbacks)
+    else:
+        raise ValueError("Unknown task.")


 def export_model(args: Optional[Dict[str, Any]] = None, max_shard_size: Optional[str] = "10GB"):
-    model_args, _, training_args, finetuning_args, _ = get_train_args(args)
+    model_args, _, training_args, finetuning_args, _, _ = get_train_args(args)
    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
    model.save_pretrained(training_args.output_dir, max_shard_size=max_shard_size)
-    tokenizer.save_pretrained(training_args.output_dir)
+    try:
+        tokenizer.save_pretrained(training_args.output_dir)
+    except:
+        logger.warning("Cannot save tokenizer, please copy the files manually.")


 if __name__ == "__main__":
				`@@ -0,0 +1 @@`
				`from llmtuner.tuner.dpo.workflow import run_dpo`