[LLM][TRL] Support DPO with Pipeline Parallel (#9039)

* support dpo/kto pp

Author: lugimzzz

llm/alignment/dpo/dpo_argument.py

@@ -17,6 +17,7 @@
 from typing import Optional
 
 from paddlenlp.trainer import TrainingArguments
+from paddlenlp.trainer.trainer_utils import IntervalStrategy
 
 
 def add_start_docstrings(*docstr):
@@ -42,9 +43,66 @@ class DPOTrainingArguments(TrainingArguments):
         default="",
         metadata={"help": "Configs to unify hybrid parallel checkpoint.\n"},
     )
-    dpo_beta: float = field(default=0.1, metadata={"help": "the beta parameter for DPO loss"})
-    dpo_label_smoothing: float = field(default=0.0, metadata={"help": "label_smoothing ratio"})
-    dpo_loss_type: str = field(default="sigmoid", metadata={"help": "DPO loss type"})
+    autotuner_benchmark: bool = field(
+        default=False,
+        metadata={"help": "Whether to run benchmark by autotuner. True for from_scratch."},
+    )
+    benchmark: bool = field(
+        default=False,
+        metadata={"help": "Whether to run benchmark by autotuner. True for from_scratch."},
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.autotuner_benchmark:
+            self.num_train_epochs = 1
+            self.max_steps = 5
+            self.do_train = True
+            self.do_export = False
+            self.do_predict = False
+            self.do_eval = False
+            self.overwrite_output_dir = True
+            self.load_best_model_at_end = False
+            self.report_to = []
+            self.save_strategy = IntervalStrategy.NO
+            self.evaluation_strategy = IntervalStrategy.NO
+            if not self.disable_tqdm:
+                self.logging_steps = 1
+                self.logging_strategy = IntervalStrategy.STEPS
+        if self.benchmark:
+            self.do_train = True
+            self.do_export = False
+            self.do_predict = False
+            self.do_eval = False
+            self.overwrite_output_dir = True
+            self.load_best_model_at_end = False
+            self.save_strategy = IntervalStrategy.NO
+            self.evaluation_strategy = IntervalStrategy.NO
+            if not self.disable_tqdm:
+                self.logging_steps = 1
+                self.logging_strategy = IntervalStrategy.STEPS
+        if self.max_steps > 0:
+            self.num_train_epochs = 1
+
+
+@dataclass
+class DPOConfig:
+    """DPOConfig"""
+
+    beta: float = field(default=0.1, metadata={"help": "the beta parameter for DPO loss"})
+    simpo_gamma: float = field(default=0.5, metadata={"help": "the gamma parameter for SimPO loss"})
+    normalize_logps: bool = field(
+        default=True,
+        metadata={"help": "Apply logprobs normalization."},
+    )
+    label_smoothing: float = field(default=0.0, metadata={"help": "label_smoothing ratio"})
+    loss_type: str = field(default="sigmoid", metadata={"help": "DPO loss type"})
+    pref_loss_ratio: float = field(default=1.0, metadata={"help": "DPO loss ratio"})
+    sft_loss_ratio: float = field(default=0.0, metadata={"help": "SFT loss ratio"})
+    dpop_lambda: float = field(default=50, metadata={"help": "SFT loss ratio"})
+    ref_model_update_steps: int = field(default=-1, metadata={"help": "Update ref model state dict "})
+    reference_free: bool = field(default=False, metadata={"help": "No reference model."})
+    lora: bool = field(default=False, metadata={"help": "Use LoRA model."})
 
 
 @dataclass
@@ -55,18 +113,16 @@ class DPODataArgument:
     dev_dataset_path: str = field(default="./data/dev.jsonl", metadata={"help": "Path to the dev dataset dir."})
     max_seq_len: int = field(default=4096, metadata={"help": "Maximum sequence length."})
     max_prompt_len: int = field(default=2048, metadata={"help": "Maximum prompt length."})
-    autotuner_benchmark: bool = field(
-        default=False,
-        metadata={"help": "Whether to run benchmark by autotuner. True for from_scratch."},
-    )
-    benchmark: bool = field(
-        default=False,
-        metadata={"help": "Whether to run benchmark by autotuner. True for from_scratch."},
-    )
     greedy_zero_padding: bool = field(
         default=False,
         metadata={"help": "Whether to use Greedy Zero Padding data stream."},
     )
+    lazy: bool = field(
+        default=False,
+        metadata={
+            "help": "Weather to return `MapDataset` or an `IterDataset`.True for `IterDataset`. False for `MapDataset`."
+        },
+    )
 
 
 @dataclass
@@ -95,3 +151,19 @@ class DPOModelArgument:
         default=False,
         metadata={"help": "whether to use sequence parallel"},
     )
+    tensor_parallel_output: bool = field(
+        default=True,
+        metadata={"help": "whether to use tensor_parallel_output"},
+    )
+    weight_quantize_algo: str = field(
+        default=None,
+        metadata={"help": "Model weight quantization algorithm including 'nf4'(qlora), 'weight_only_int8'."},
+    )
+    # LoRA
+    lora_rank: int = field(default=8, metadata={"help": "Lora rank."})
+    lora_path: str = field(default=None, metadata={"help": "Initialize lora state dict."})
+    rslora: bool = field(default=False, metadata={"help": "Whether to use RsLoRA"})
+    lora_plus_scale: float = field(default=1.0, metadata={"help": "Lora B scale in LoRA+ technique"})
+    lora_alpha: int = field(default=-1, metadata={"help": "lora_alpha"})
+    rslora_plus: bool = field(default=False, metadata={"help": "Strengthen lora performance"})
+    use_quick_lora: bool = field(default=True, metadata={"help": "quick lora"})

llm/alignment/dpo/run_dpo.py

@@ -20,18 +20,24 @@
 from functools import partial
 
 import paddle
-from dpo_argument import DPODataArgument, DPOModelArgument, DPOTrainingArguments
-
-from paddlenlp.datasets import ZeroPaddingMapDataset, load_dataset
-from paddlenlp.trainer import (
-    IntervalStrategy,
-    PdArgumentParser,
-    get_last_checkpoint,
-    set_seed,
+from dpo_argument import (
+    DPOConfig,
+    DPODataArgument,
+    DPOModelArgument,
+    DPOTrainingArguments,
 )
+
+from paddlenlp.datasets import (
+    ZeroPaddingIterableDataset,
+    ZeroPaddingMapDataset,
+    load_dataset,
+)
+from paddlenlp.peft import LoRAConfig, LoRAModel
+from paddlenlp.trainer import PdArgumentParser, get_last_checkpoint, set_seed
 from paddlenlp.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
+    AutoModelForCausalLMPipe,
     AutoTokenizer,
     LlamaForCausalLM,
     LlamaForCausalLMPipe,
@@ -43,47 +49,34 @@
     preference_collate_fn,
     preprocess_preference_data,
 )
+from paddlenlp.utils.llm_utils import get_lora_target_modules
 from paddlenlp.utils.log import logger
 
 flash_mask_support_list = [LlamaForCausalLM, LlamaForCausalLMPipe]
 
 
 def main():
     """main"""
-    parser = PdArgumentParser((DPOModelArgument, DPODataArgument, DPOTrainingArguments))
-    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
-        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    parser = PdArgumentParser((DPOModelArgument, DPODataArgument, DPOTrainingArguments, DPOConfig))
+    if len(sys.argv) >= 2 and sys.argv[1].endswith(".json"):
+        model_args, data_args, training_args, dpo_config = parser.parse_json_file_and_cmd_lines()
     else:
-        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
-
-    training_args.print_config(model_args, "Model")
-    training_args.print_config(data_args, "Data")
-    if training_args.max_steps > 0:
-        training_args.num_train_epochs = 1
-    if data_args.autotuner_benchmark:
-        training_args.num_train_epochs = 1
-        training_args.max_steps = 5
-        training_args.do_train = True
-        training_args.do_export = False
-        training_args.do_predict = False
-        training_args.do_eval = False
-        training_args.overwrite_output_dir = True
-        training_args.load_best_model_at_end = False
-        training_args.report_to = []
-        training_args.save_strategy = IntervalStrategy.NO
-        training_args.evaluation_strategy = IntervalStrategy.NO
-    if data_args.benchmark:
-        training_args.do_train = True
-        training_args.do_export = False
-        training_args.do_predict = False
-        training_args.do_eval = False
-        training_args.overwrite_output_dir = True
-        training_args.load_best_model_at_end = False
-        training_args.save_strategy = IntervalStrategy.NO
-        training_args.evaluation_strategy = IntervalStrategy.NO
+        model_args, data_args, training_args, dpo_config = parser.parse_args_into_dataclasses()
 
     paddle.set_device(training_args.device)
     set_seed(training_args.seed)
+    if dpo_config.loss_type == "orpo":
+        dpo_config.reference_free = True
+        dpo_config.sft_loss_ratio = 1.0
+        dpo_config.loss_type = "or"
+        logger.info("orpo loss_type is equal to sft_loss + pref_loss_ratio * or_loss.")
+    if dpo_config.loss_type in ["or", "simpo"] and not dpo_config.reference_free:
+        dpo_config.reference_free = True
+        logger.warning(f"{dpo_config.loss_type} loss_type only supports reference_free. Set reference_free to True.")
+
+    training_args.print_config(model_args, "Model")
+    training_args.print_config(data_args, "Data")
+    training_args.print_config(dpo_config, "DPOConfig")
 
     logger.warning(
         f"Process rank: {training_args.local_rank}, device: {training_args.device}, world_size: "
@@ -116,51 +109,102 @@ def main():
         tensor_parallel_rank=training_args.tensor_parallel_rank,
         recompute_granularity=model_args.recompute_granularity,
         use_flash_attention=model_args.use_flash_attention,
-        tensor_parallel_output=True,
+        tensor_parallel_output=model_args.tensor_parallel_output,
     )
     if training_args.pipeline_parallel_degree > 1:
         raise ValueError("DPO does not support pipeline parallelism yet.")
-
-    if not data_args.autotuner_benchmark:
-        ref_model = AutoModelForCausalLM.from_pretrained(**model_kwargs)
-        config = AutoConfig.from_pretrained(**model_kwargs)
-        model = AutoModelForCausalLM.from_config(config)
-        model.set_state_dict(ref_model.state_dict())
+    if training_args.pipeline_parallel_degree > 1:
+        model_class = AutoModelForCausalLMPipe
+    else:
+        model_class = AutoModelForCausalLM
+    if not training_args.autotuner_benchmark or model_args.weight_quantize_algo is not None:
+        model = model_class.from_pretrained(**model_kwargs)
+        # for DPO save
+        model.config.dpo_config = None
+        if not dpo_config.reference_free and not dpo_config.lora:
+            config = AutoConfig.from_pretrained(**model_kwargs)
+            ref_model = model_class.from_config(config, dtype=dtype)
+            ref_model.set_state_dict(model.state_dict())
+        else:
+            ref_model = None
     else:
         config = AutoConfig.from_pretrained(**model_kwargs)
-        model = AutoModelForCausalLM.from_config(config)
-        ref_config = AutoConfig.from_pretrained(**model_kwargs)
-        ref_model = AutoModelForCausalLM.from_config(ref_config)
-        model.set_state_dict(ref_model.state_dict())
+        model = model_class.from_config(config, dtype=dtype)
+        if not dpo_config.reference_free and not dpo_config.lora:
+            ref_model = model_class.from_config(config, dtype=dtype)
+        else:
+            ref_model = None
 
     if model_args.flash_mask and not model.config.use_flash_attention:
         logger.warning("`flash_mask` must use with zero padding and flash attention.")
         model.config.use_flash_attention = True
 
     if model_args.flash_mask and not any(isinstance(model, cls) for cls in flash_mask_support_list):
         raise NotImplementedError(f"{model.__class__} not support flash mask.")
-    if training_args.sequence_parallel:
+
+    if model_args.sequence_parallel:
         register_sequence_parallel_allreduce_hooks(
-            model, training_args.gradient_accumulation_steps, training_args.fuse_sequence_parallel_allreduce
+            model, training_args.gradient_accumulation_steps, model_args.fuse_sequence_parallel_allreduce
         )
     if model_args.tokenizer_name_or_path is not None:
         tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
     else:
         tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
     # TODO: support chat template in next pr
-    # tokenizer.chat_template = None
+    tokenizer.chat_template = None
     logger.info("Loading model & tokenizer successfully !")
 
+    if dpo_config.lora:
+        if training_args.sharding_parallel_degree > 1:
+            assert (
+                "enable_stage1_overlap" not in training_args.sharding_parallel_config
+            ), "Currently not support enabling sharding_stage1_overlap in lora mode."
+        if model_args.lora_path is None:
+            target_modules = get_lora_target_modules(model)
+            if model_args.rslora_plus:
+                model_args.rslora = True
+                model_args.lora_plus_scale = 4
+                model_args.lora_alpha = 4
+            if model_args.weight_quantize_algo is not None:
+                if model_args.rslora or model_args.lora_plus_scale != 1.0:
+                    logger.info("Weight quantization is not supported in LoRA+ and RsLoRA.")
+            if model_args.lora_alpha == -1:
+                if model_args.rslora:
+                    model_args.lora_alpha = 4
+                else:
+                    model_args.lora_alpha = 2 * model_args.lora_rank
+            lora_config = LoRAConfig(
+                target_modules=target_modules,
+                r=model_args.lora_rank,
+                lora_alpha=2 * model_args.lora_rank if not model_args.rslora else 4,
+                rslora=model_args.rslora,
+                lora_plus_scale=model_args.lora_plus_scale,
+                tensor_parallel_degree=training_args.tensor_parallel_degree,
+                dtype=dtype,
+                base_model_name_or_path=model_args.model_name_or_path,
+                use_quick_lora=model_args.use_quick_lora,
+            )
+            model = LoRAModel(model, lora_config)
+        else:
+            model = LoRAModel.from_pretrained(model=model, lora_path=model_args.lora_path)
+
+        model.print_trainable_parameters()
+
     logger.info("Start to create dataset")
     trans_func = partial(preprocess_preference_data, tokenizer=tokenizer, data_args=data_args, model_args=model_args)
+    if data_args.lazy:
+        zero_padding_dataset = ZeroPaddingIterableDataset
+    else:
+        zero_padding_dataset = ZeroPaddingMapDataset
     if training_args.do_train and training_args.should_load_dataset:
         train_ds = load_dataset(
             "json",
             data_files=data_args.train_dataset_path,
+            lazy=data_args.lazy,
         )[0]
         logger.info("Creating train Zero Padding Data Stream. This may take a few minutes.")
         train_ds = (
-            ZeroPaddingMapDataset(
+            zero_padding_dataset(
                 train_ds.map(trans_func),
                 tokenizer=tokenizer,
                 max_length=data_args.max_seq_len,
@@ -176,10 +220,11 @@ def main():
         eval_ds = load_dataset(
             "json",
             data_files=data_args.dev_dataset_path,
+            lazy=data_args.lazy,
         )[0]
         logger.info("Creating dev Zero Padding Data Stream. This may take a few minutes.")
         eval_ds = (
-            ZeroPaddingMapDataset(
+            zero_padding_dataset(
                 eval_ds.map(trans_func),
                 tokenizer=tokenizer,
                 max_length=data_args.max_seq_len,
@@ -194,6 +239,7 @@ def main():
     trainer = DPOTrainer(
         model=model,
         ref_model=ref_model,
+        dpo_config=dpo_config,
         args=training_args,
         train_dataset=train_ds,
         eval_dataset=eval_ds,
@@ -202,17 +248,18 @@ def main():
             preference_collate_fn,
             max_seq_len=data_args.max_seq_len,
         ),
+        ignore_eos_token=True,
     )
 
     if training_args.do_train:
         train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
 
-        if not data_args.autotuner_benchmark and not data_args.benchmark:
+        if not training_args.autotuner_benchmark and not training_args.benchmark:
             trainer.save_model(merge_tensor_parallel=training_args.tensor_parallel_degree > 1)
             trainer.log_metrics("train", train_result.metrics)
             trainer.save_metrics("train", train_result.metrics)
             trainer.save_state()
-        if data_args.benchmark:
+        if training_args.benchmark:
             total_effective_tokens, total_tokens = calculate_effective_tokens(
                 training_args, train_ds, data_args.max_seq_len
             )

llm/config/baichuan/dpo_argument.json

@@ -24,7 +24,6 @@
     "disable_tqdm": true,
     "load_best_model_at_end": true,
     "tensor_parallel_degree": 8,
-    "sharding_parallel_degree": 1,
     "sharding": "stage1",
     "use_flash_attention": true,
     "recompute": false,