feat: add Direct Preference Optimization (DPO) training support

ashrane111 · ashrane111 · commit 1df8ad4a0236 · 2026-02-19T22:40:20.000-06:00
Add DPO as a new alignment method alongside the existing SFT pipeline.
DPO enables fine-tuning LLMs using human preference pairs (chosen vs
rejected responses) without requiring a separate reward model.

New components:
- PreferenceDataset: dataset class for prompt/chosen/rejected triplets
- PreferenceDataCollator: tokenizer that prepares paired sequences
- DPOTrainer: PyTorch Lightning trainer with DPO loss and reference model
- CausalModel.dpo_finetune(): user-facing API for DPO training

All components integrate via the existing registry pattern and work with
every model variant (full, LoRA, INT8, LoRA+INT8, LoRA+Kbit).

Includes tests and an example script.

README.md is not modified by this change.
diff --git a/examples/features/dpo/dpo_finetune.py b/examples/features/dpo/dpo_finetune.py
@@ -0,0 +1,67 @@
+"""Minimal example showing how to align a model using Direct Preference
+Optimization (DPO) with xTuring.
+
+DPO fine-tunes a language model using pairs of preferred and dispreferred
+responses so that the model learns to produce outputs that match human
+preferences without requiring a separate reward model.
+"""
+
+from pathlib import Path
+
+from xturing.datasets.preference_dataset import PreferenceDataset
+from xturing.models import BaseModel
+
+OUTPUT_DIR = Path(__file__).parent / "dpo_weights"
+
+
+def main():
+    # Build a small preference dataset. Each sample needs a prompt, a chosen
+    # (preferred) response, and a rejected (dispreferred) response.
+    preference_data = {
+        "prompt": [
+            "Explain quantum computing in simple terms.",
+            "What is the capital of France?",
+            "How do I make pasta?",
+            "What causes rain?",
+        ],
+        "chosen": [
+            "Quantum computing uses qubits that can be 0, 1, or both at once, "
+            "letting it solve certain problems much faster than regular computers.",
+            "The capital of France is Paris.",
+            "Boil salted water, cook pasta until al dente, then drain and toss "
+            "with your favorite sauce.",
+            "Rain forms when water evaporates, rises, cools into clouds, and "
+            "falls back as droplets when clouds become saturated.",
+        ],
+        "rejected": [
+            "Quantum computing is basically magic computers that can do "
+            "everything instantly.",
+            "France doesn't have a capital, it's a collective.",
+            "Just put some noodles in a microwave with ketchup.",
+            "Rain happens because the sky is sad.",
+        ],
+    }
+
+    dataset = PreferenceDataset(preference_data)
+
+    # Initialise a model with a LoRA adapter. DPO works with any model
+    # variant, but LoRA is recommended to keep memory usage low since DPO
+    # requires a frozen reference model in addition to the policy model.
+    model = BaseModel.create("qwen3_0_6b_lora")
+
+    # Run DPO fine-tuning. The beta parameter controls how strongly the model
+    # is penalised for deviating from the reference policy (higher = more
+    # conservative).
+    model.dpo_finetune(dataset=dataset, beta=0.1)
+
+    # Verify the aligned model generates reasonable output.
+    output = model.generate(texts=["Explain gravity in simple terms."])
+    print(f"Generated output: {output}")
+
+    # Save the fine-tuned adapter weights.
+    model.save(str(OUTPUT_DIR))
+    print(f"Saved DPO fine-tuned weights to {OUTPUT_DIR}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/xturing/datasets/__init__.py b/src/xturing/datasets/__init__.py
@@ -3,9 +3,11 @@
     InstructionDataset,
     InstructionDatasetMeta,
 )
+from xturing.datasets.preference_dataset import PreferenceDataset, PreferenceDatasetMeta
 from xturing.datasets.text2image_dataset import Text2ImageDataset
 from xturing.datasets.text_dataset import TextDataset, TextDatasetMeta
 
 BaseDataset.add_to_registry(TextDataset.config_name, TextDataset)
 BaseDataset.add_to_registry(InstructionDataset.config_name, InstructionDataset)
 BaseDataset.add_to_registry(Text2ImageDataset.config_name, Text2ImageDataset)
+BaseDataset.add_to_registry(PreferenceDataset.config_name, PreferenceDataset)
diff --git a/src/xturing/datasets/preference_dataset.py b/src/xturing/datasets/preference_dataset.py
@@ -0,0 +1,95 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+
+from xturing.datasets.base import BaseDataset
+
+
+@dataclass
+class PreferenceDatasetMeta:
+    """Metadata for preference datasets used in DPO training."""
+
+
+class PreferenceDataset(BaseDataset):
+    """Dataset for Direct Preference Optimization (DPO) training.
+
+    Each sample contains a prompt, a chosen (preferred) response, and a
+    rejected (dispreferred) response. The dataset must have exactly three
+    columns: ``prompt``, ``chosen``, and ``rejected``.
+
+    Args:
+        path: A local directory saved with ``datasets.save_to_disk``, a path
+            to a ``.jsonl`` file, a HuggingFace ``Dataset``/``DatasetDict``,
+            or a plain dictionary with the required keys.
+    """
+
+    config_name: str = "preference_dataset"
+
+    def __init__(self, path: Union[str, Path, HFDataset, DatasetDict, dict]):
+        if isinstance(path, HFDataset) or isinstance(path, DatasetDict):
+            self.data = path
+        elif isinstance(path, dict):
+            self.data = {"train": HFDataset.from_dict(path)}
+        else:
+            path = Path(path)
+            assert path.exists(), "path does not exist"
+            if path.is_dir():
+                self.data = load_from_disk(str(path))
+            elif path.suffix == ".jsonl":
+                self.data = {"train": HFDataset.from_dict(self._from_jsonl(path))}
+            else:
+                raise ValueError(
+                    f"Unsupported file format: {path.suffix}. Use a directory or .jsonl file."
+                )
+
+        self._validate()
+        self._meta = PreferenceDatasetMeta()
+
+    def _from_jsonl(self, path: Path):
+        data = {
+            "prompt": [],
+            "chosen": [],
+            "rejected": [],
+        }
+        try:
+            for line in open(path):
+                json_line = json.loads(line)
+                data["prompt"].append(json_line["prompt"])
+                data["chosen"].append(json_line["chosen"])
+                data["rejected"].append(json_line["rejected"])
+        except KeyError:
+            raise ValueError(
+                "The jsonl file should have keys: prompt, chosen, and rejected"
+            )
+        return data
+
+    def _validate(self):
+        assert "train" in self.data, "The dataset should have a train split"
+        assert (
+            "prompt" in self.data["train"].column_names
+        ), "The dataset should have a column named prompt"
+        assert (
+            "chosen" in self.data["train"].column_names
+        ), "The dataset should have a column named chosen"
+        assert (
+            "rejected" in self.data["train"].column_names
+        ), "The dataset should have a column named rejected"
+        assert (
+            len(self.data["train"].column_names) == 3
+        ), "The dataset should have only three columns: prompt, chosen, and rejected"
+
+    def __len__(self):
+        return len(self.data["train"])
+
+    def __iter__(self):
+        return iter(self.data["train"])
+
+    def __getitem__(self, idx):
+        return self.data["train"][idx]
+
+    def save(self, path):
+        return self.data["train"].save_to_disk(path)
diff --git a/src/xturing/models/causal.py b/src/xturing/models/causal.py
@@ -12,11 +12,13 @@
 from xturing.config.config_data_classes import FinetuningConfig, GenerationConfig
 from xturing.config.read_config import load_config
 from xturing.datasets.instruction_dataset import InstructionDataset
+from xturing.datasets.preference_dataset import PreferenceDataset
 from xturing.datasets.text_dataset import TextDataset
 from xturing.engines.base import BaseEngine
 from xturing.models import BaseModel
 from xturing.preprocessors.base import BasePreprocessor
 from xturing.trainers.base import BaseTrainer
+from xturing.trainers.dpo_trainer import DPOTrainer
 from xturing.trainers.lightning_trainer import LightningTrainer
 from xturing.utils.logging import configure_logger
 from xturing.utils.prompt import OpenAICreateChatPrompt, OpenAICreatePrompt, Prompt
@@ -118,6 +120,54 @@ def finetune(
         trainer = self._make_trainer(dataset, logger)
         trainer.fit()
 
+    def _make_dpo_collate_fn(self, dataset: PreferenceDataset):
+        return BasePreprocessor.create(
+            dataset.config_name,
+            self.engine.tokenizer,
+            int(self.finetuning_args.max_length),
+            dataset.meta,
+        )
+
+    def _make_dpo_trainer(
+        self,
+        dataset: PreferenceDataset,
+        beta: float = 0.1,
+        logger: Union[Logger, Iterable[Logger], bool] = True,
+    ):
+        return BaseTrainer.create(
+            DPOTrainer.config_name,
+            self.engine,
+            dataset,
+            self._make_dpo_collate_fn(dataset),
+            int(self.finetuning_args.num_train_epochs),
+            int(self.finetuning_args.batch_size),
+            float(self.finetuning_args.learning_rate),
+            self.finetuning_args.optimizer_name,
+            beta,
+            logger=logger,
+        )
+
+    def dpo_finetune(
+        self,
+        dataset: PreferenceDataset,
+        beta: float = 0.1,
+        logger: Union[Logger, Iterable[Logger], bool] = True,
+    ):
+        """Fine-tune the model using Direct Preference Optimization (DPO).
+
+        Args:
+            dataset: A :class:`PreferenceDataset` containing prompt, chosen,
+                and rejected columns.
+            beta: Temperature parameter for DPO.  Higher values keep the model
+                closer to the reference policy.
+            logger: PyTorch Lightning logger(s) for tracking training metrics.
+        """
+        assert (
+            dataset.config_name == "preference_dataset"
+        ), "Please provide a PreferenceDataset for DPO training"
+        trainer = self._make_dpo_trainer(dataset, beta, logger)
+        trainer.fit()
+
     def _generate_from_iterable(
         self, data_iterator: Iterable, do_tokenization=False, show_tqdm_bar=True
     ):
diff --git a/src/xturing/preprocessors/__init__.py b/src/xturing/preprocessors/__init__.py
@@ -1,3 +1,4 @@
 from xturing.preprocessors.base import BasePreprocessor
 from xturing.preprocessors.instruction_collator import InstructionDataCollator
+from xturing.preprocessors.preference_collator import PreferenceDataCollator
 from xturing.preprocessors.text_collator import TextDataCollator
diff --git a/src/xturing/preprocessors/base.py b/src/xturing/preprocessors/base.py
@@ -1,4 +1,5 @@
 from xturing.preprocessors.instruction_collator import InstructionDataCollator
+from xturing.preprocessors.preference_collator import PreferenceDataCollator
 from xturing.preprocessors.text_collator import TextDataCollator
 from xturing.registry import BaseParent
 
@@ -11,3 +12,6 @@ class BasePreprocessor(BaseParent):
     InstructionDataCollator.config_name, InstructionDataCollator
 )
 BasePreprocessor.add_to_registry(TextDataCollator.config_name, TextDataCollator)
+BasePreprocessor.add_to_registry(
+    PreferenceDataCollator.config_name, PreferenceDataCollator
+)
diff --git a/src/xturing/preprocessors/preference_collator.py b/src/xturing/preprocessors/preference_collator.py
@@ -0,0 +1,117 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn.functional as F
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+
+from xturing.datasets.preference_dataset import PreferenceDatasetMeta
+
+
+class PreferenceDataCollator:
+    """Collator for preference datasets used in DPO training.
+
+    For each sample, this collator tokenizes two sequences:
+    - ``prompt + chosen`` (the preferred completion)
+    - ``prompt + rejected`` (the dispreferred completion)
+
+    The resulting batch contains ``chosen_input_ids``, ``chosen_attention_mask``,
+    ``chosen_labels``, and the corresponding ``rejected_*`` tensors. Labels are
+    masked so that the loss is only computed over the response tokens (not the
+    prompt).
+    """
+
+    config_name = "preference_dataset"
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        max_length: Optional[int] = None,
+        meta: PreferenceDatasetMeta = PreferenceDatasetMeta(),
+    ):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.meta = meta
+
+    def _tokenize_pair(self, prompt: str, response: str):
+        """Tokenize a prompt-response pair and return input_ids with a label
+        mask that marks only the response tokens as trainable."""
+        prompt_tokens = self.tokenizer(prompt)
+        response_tokens = self.tokenizer(response)
+
+        input_ids = prompt_tokens["input_ids"] + response_tokens["input_ids"]
+        # Labels: -100 for prompt tokens (ignored by loss), actual ids for response
+        label_mask = [False] * len(prompt_tokens["input_ids"]) + [True] * len(
+            response_tokens["input_ids"]
+        )
+
+        # Truncate to max_length - 1 to leave room for eos token
+        input_ids = input_ids[: self.max_length - 1]
+        input_ids.append(self.tokenizer.eos_token_id)
+        attention_mask = [1] * len(input_ids)
+
+        label_mask = label_mask[: self.max_length - 1]
+        label_mask.append(True)
+
+        return {
+            "input_ids": torch.tensor(input_ids).long(),
+            "attention_mask": torch.tensor(attention_mask).long(),
+            "label_mask": label_mask,
+        }
+
+    def _pad_and_stack(self, samples: List[Dict]):
+        """Pad a list of tokenized samples and stack into batch tensors."""
+        padded = self.tokenizer.pad(
+            [
+                {"input_ids": s["input_ids"], "attention_mask": s["attention_mask"]}
+                for s in samples
+            ],
+            padding=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+
+        dim = padded["input_ids"].shape[-1]
+        label_masks = torch.stack(
+            [
+                F.pad(
+                    torch.tensor(s["label_mask"]),
+                    (0, dim - len(s["label_mask"])),
+                    value=False,
+                )
+                for s in samples
+            ]
+        )
+
+        # Build labels: copy input_ids shifted by 1, masked with -100 for prompt tokens
+        labels = padded["input_ids"].clone()
+        labels[~label_masks] = -100
+
+        return {
+            "input_ids": padded["input_ids"],
+            "attention_mask": padded["attention_mask"],
+            "labels": labels,
+        }
+
+    def __call__(self, batches: List[Dict]):
+        chosen_samples = []
+        rejected_samples = []
+
+        for sample in batches:
+            chosen_samples.append(
+                self._tokenize_pair(sample["prompt"], sample["chosen"])
+            )
+            rejected_samples.append(
+                self._tokenize_pair(sample["prompt"], sample["rejected"])
+            )
+
+        chosen_batch = self._pad_and_stack(chosen_samples)
+        rejected_batch = self._pad_and_stack(rejected_samples)
+
+        return {
+            "chosen_input_ids": chosen_batch["input_ids"],
+            "chosen_attention_mask": chosen_batch["attention_mask"],
+            "chosen_labels": chosen_batch["labels"],
+            "rejected_input_ids": rejected_batch["input_ids"],
+            "rejected_attention_mask": rejected_batch["attention_mask"],
+            "rejected_labels": rejected_batch["labels"],
+        }
diff --git a/src/xturing/trainers/__init__.py b/src/xturing/trainers/__init__.py
@@ -1,2 +1,3 @@
 from xturing.trainers.base import BaseTrainer
+from xturing.trainers.dpo_trainer import DPOTrainer
 from xturing.trainers.lightning_trainer import LightningTrainer
diff --git a/src/xturing/trainers/base.py b/src/xturing/trainers/base.py
@@ -1,4 +1,5 @@
 from xturing.registry import BaseParent
+from xturing.trainers.dpo_trainer import DPOTrainer
 from xturing.trainers.lightning_trainer import LightningTrainer
 
 
@@ -7,3 +8,4 @@ class BaseTrainer(BaseParent):
 
 
 BaseTrainer.add_to_registry(LightningTrainer.config_name, LightningTrainer)
+BaseTrainer.add_to_registry(DPOTrainer.config_name, DPOTrainer)
diff --git a/src/xturing/trainers/dpo_trainer.py b/src/xturing/trainers/dpo_trainer.py
diff --git a/tests/xturing/trainers/__init__.py b/tests/xturing/trainers/__init__.py
diff --git a/tests/xturing/trainers/test_dpo_trainer.py b/tests/xturing/trainers/test_dpo_trainer.py

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`from xturing.trainers.base import BaseTrainer`
	`2`	`+from xturing.trainers.dpo_trainer import DPOTrainer`
`2`	`3`	`from xturing.trainers.lightning_trainer import LightningTrainer`