diff --git a/examples/dpo_train.py b/examples/dpo_train.py
new file mode 100644
index 000000000..ae3d0f819
--- /dev/null
+++ b/examples/dpo_train.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 7/4/2024 20:31
+# @Author  : Yu Li
+# @Site    : 
+# @File    : dpo_train.py
+# 0. imports
+import logging
+import os
+import sys
+
+sys.path.remove(os.path.abspath(os.path.dirname(sys.argv[0])))
+from dataclasses import dataclass, field
+from typing import Optional
+
+from transformers import HfArgumentParser, pipeline, AutoTokenizer
+
+from lmflow.args import (
+    ModelArguments,
+    DatasetArguments,
+    AutoArguments,
+)
+import torch
+from lmflow.datasets.dataset import Dataset
+from lmflow.models.auto_model import AutoModel
+from lmflow.pipeline.auto_pipeline import AutoPipeline
+
+if __name__ == "__main__":
+    # Parses arguments
+    pipeline_name = "dpo_aligner"
+    PipelineArguments = AutoArguments.get_pipeline_args_class(pipeline_name)
+    parser = HfArgumentParser((
+        ModelArguments,
+        DatasetArguments,
+        PipelineArguments,
+    ))
+
+    model_args, data_args, pipeline_args = parser.parse_args_into_dataclasses()
+
+    # Initializes pipeline, dataset and model for reward training
+    aligner = AutoPipeline.get_pipeline(
+        pipeline_name=pipeline_name,
+        model_args=model_args,
+        data_args=data_args,
+        pipeline_args=pipeline_args,
+    )
+    model = AutoModel.get_model(model_args)
+
+    # Aligns model with rewards
+    aligned_model = aligner.align(
+        model=model,
+        dataset=None,
+        reward_model=None
+    )
\ No newline at end of file
diff --git a/scripts/run_dpo_align.sh b/scripts/run_dpo_align.sh
new file mode 100644
index 000000000..6f80883c0
--- /dev/null
+++ b/scripts/run_dpo_align.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+# Please run this script under ${project_id} in project directory of
+
+# Parses arguments
+model_name_or_path=meta-llama/Llama-2-7b-hf
+dataset_path=data/stack-exchange-paired-dummy
+output_dir=output_models/dpo
+deepspeed_args="--master_port=11000"
+# specify gpus/single gpu here by 
+# `--include localhost:0,1` or `--include localhost:0`
+
+while [[ $# -ge 1 ]]; do
+  key="$1"
+  case ${key} in
+    -m|--model_name_or_path)
+      model_name_or_path="$2"
+      shift
+      ;;
+    -d|--dataset_path)
+      dataset_path="$2"
+      shift
+      ;;
+    -o|--output_lora_path)
+      output_dir="$2"
+      shift
+      ;;
+    --deepspeed_args)
+      deepspeed_args="$2"
+      shift
+      ;;
+    *)
+      echo "error: unknown option \"${key}\"" 1>&2
+      exit 1
+  esac
+  shift
+done
+exp_id=dpo
+project_dir=$(cd "$(dirname $0)"/..; pwd)
+log_dir=${project_dir}/log/${exp_id}
+mkdir -p ${output_dir} ${log_dir}
+
+deepspeed ${deepspeed_args} \
+  examples/dpo_train.py \
+    --model_name_or_path ${model_name_or_path} \
+    --dataset_path ${dataset_path} \
+    --output_dir ${output_dir} \
+    --max_steps 200 \
+    --learning_rate 1e-4 \
+    --use_lora 1 \
+    --lora_r 8 \
+    --sanity_check True \
+    --save_aggregated_lora 0\
+    --logging_steps 20 \
+    | tee ${log_dir}/train.log \
+    2> ${log_dir}/train.err
diff --git a/src/lmflow/args.py b/src/lmflow/args.py
index f8ebf737d..a50e824ec 100644
--- a/src/lmflow/args.py
+++ b/src/lmflow/args.py
@@ -31,7 +31,7 @@ class ModelArguments:
     """
     Define a class ModelArguments using the dataclass decorator. 
     The class contains several optional parameters that can be used to configure a model. 
-    
+
     model_name_or_path : str
         a string representing the path or name of a pretrained
         model checkpoint for weights initialization. If None, a model will be trained from scratch.
@@ -39,15 +39,15 @@ class ModelArguments:
     model_type :  str
         a string representing the type of model to use if training from
         scratch. If not provided, a pretrained model will be used.
-    
+
     config_overrides :  str
         a string representing the default config settings to override
         when training a model from scratch.
-    
+
     config_name : str
         a string representing the name or path of the pretrained config to
         use, if different from the model_name_or_path.
-    
+
     tokenizer_name :  str
         a string representing the name or path of the pretrained tokenizer
         to use, if different from the model_name_or_path.
@@ -77,6 +77,8 @@ class ModelArguments:
         enough.
     use_int8 : bool
         a boolean indicating whether to load int8 quantization for inference.
+    load_in_4bit : bool
+        whether to load the model in 4bit
     """
 
     model_name_or_path: Optional[str] = field(
@@ -151,7 +153,7 @@ class ModelArguments:
             )
         },
     )
-    trust_remote_code : bool = field(
+    trust_remote_code: bool = field(
         default=False,
         metadata={
             "help": (
@@ -180,12 +182,12 @@ class ModelArguments:
     bits: int = field(
         default=4,
         metadata={"help": "The number of bits for quantization.",
-                  "choices": [4, 8],},
+                  "choices": [4, 8], },
     )
     quant_type: str = field(
         default='nf4',
         metadata={"help": "The quantization type for quantization.",
-                  "choices": ["nf4", "fp4"],},
+                  "choices": ["nf4", "fp4"], },
     )
     double_quant: bool = field(
         default=True,
@@ -197,11 +199,12 @@ class ModelArguments:
     )
     lora_alpha: int = field(
         default=32,
-        metadata={"help": "Merging ratio between the fine-tuned model and the original. This is controlled by a parameter called alpha in the paper."},
+        metadata={
+            "help": "Merging ratio between the fine-tuned model and the original. This is controlled by a parameter called alpha in the paper."},
     )
     lora_target_modules: List[str] = field(
         default=None, metadata={"help": "Pretrained config name or path if not the same as model_name",
-                              }
+                                }
     )
     lora_dropout: float = field(
         default=0.1,
@@ -210,7 +213,7 @@ class ModelArguments:
     save_aggregated_lora: bool = field(
         default=False,
         metadata={"help": "Whether to save aggregated lora."},
-        )
+    )
     use_ram_optimized_load: bool = field(
         default=True,
         metadata={"help": "Whether use disk mapping when memory is not enough."}
@@ -233,7 +236,7 @@ class ModelArguments:
         }
     )
     do_rope_scaling: bool = field(
-        default = False,
+        default=False,
         metadata={
             "help": (
                 "whether do ROPE scaling for llama model."
@@ -242,7 +245,7 @@ class ModelArguments:
                 "NTK_scaling credits to the Reddit users /u/bloc97 and /u/emozilla."
                 "https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/"
             )
-        }   
+        }
     )
     rope_pi_ratio: int = field(
         default=1,
@@ -264,6 +267,12 @@ class ModelArguments:
         default=False,
         metadata={"help": "whether to load int8 quantization for inference"}
     )
+    load_in_4bit: Optional[bool] = field(
+        default=True,
+        metadata={
+            "help": "whether to load the model in 4bit"
+        },
+    )
 
     def __post_init__(self):
         if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None):
@@ -345,12 +354,13 @@ class VisModelArguments(ModelArguments):
         metadata={"help": "Path to pretrained model."},
     )
 
+
 @dataclass
 class DatasetArguments:
     """
-    Define a class DatasetArguments using the dataclass decorator. 
-    The class contains several optional parameters that can be used to configure a dataset for a language model. 
-    
+    Define a class DatasetArguments using the dataclass decorator.
+    The class contains several optional parameters that can be used to configure a dataset for a language model.
+
 
     dataset_path : str
         a string representing the path of the dataset to use.
@@ -374,26 +384,26 @@ class DatasetArguments:
         a string representing the path to the input evaluation data file to evaluate the perplexity on (a text file).
 
     max_train_samples : int
-        an integer indicating the maximum number of training examples to use for debugging or quicker training. 
+        an integer indicating the maximum number of training examples to use for debugging or quicker training.
         If set, the training dataset will be truncated to this number.
 
     max_eval_samples: int
-        an integer indicating the maximum number of evaluation examples to use for debugging or quicker training. 
+        an integer indicating the maximum number of evaluation examples to use for debugging or quicker training.
         If set, the evaluation dataset will be truncated to this number.
 
     streaming : bool
         a boolean indicating whether to enable streaming mode.
 
     block_size: int
-        an integer indicating the optional input sequence length after tokenization. The training dataset will be 
+        an integer indicating the optional input sequence length after tokenization. The training dataset will be
         truncated in blocks of this size for training.
-        
+
     train_on_prompt: bool
         a boolean indicating whether to train on prompt for conversation datasets such as ShareGPT.
-        
+
     disable_conversation_bos_token: bool
         a boolean indicating whether to disable the bos token for conversation datasets.
-        
+
     disable_conversation_eos_token: bool
         a boolean indicating whether to disable the eos token for conversation datasets.
 
@@ -402,7 +412,7 @@ class DatasetArguments:
     `keep_linebreaks`, and `prompt_structure`.
 
     The field function is used to set default values and provide help messages for each parameter. The Optional type hint is
-    used to indicate that a parameter is optional. The metadata argument is used to provide additional information about 
+    used to indicate that a parameter is optional. The metadata argument is used to provide additional information about
     each parameter, such as a help message.
     """
 
@@ -480,21 +490,16 @@ class DatasetArguments:
         }
     )
     disable_group_texts: bool = field(
-        default=True,
+        default=False,
         metadata={
             "help": (
-                "Whether we disable group of original samples together to"
-                " generate sample sequences of length `block_size`"
-                " By Default, it is True, which means the long samples"
-                " are truncated to `block_size` tokens"
-                " and short samples are padded to `block_size` tokens."
-                " If set to False, we group every 1000 tokenized"
-                " sequences together, divide them into"
-                " [{total_num_tokens} / {block_size}] sequences,"
-                " each with `block_size` tokens"
-                " (the remaining tokens are ommited."
-                " This group text behavior is useful"
-                " for continual pretrain or pretrain."
+                "Whether we group original samples together to generate sample"
+                " sequences of length `block_size`. By default, we group every"
+                " 1000 tokenized sequences together, divide them into "
+                " [{total_num_tokens} / {block_size}] sequences, each with"
+                " `block_size` tokens (the remaining tokens are ommited."
+                " If this flag is set to True, we only group 1 tokenized"
+                " sequence, i.e. cutting long sequence into chunks."
             )
         },
     )
@@ -552,7 +557,6 @@ class MultiModalDatasetArguments(DatasetArguments):
     )
 
 
-
 @dataclass
 class FinetunerArguments(TrainingArguments):
     """
@@ -624,28 +628,28 @@ class EvaluatorArguments:
     mixed_precision : str, choice from ["bf16","fp16"].
         mixed precision mode, whether to use bf16 or fp16
 
-    deepspeed : 
+    deepspeed :
         Enable deepspeed and pass the path to deepspeed json config file (e.g. ds_config.json) or an already
         loaded json file as a dict
-        
+
     temperature : float
         An argument of model.generate in huggingface to control the diversity of generation.
-        
+
     repetition_penalty : float
         An argument of model.generate in huggingface to penalize repetitions.
     """
     local_rank: int = field(
         default=-1,
         metadata={"help": "For distributed training: local_rank"
-        }
+                  }
     )
 
     random_shuffle: Optional[bool] = field(
-        default=False, 
+        default=False,
         metadata={"help": ""
-        }
+                  }
     )
-    
+
     use_wandb: Optional[bool] = field(
         default=False,
         metadata={
@@ -672,7 +676,7 @@ class EvaluatorArguments:
             "help": (
                 "mixed precision mode, whether to use bf16 or fp16"
             ),
-            "choices": ["bf16","fp16"],
+            "choices": ["bf16", "fp16"],
         },
     )
     deepspeed: Optional[str] = field(
@@ -749,22 +753,23 @@ class EvaluatorArguments:
     use_accelerator_for_evaluator: bool = field(
         default=False, metadata={"help": "Whether to use Huggingface Accelerator instead of Deepspeed"},
     )
-        
+
     temperature: float = field(
         default=0,
         metadata={"help": "Temperature during inference."},
     )
-    
+
     repetition_penalty: float = field(
         default=1,
         metadata={"help": "Repetition_penalty during inference."},
     )
-        
+
     max_new_tokens: int = field(
         default=100,
         metadata={"help": "Maximum length during inference."},
     )
-    
+
+
 @dataclass
 class InferencerArguments:
     """
@@ -781,10 +786,10 @@ class InferencerArguments:
         loaded json file as a dict
     mixed_precision : str, choice from ["bf16","fp16"].
         mixed precision mode, whether to use bf16 or fp16
-    
+
     temperature : float
         An argument of model.generate in huggingface to control the diversity of generation.
-        
+
     repetition_penalty : float
         An argument of model.generate in huggingface to penalize repetitions.
     """
@@ -798,19 +803,19 @@ class InferencerArguments:
     local_rank: int = field(
         default=-1,
         metadata={"help": "For distributed training: local_rank"
-        },
+                  },
     )
-        
+
     temperature: float = field(
         default=0.0,
         metadata={"help": "Temperature during inference."},
     )
-    
+
     repetition_penalty: float = field(
         default=1,
         metadata={"help": "Repetition_penalty during inference."},
     )
-        
+
     max_new_tokens: int = field(
         default=100,
         metadata={"help": "Maximum length during inference."},
@@ -839,7 +844,7 @@ class InferencerArguments:
             "help": (
                 "mixed precision mode, whether to use bf16 or fp16"
             ),
-            "choices": ["bf16","fp16"],
+            "choices": ["bf16", "fp16"],
         },
     )
     do_sample: Optional[bool] = field(
@@ -926,6 +931,7 @@ class RaftAlignerArguments(TrainingArguments):
         },
     )
 
+
 @dataclass
 class BenchmarkingArguments:
     dataset_name: Optional[str] = field(
@@ -939,11 +945,154 @@ class BenchmarkingArguments:
         metadata={
             "help": "the metric the model will be evaluated on",
             "choices": ["acc", "acc_norm", "bleu", "chrf", "em", "f1", "ppl", \
-                "ter", "r@1", "r@2", "mrr", "mc1", "mc2", "word_perplexity", \
-                    "byte_perplexity", "bits_per_byte"],
+                        "ter", "r@1", "r@2", "mrr", "mc1", "mc2", "word_perplexity", \
+                        "byte_perplexity", "bits_per_byte"],
+        },
+    )
+
+
+@dataclass
+class DPOAlignerArguments:
+    """
+    The arguments for the DPO training script.
+    """
+    local_rank: int = field(
+        default=-1,
+        metadata={"help": "For distributed training: local_rank"
+                  },
+    )
+    # data parameters
+    beta: Optional[float] = field(
+        default=0.1,
+        metadata={
+            "help": "the beta parameter for DPO loss"
+        }
+    )
+    # # training parameters
+    learning_rate: Optional[float] = field(
+        default=5e-4,
+        metadata={
+            "help": "optimizer learning rate"
+        }
+    )
+    lr_scheduler_type: Optional[str] = field(
+        default="cosine",
+        metadata={
+            "help": "the lr scheduler type"
+        }
+    )
+    warmup_steps: Optional[int] = field(
+        default=100, metadata={
+            "help": "the number of warmup steps"
+        }
+    )
+    weight_decay: Optional[float] = field(
+        default=0.05, metadata={
+            "help": "the weight decay"
+        }
+    )
+    optimizer_type: Optional[str] = field(
+        default="paged_adamw_32bit",
+        metadata={
+            "help": "the optimizer type"
+        }
+    )
+
+    per_device_train_batch_size: Optional[int] = field(
+        default=4,
+        metadata={
+            "help": "train batch size per device"
+        }
+    )
+    per_device_eval_batch_size: Optional[int] = field(
+        default=1, metadata={
+            "help": "eval batch size per device"
+        }
+    )
+    gradient_accumulation_steps: Optional[int] = field(
+        default=4,
+        metadata={
+            "help": "the number of gradient accumulation steps"
+        },
+    )
+    gradient_checkpointing: Optional[bool] = field(
+        default=True,
+        metadata={
+            "help": "whether to use gradient checkpointing"
         },
     )
 
+    gradient_checkpointing_use_reentrant: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "whether to use reentrant for gradient checkpointing"
+        },
+    )
+    max_prompt_length: Optional[int] = field(
+        default=512,
+        metadata={
+            "help": "the maximum prompt length"
+        },
+    )
+    max_length: Optional[int] = field(
+        default=1024,
+        metadata={
+            "help": "the maximum sequence length"
+        },
+    )
+    max_steps: Optional[int] = field(
+        default=1000,
+        metadata={
+            "help": "max number of training steps"
+        },
+    )
+    logging_steps: Optional[int] = field(
+        default=10,
+        metadata={
+            "help": "the logging frequency"
+        },
+    )
+    save_steps: Optional[int] = field(
+        default=100,
+        metadata={
+            "help": "the saving frequency"
+        },
+    )
+    eval_steps: Optional[int] = field(
+        default=100,
+        metadata={
+            "help": "the evaluation frequency"
+        },
+    )
+    output_dir: Optional[str] = field(
+        default="./results",
+        metadata={
+            "help": "the output directory"
+        },
+    )
+    log_freq: Optional[int] = field(
+        default=1,
+        metadata={
+            "help": "the logging frequency"
+        },
+    )
+    sanity_check: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "only train on 1000 samples"
+        }
+    )
+    report_to: Optional[str] = field(
+        default="wandb",
+        metadata={
+            "help": 'The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`,'
+                    '`"comet_ml"`, `"mlflow"`, `"neptune"`, `"tensorboard"`,`"clearml"` and `"wandb"`. '
+                    'Use `"all"` to report to all integrations installed, `"none"` for no integrations.'
+        },
+    )
+    seed: Optional[int] = field(
+        default=0, metadata={"help": "Random seed that will be set at the beginning of training."}
+    )
 
 
 PIPELINE_ARGUMENT_MAPPING = {
@@ -951,6 +1100,7 @@ class BenchmarkingArguments:
     "evaluator": EvaluatorArguments,
     "inferencer": InferencerArguments,
     "raft_aligner": RaftAlignerArguments,
+    "dpo_aligner": DPOAlignerArguments,
 }
 
 
@@ -958,5 +1108,6 @@ class AutoArguments:
     """
     Automatically choose arguments from FinetunerArguments or EvaluatorArguments.
     """
+
     def get_pipeline_args_class(pipeline_name: str):
         return PIPELINE_ARGUMENT_MAPPING[pipeline_name]
diff --git a/src/lmflow/pipeline/auto_pipeline.py b/src/lmflow/pipeline/auto_pipeline.py
index 524338e8d..1e74abe8d 100644
--- a/src/lmflow/pipeline/auto_pipeline.py
+++ b/src/lmflow/pipeline/auto_pipeline.py
@@ -17,11 +17,12 @@ def is_package_version_at_least(package_name, min_version):
 from lmflow.pipeline.evaluator import Evaluator
 from lmflow.pipeline.finetuner import Finetuner
 from lmflow.pipeline.inferencer import Inferencer
-
+from lmflow.pipeline.dpo_aligner import DPOAligner
 PIPELINE_MAPPING = {
     "evaluator": Evaluator,
     "finetuner": Finetuner,
     "inferencer": Inferencer,
+    "dpo_aligner": DPOAligner,
 }
 
 if not is_package_version_at_least('transformers', '4.35.0'):
diff --git a/src/lmflow/pipeline/dpo_aligner.py b/src/lmflow/pipeline/dpo_aligner.py
new file mode 100644
index 000000000..6c0a5347e
--- /dev/null
+++ b/src/lmflow/pipeline/dpo_aligner.py
@@ -0,0 +1,158 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 7/4/2024 21:12
+# @Author  : Yu Li
+# @Site    :
+# @File    : dpo_pipeline.py
+from lmflow.pipeline.base_aligner import BaseAligner
+from trl import DPOTrainer
+from typing import Dict, Optional
+from datasets import Dataset, load_dataset
+import os
+from peft import LoraConfig
+from transformers import TrainingArguments
+
+
+def get_paired_dataset(
+        data_dir: str = "rl",
+        data_root: str = None,
+        sanity_check: bool = False,
+        cache_dir: Optional[str] = None,
+        num_proc=24,
+) -> Dataset:
+    """Load dataset and convert it to the necessary format.
+
+    The dataset is converted to a dictionary with the following structure:
+    {
+        'prompt': List[str],
+        'chosen': List[str],
+        'rejected': List[str],
+    }
+
+    Prompts are structured as follows:
+      "Question: " + <prompt> + "\n\nAnswer: "
+    """
+    dataset = load_dataset(
+        data_root,
+        split="train",
+        cache_dir=cache_dir,
+        data_dir=data_dir,
+    )
+    original_columns = dataset.column_names
+
+    if sanity_check:
+        dataset = dataset.select(range(min(len(dataset), 1000)))
+
+    def return_prompt_and_responses(samples) -> Dict[str, str]:
+        return {
+            "prompt": ["Question: " + question + "\n\nAnswer: " for question in samples["question"]],
+            "chosen": samples["response_j"],
+            "rejected": samples["response_k"],
+        }
+
+    return dataset.map(
+        return_prompt_and_responses,
+        batched=True,
+        num_proc=num_proc,
+        remove_columns=original_columns,
+    )
+
+
+class DPOAligner(BaseAligner):
+    def __init__(self, model_args, data_args, aligner_args):
+        self.model_args = model_args
+        self.data_args = data_args
+        self.aligner_args = aligner_args
+
+    def _initialize_trainer(self, model, tokenizer):
+        peft_config = LoraConfig(
+            r=self.model_args.lora_r,
+            lora_alpha=self.model_args.lora_alpha,
+            lora_dropout=self.model_args.lora_dropout,
+            target_modules=[
+                "q_proj",
+                "v_proj",
+                "k_proj",
+                "out_proj",
+                "fc_in",
+                "fc_out",
+                "wte",
+            ],
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+        training_args = TrainingArguments(
+            per_device_train_batch_size=self.aligner_args.per_device_train_batch_size,
+            per_device_eval_batch_size=self.aligner_args.per_device_eval_batch_size,
+            max_steps=self.aligner_args.max_steps,
+            logging_steps=self.aligner_args.logging_steps,
+            save_steps=self.aligner_args.save_steps,
+            gradient_accumulation_steps=self.aligner_args.gradient_accumulation_steps,
+            gradient_checkpointing=self.aligner_args.gradient_checkpointing,
+            learning_rate=self.aligner_args.learning_rate,
+            evaluation_strategy="steps",
+            eval_steps=self.aligner_args.eval_steps,
+            output_dir=self.aligner_args.output_dir,
+            report_to=self.aligner_args.report_to,
+            lr_scheduler_type=self.aligner_args.lr_scheduler_type,
+            warmup_steps=self.aligner_args.warmup_steps,
+            optim=self.aligner_args.optimizer_type,
+            bf16=True,
+            remove_unused_columns=False,
+            run_name="dpo_llama2",
+            ddp_find_unused_parameters=False,
+            # gradient_checkpointing_kwargs=dict(use_reentrant=self.aligner_args.gradient_checkpointing_use_reentrant),
+            seed=self.aligner_args.seed,
+        )
+        dpo_trainer = DPOTrainer(
+            model,
+            ref_model=None,
+            args=training_args,
+            beta=self.aligner_args.beta,
+            train_dataset=self.train_dataset,
+            eval_dataset=self.eval_dataset,
+            tokenizer=tokenizer,
+            peft_config=peft_config,
+            max_prompt_length=self.aligner_args.beta,
+            max_length=self.aligner_args.max_length,
+        )
+        return dpo_trainer
+
+    def _load_dataset(self):
+        # load training set
+        self.train_dataset = get_paired_dataset(data_root=self.data_args.dataset_path,
+                                                data_dir="data/rl",
+                                                sanity_check=self.aligner_args.sanity_check)
+        self.train_dataset = self.train_dataset.filter(
+            lambda x: len(x["prompt"]) + len(x["chosen"]) <= self.aligner_args.max_length
+                      and len(x["prompt"]) + len(x["rejected"]) <= self.aligner_args.max_length
+        )
+        # load evaluation set
+        self.eval_dataset = get_paired_dataset(data_root=self.data_args.dataset_path,
+                                               data_dir="data/evaluation",
+                                               sanity_check=True)
+        self.eval_dataset = self.eval_dataset.filter(
+            lambda x: len(x["prompt"]) + len(x["chosen"]) <= self.aligner_args.max_length
+                      and len(x["prompt"]) + len(x["rejected"]) <= self.aligner_args.max_length
+        )
+
+    def align(self, model, dataset, reward_model):
+        tokenizer = model.get_tokenizer()
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        self._load_dataset()
+
+        wrapped_model = model
+        model = model.get_backend_model()
+
+        dpo_trainer = self._initialize_trainer(model, tokenizer)
+        dpo_trainer.train()
+        dpo_trainer.save_model(self.aligner_args.output_dir)
+
+        # 7. save
+        output_dir = os.path.join(self.aligner_args.output_dir, "final_checkpoint")
+        dpo_trainer.model.save_pretrained(output_dir)
+
+
+
+