[v1][feature] add dpo trainer (#10544)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

examples/v1/train_lora/train_lora_dpo.yaml

@@ -0,0 +1,36 @@
+model: Qwen/Qwen3-4B
+model_class: llm
+
+template: qwen3_nothink
+
+# PEFT Configuration
+peft_config:
+  name: lora
+  r: 16
+  lora_alpha: 32
+  lora_dropout: 0.05
+  target_modules: all
+
+# Kernel Config
+kernel_config:
+  name: auto
+  include_kernels: auto
+
+# FSDP Config
+dist_config:
+  name: fsdp2
+  dcp_path: null
+
+### data
+train_dataset: data/v1_dpo_demo.yaml
+
+### training
+output_dir: ./outputs/test_lora
+micro_batch_size: 1
+cutoff_len: 2048
+learning_rate: 1.0e-5
+max_steps: 10
+
+### sample
+sample_backend: hf
+max_new_tokens: 128

src/llamafactory/v1/config/training_args.py

@@ -14,6 +14,7 @@
 
 import os
 from dataclasses import dataclass, field
+from typing import Literal
 from uuid import uuid4
 
 from .arg_utils import BatchingStrategy, PluginConfig, get_plugin_config
@@ -115,6 +116,30 @@ class TrainingArguments:
         default=1,
         metadata={"help": "Log metrics every N optimizer steps."},
     )
+    pref_loss: Literal["sigmoid", "orpo", "simpo"] = field(
+        default="sigmoid",
+        metadata={"help": "The type of DPO loss to use."},
+    )
+    pref_beta: float = field(
+        default=0.1,
+        metadata={"help": "The beta parameter in the preference loss."},
+    )
+    pref_ftx: float = field(
+        default=0.0,
+        metadata={"help": "The supervised fine-tuning loss coefficient in DPO training."},
+    )
+    simpo_gamma: float = field(
+        default=0.5,
+        metadata={"help": "The target reward margin term in SimPO loss."},
+    )
+    dpo_label_smoothing: float = field(
+        default=0.0,
+        metadata={"help": "The robust DPO label smoothing parameter in cDPO that should be between 0 and 0.5."},
+    )
+    ld_alpha: float | None = field(
+        default=None,
+        metadata={"help": "Alpha parameter from LD-DPO, controls weighting of verbose token log-probabilities."},
+    )
 
     def __post_init__(self) -> None:
         self.dist_config = get_plugin_config(self.dist_config)

src/llamafactory/v1/core/base_trainer.py

@@ -27,6 +27,7 @@ Train Phase:
 
 """
 
+import os
 from abc import abstractmethod
 
 import torch
@@ -318,6 +319,10 @@ class BaseTrainer:
                         "grad_norm": grad_norm,
                         "learning_rate": current_lr,
                     }
+                    # Merge per-step trainer metrics (e.g. DPO rewards/logps/logits)
+                    step_metrics = getattr(self, "_step_metrics", None)
+                    if step_metrics:
+                        logs.update(step_metrics)
                     self.callback_handler.on_log(self.args, self.state, logs)
 
                 if self.args.save_steps and self.global_step % self.args.save_steps == 0:

src/llamafactory/v1/launcher.py

@@ -146,7 +146,9 @@ def launch():
 
             run_sft()
         elif command == "dpo":
-            raise NotImplementedError("DPO trainer is not implemented yet.")
+            from llamafactory.v1.trainers.dpo_trainer import run_dpo
+
+            run_dpo()
         elif command == "rm":
             from llamafactory.v1.trainers.rm_trainer import run_rm
 
@@ -173,9 +175,9 @@ def main():
 
         run_sft()
     elif command == "dpo":
-        # from llamafactory.v1.trainers.dpo_trainer import run_dpo
-        # run_dpo()
-        raise NotImplementedError("DPO trainer is not implemented yet.")
+        from llamafactory.v1.trainers.dpo_trainer import run_dpo
+
+        run_dpo()
     elif command == "rm":
         from llamafactory.v1.trainers.rm_trainer import run_rm
 

src/llamafactory/v1/trainers/dpo_trainer.py

@@ -0,0 +1,450 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import copy
+import os
+
+import torch
+import torch.nn.functional as F
+
+from ..accelerator.interface import Dim, DistributedInterface
+from ..config import InputArgument, TrainingArguments, get_args
+from ..core.base_trainer import BaseTrainer
+from ..core.data_engine import DataEngine
+from ..core.model_engine import ModelEngine
+from ..utils import logging
+from ..utils.constants import IGNORE_INDEX
+from ..utils.types import BatchInput, HFModel, Tensor
+
+
+logger = logging.get_logger(__name__)
+
+
+def compute_sigmoid_dpo_loss(
+    policy_chosen_logps: Tensor,
+    policy_rejected_logps: Tensor,
+    ref_chosen_logps: Tensor,
+    ref_rejected_logps: Tensor,
+    beta: float = 0.1,
+    label_smoothing: float = 0.0,
+) -> Tensor:
+    r"""Standalone pure function for sigmoid DPO loss (Rafailov et al. 2023).
+
+    .. math::
+        \text{logits} = (\log\pi_\theta(y_c) - \log\pi_\text{ref}(y_c))
+                      - (\log\pi_\theta(y_r) - \log\pi_\text{ref}(y_r))
+        \mathcal{L} = -(1-\varepsilon)\log\sigma(\beta\cdot\text{logits})
+                      - \varepsilon\log\sigma(-\beta\cdot\text{logits})
+
+    Args:
+        policy_chosen_logps: Log-probabilities from the policy model for chosen responses.
+        policy_rejected_logps: Log-probabilities from the policy model for rejected responses.
+        ref_chosen_logps: Log-probabilities from the reference model for chosen responses.
+        ref_rejected_logps: Log-probabilities from the reference model for rejected responses.
+        beta: Temperature / scaling factor for the DPO loss.
+        label_smoothing: Label smoothing factor in [0, 1].
+
+    Returns:
+        Per-sample element-wise loss tensor.
+    """
+    chosen_logratios = policy_chosen_logps - ref_chosen_logps
+    rejected_logratios = policy_rejected_logps - ref_rejected_logps
+    logits = chosen_logratios - rejected_logratios
+    return (
+        -F.logsigmoid(beta * logits) * (1 - label_smoothing)
+        - F.logsigmoid(-beta * logits) * label_smoothing
+    )
+
+
+def _validate_dpo_dataset_format(train_dataset: DataEngine, dataset_path: str) -> None:
+    if train_dataset.streaming:
+        return
+
+    if len(train_dataset) == 0:
+        raise ValueError(f"DPO training dataset is empty: {dataset_path}")
+
+    sample = train_dataset[0]
+    if "chosen_messages" in sample and "rejected_messages" in sample:
+        return
+
+    dataset_name = sample.get("_dataset_name", "unknown")
+    sample_keys = sorted(sample.keys())
+    raise ValueError(
+        "DPO training requires pair-format samples containing chosen/rejected responses. "
+        f"First sample from dataset '{dataset_name}' has keys: {sample_keys}. "
+        "Please use pair data (e.g. a dataset with chosen_messages/rejected_messages)."
+    )
+
+
+class DPOTrainer(BaseTrainer):
+    def __init__(
+        self,
+        args: TrainingArguments,
+        model: HFModel,
+        renderer,
+        train_dataset,
+        callbacks=None,
+    ) -> None:
+        cp_size = args.dist_config.get("cp_size", 1) if args.dist_config is not None else 1
+        if cp_size > 1:
+            raise NotImplementedError("DPO trainer currently only supports cp_size == 1.")
+
+        self.pref_loss = args.pref_loss
+        self.pref_beta = args.pref_beta
+        self.pref_ftx = args.pref_ftx
+        self.simpo_gamma = args.simpo_gamma
+        self.ld_alpha = args.ld_alpha
+        self.dpo_label_smoothing = args.dpo_label_smoothing
+
+        # ref_model must be created AFTER super().__init__() because FSDP2 with
+        # init_on_meta materialises the model during _shard_model().  We defer
+        # creation to _init_ref_model() below.
+        self.ref_model = None
+
+        super().__init__(args, model, renderer, train_dataset, callbacks)
+
+        if self.pref_loss == "sigmoid":
+            self._init_ref_model()
+
+    def _shard_model(self) -> None:
+        if self.args.dist_config is None:
+            if DistributedInterface().get_world_size(Dim.DP) > 1:
+                from torch.nn.parallel import DistributedDataParallel as DDP
+
+                device_ids = None if self.device.type == "cpu" else [self.device.index]
+                self.model = DDP(self.model, device_ids=device_ids, find_unused_parameters=True)
+        else:
+            super()._shard_model()
+
+    @property
+    def _unwrapped_model(self):
+        model = self.model
+        if hasattr(model, "module"):
+            model = model.module
+        return model
+
+    # ------------------------------------------------------------------
+    # Reference model (frozen snapshot for sigmoid DPO)
+    # ------------------------------------------------------------------
+
+    @property
+    def _use_lora_ref(self) -> bool:
+        """Whether the policy model supports disable_adapter() for ref forward."""
+        unwrapped = self._unwrapped_model
+        return hasattr(unwrapped, "disable_adapter")
+
+    def _init_ref_model(self) -> None:
+        """Create a frozen copy of the initial model to serve as reference.
+
+        For LoRA / PEFT models the base weights are already frozen, so we
+        reuse the policy model with ``disable_adapter()`` instead of copying.
+        For full fine-tuning a deep copy is required because the policy model's
+        base weights change during training.
+
+        Must be called AFTER super().__init__() so that FSDP2 / DDP sharding
+        has materialised the model onto real devices.
+        """
+        if self._use_lora_ref:
+            self.ref_model = None
+            logger.info_rank0("LoRA detected — reference log-probs will reuse the base model via disable_adapter().")
+            return
+
+        unwrapped = self._unwrapped_model
+        self.ref_model = copy.deepcopy(unwrapped)
+        self.ref_model.eval()
+        for param in self.ref_model.parameters():
+            param.requires_grad_(False)
+        logger.info_rank0("Full fine-tuning — created independent reference model via deep copy.")
+
+    # ------------------------------------------------------------------
+    # Shared log-probability extraction from logits
+    # ------------------------------------------------------------------
+
+    def _extract_chosen_rejected_logps(
+        self,
+        logits: Tensor,
+        labels: Tensor,
+        token_type_ids: Tensor,
+        use_ld: bool = True,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        """Extract chosen / rejected log-probabilities (sum and average) from logits.
+
+        Args:
+            logits: (batch_size, seq_len, vocab_size)
+            labels: (batch_size, seq_len)
+            token_type_ids: (batch_size, seq_len) – 1=chosen, 2=rejected
+            use_ld: Whether to apply LD-DPO length-dependent weighting. Should be
+                ``False`` for the reference model to match the v0 behaviour where
+                ``ld_alpha`` is only applied to the policy log-probs.
+
+        Returns:
+            chosen_logps:   (batch_size,) sum of per-token log-probs for chosen
+            rejected_logps: (batch_size,) sum of per-token log-probs for rejected
+            chosen_logps_avg:   (batch_size,) length-normalised chosen log-probs
+            rejected_logps_avg: (batch_size,) length-normalised rejected log-probs
+        """
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        shift_token_type_ids = token_type_ids[..., 1:]
+
+        per_token_logps = -F.cross_entropy(
+            shift_logits.view(-1, shift_logits.size(-1)),
+            shift_labels.view(-1),
+            reduction="none",
+            ignore_index=IGNORE_INDEX,
+        ).view(shift_labels.size(0), shift_labels.size(1))
+
+        loss_mask = shift_labels != IGNORE_INDEX
+        chosen_mask = (shift_token_type_ids == 1) & loss_mask
+        rejected_mask = (shift_token_type_ids == 2) & loss_mask
+
+        chosen_valid_len = chosen_mask.sum(dim=-1)
+        rejected_valid_len = rejected_mask.sum(dim=-1)
+
+        ld_alpha = self.ld_alpha if use_ld else None
+        if ld_alpha is not None:
+            min_lengths = torch.min(chosen_valid_len, rejected_valid_len)
+            chosen_starts = torch.argmax(chosen_mask.int(), dim=1)
+            rejected_starts = torch.argmax(rejected_mask.int(), dim=1)
+
+            chosen_public_lengths = chosen_starts + min_lengths
+            rejected_public_lengths = rejected_starts + min_lengths
+
+            seq_len = shift_labels.size(1)
+            position_ids = torch.arange(seq_len, device=self.device).unsqueeze(0)
+
+            chosen_ld_mask = position_ids < chosen_public_lengths.unsqueeze(1)
+            rejected_ld_mask = position_ids < rejected_public_lengths.unsqueeze(1)
+
+            chosen_front_mask = (chosen_ld_mask * chosen_mask).float()
+            chosen_rear_mask = ((~chosen_ld_mask) * chosen_mask).float()
+            rejected_front_mask = (rejected_ld_mask * rejected_mask).float()
+            rejected_rear_mask = ((~rejected_ld_mask) * rejected_mask).float()
+
+            chosen_logps = (per_token_logps * chosen_front_mask).sum(dim=-1) + ld_alpha * (
+                per_token_logps * chosen_rear_mask
+            ).sum(dim=-1)
+            rejected_logps = (per_token_logps * rejected_front_mask).sum(dim=-1) + ld_alpha * (
+                per_token_logps * rejected_rear_mask
+            ).sum(dim=-1)
+        else:
+            chosen_logps = (per_token_logps * chosen_mask.float()).sum(dim=-1)
+            rejected_logps = (per_token_logps * rejected_mask.float()).sum(dim=-1)
+
+        chosen_logps_avg = chosen_logps / (chosen_valid_len + 1e-6)
+        rejected_logps_avg = rejected_logps / (rejected_valid_len + 1e-6)
+
+        return chosen_logps, rejected_logps, chosen_logps_avg, rejected_logps_avg
+
+    # ------------------------------------------------------------------
+    # Model inputs (block-diagonal attention + per-document position_ids)
+    # ------------------------------------------------------------------
+
+    def _prepare_model_inputs(self, input_ids: Tensor, token_type_ids: Tensor) -> dict[str, Tensor]:
+        """Build model inputs with block-diagonal attention and per-document position IDs.
+
+        In the v1 concatenated format each sample is::
+
+            [chosen prompt | chosen response | rejected prompt | rejected response]
+
+        with ``token_type_ids`` 1 / 2 marking the two documents.  A plain causal
+        mask would let the rejected half attend to the chosen half and produce
+        contiguous RoPE positions across the boundary, biasing the DPO objective.
+
+        We instead:
+
+        * pass ``token_type_ids`` as the attention mask so that Transformers v5
+          builds a **block-diagonal** causal mask (each document only attends to
+          itself — see :class:`RMTrainer` for the same pattern).
+        * compute ``position_ids`` that **reset at each document boundary** so
+          that every document gets its own RoPE positions starting from 0.
+        """
+        batch_size, seq_len = token_type_ids.shape
+        arange = torch.arange(seq_len, device=self.device).unsqueeze(0).expand(batch_size, -1)
+
+        chosen_mask = token_type_ids == 1
+        rejected_mask = token_type_ids == 2
+        chosen_lens = chosen_mask.sum(dim=1, keepdim=True)
+
+        position_ids = torch.zeros_like(token_type_ids)
+        position_ids[chosen_mask] = arange[chosen_mask]
+        position_ids[rejected_mask] = (arange - chosen_lens)[rejected_mask]
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": token_type_ids,  # block-diagonal doc mask (v5)
+            "position_ids": position_ids,
+        }
+
+    # ------------------------------------------------------------------
+    # Reference log-probabilities (frozen model, no grad)
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def _compute_ref_logps(self, batch: BatchInput) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        """Forward the frozen reference model and return chosen/rejected log-probs.
+
+        For LoRA models the base weights are frozen, so we reuse the policy
+        model with adapters disabled instead of maintaining a separate copy.
+        """
+        input_ids = batch["input_ids"].to(self.device, non_blocking=True)
+        labels = batch["labels"].to(self.device, non_blocking=True)
+        token_type_ids = batch["token_type_ids"].to(self.device, non_blocking=True)
+
+        model_inputs = self._prepare_model_inputs(input_ids, token_type_ids)
+
+        if self._use_lora_ref:
+            unwrapped = self._unwrapped_model
+            with unwrapped.disable_adapter():
+                ref_logits = unwrapped(**model_inputs, use_cache=False, return_dict=True).logits.float()
+        else:
+            ref_logits = self.ref_model(**model_inputs, use_cache=False, return_dict=True).logits.float()
+
+        return self._extract_chosen_rejected_logps(ref_logits, labels, token_type_ids, use_ld=False)
+
+    # ------------------------------------------------------------------
+    # Loss functions
+    # ------------------------------------------------------------------
+
+    def _sigmoid_dpo_loss(
+        self,
+        policy_chosen_logps: Tensor,
+        policy_rejected_logps: Tensor,
+        ref_chosen_logps: Tensor,
+        ref_rejected_logps: Tensor,
+    ) -> Tensor:
+        """Compute sigmoid DPO loss — delegates to :func:`compute_sigmoid_dpo_loss`."""
+        return compute_sigmoid_dpo_loss(
+            policy_chosen_logps,
+            policy_rejected_logps,
+            ref_chosen_logps,
+            ref_rejected_logps,
+            beta=self.pref_beta,
+            label_smoothing=self.dpo_label_smoothing,
+        )
+
+    def _odds_ratio_loss(self, chosen_logps_avg: Tensor, rejected_logps_avg: Tensor) -> Tensor:
+        log_odds = (chosen_logps_avg - rejected_logps_avg) - (
+            torch.log1p(-torch.exp(chosen_logps_avg)) - torch.log1p(-torch.exp(rejected_logps_avg))
+        )
+        sft_loss = -chosen_logps_avg
+        odds_ratio_loss = -F.logsigmoid(log_odds)
+        return sft_loss + self.pref_beta * odds_ratio_loss
+
+    def _simpo_loss(self, chosen_logps_avg: Tensor, rejected_logps_avg: Tensor) -> Tensor:
+        pi_logratios = chosen_logps_avg - rejected_logps_avg
+        gamma_logratios = self.simpo_gamma / self.pref_beta
+        logits = pi_logratios - gamma_logratios
+        simpo_loss = -F.logsigmoid(self.pref_beta * logits)
+        return simpo_loss
+
+    # ------------------------------------------------------------------
+    # Main compute_loss
+    # ------------------------------------------------------------------
+
+    def compute_loss(self, batch: BatchInput) -> Tensor:
+        input_ids = batch["input_ids"].to(self.device, non_blocking=True)
+        labels = batch["labels"].to(self.device, non_blocking=True)
+        token_type_ids = batch["token_type_ids"].to(self.device, non_blocking=True)
+
+        # Block-diagonal attention (token_type_ids as doc mask) + per-document position_ids
+        model_inputs = self._prepare_model_inputs(input_ids, token_type_ids)
+
+        # --- Policy forward ---
+        model_output = self.model(**model_inputs, use_cache=False, return_dict=True)
+        logits = model_output.logits.float()
+
+        # Split logits into chosen / rejected for metrics
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_token_type_ids = token_type_ids[..., 1:]
+        chosen_logit_mask = (shift_token_type_ids == 1).float()
+        rejected_logit_mask = (shift_token_type_ids == 2).float()
+
+        policy_chosen_logps, policy_rejected_logps, chosen_logps_avg, rejected_logps_avg = (
+            self._extract_chosen_rejected_logps(logits, labels, token_type_ids)
+        )
+
+        # Raw logits means (for logging)
+        chosen_logits_mean = (shift_logits.mean(dim=-1) * chosen_logit_mask).sum() / (chosen_logit_mask.sum() + 1e-6)
+        rejected_logits_mean = (shift_logits.mean(dim=-1) * rejected_logit_mask).sum() / (rejected_logit_mask.sum() + 1e-6)
+
+        if self.pref_loss == "sigmoid":
+            if not self._use_lora_ref and self.ref_model is None:
+                raise RuntimeError(
+                    "Reference model is required for sigmoid DPO loss but ref_model is None. "
+                    "This should not happen; the ref model is created at __init__ for sigmoid loss."
+                )
+
+            ref_chosen_logps, ref_rejected_logps, _, _ = self._compute_ref_logps(batch)
+            losses = self._sigmoid_dpo_loss(
+                policy_chosen_logps,
+                policy_rejected_logps,
+                ref_chosen_logps,
+                ref_rejected_logps,
+            )
+            # DPO rewards: beta * (policy_logps - ref_logps)
+            chosen_rewards = (self.pref_beta * (policy_chosen_logps - ref_chosen_logps)).detach()
+            rejected_rewards = (self.pref_beta * (policy_rejected_logps - ref_rejected_logps)).detach()
+        elif self.pref_loss == "orpo":
+            losses = self._odds_ratio_loss(chosen_logps_avg, rejected_logps_avg)
+            chosen_rewards = (self.pref_beta * chosen_logps_avg).detach()
+            rejected_rewards = (self.pref_beta * rejected_logps_avg).detach()
+        elif self.pref_loss == "simpo":
+            losses = self._simpo_loss(chosen_logps_avg, rejected_logps_avg)
+            chosen_rewards = (self.pref_beta * chosen_logps_avg).detach()
+            rejected_rewards = (self.pref_beta * rejected_logps_avg).detach()
+        else:
+            raise ValueError(f"Unknown pref_loss: {self.pref_loss}")
+
+        if self.pref_ftx > 1e-6:
+            sft_loss = -chosen_logps_avg
+            losses = losses + self.pref_ftx * sft_loss
+
+        # --- Per-step DPO metrics (matches v0 logging) ---
+        self._step_metrics = {
+            "rewards/chosen": chosen_rewards.mean().item(),
+            "rewards/rejected": rejected_rewards.mean().item(),
+            "rewards/accuracies": (chosen_rewards > rejected_rewards).float().mean().item(),
+            "rewards/margins": (chosen_rewards - rejected_rewards).mean().item(),
+            "logps/chosen": policy_chosen_logps.mean().item(),
+            "logps/rejected": policy_rejected_logps.mean().item(),
+            "logits/chosen": chosen_logits_mean.item(),
+            "logits/rejected": rejected_logits_mean.item(),
+        }
+
+        return losses.mean()
+
+
+def run_dpo(args: InputArgument = None):
+    model_args, data_args, training_args, _ = get_args(args)
+    if getattr(training_args, "use_cpu", False):
+        os.environ["FORCE_V1_CPU"] = "1"
+    DistributedInterface(training_args.dist_config)
+    train_dataset = DataEngine(data_args.train_dataset)
+    _validate_dpo_dataset_format(train_dataset, data_args.train_dataset)
+    model_engine = ModelEngine(model_args, is_train=True)
+    trainer = DPOTrainer(
+        args=training_args,
+        model=model_engine.model,
+        renderer=model_engine.renderer,
+        train_dataset=train_dataset,
+    )
+    trainer.fit()
+    trainer.save_model()
+    DistributedInterface().destroy()
+
+
+if __name__ == "__main__":
+    run_dpo()

tests_v1/trainers/test_dpo_loss_precision.py

@@ -0,0 +1,240 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Precision tests for v1 sigmoid-based DPO loss."""
+
+from types import SimpleNamespace
+
+import torch
+import torch.nn.functional as F
+
+from llamafactory.extras.constants import IGNORE_INDEX
+from llamafactory.train.dpo.trainer import CustomDPOTrainer
+from llamafactory.v1.trainers.dpo_trainer import DPOTrainer, compute_sigmoid_dpo_loss
+
+
+# ==============================================================================
+# Mock helpers
+# ==============================================================================
+
+def _make_mock_v1(
+    pref_beta: float = 0.1,
+    dpo_label_smoothing: float = 0.0,
+    ld_alpha: float | None = None,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        pref_beta=pref_beta,
+        dpo_label_smoothing=dpo_label_smoothing,
+        ld_alpha=ld_alpha,
+        device=torch.device("cpu"),
+    )
+
+
+def _make_mock_v0_dpo(beta: float = 0.1, label_smoothing: float = 0.0) -> SimpleNamespace:
+    mock = SimpleNamespace()
+    mock.beta = beta
+    mock.label_smoothing = label_smoothing
+    mock.reference_free = False
+    mock.f_divergence_type = "reverse_kl"
+    mock.f_divergence_params = None
+    mock.accelerator = SimpleNamespace()
+    mock.accelerator.device = torch.device("cpu")
+    return mock
+
+
+# ==============================================================================
+# Fixed test inputs
+# ==============================================================================
+
+P_CHOSEN = torch.tensor([-3.0, -2.5, -4.0, -1.5])
+P_REJECTED = torch.tensor([-5.0, -3.5, -6.0, -2.5])
+R_CHOSEN = torch.tensor([-2.8, -2.3, -3.8, -1.4])
+R_REJECTED = torch.tensor([-3.2, -2.7, -4.2, -1.8])
+
+
+# ==============================================================================
+# Test 1 — Core loss correctness (pure function ↔ v1 instance ↔ v0/TRL)
+# ==============================================================================
+
+def test_sigmoid_dpo_loss_correctness():
+    """Comprehensive correctness check for compute_sigmoid_dpo_loss and its wrapper."""
+    # ---- 1a: pure function matches instance method ----
+    v1 = _make_mock_v1(pref_beta=0.1)
+    actual = DPOTrainer._sigmoid_dpo_loss(v1, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED)
+    expected = compute_sigmoid_dpo_loss(P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, beta=0.1)
+    torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6)
+
+    # ---- 1b: v1 matches v0 (TRL) on fixed inputs ----
+    v0 = _make_mock_v0_dpo(beta=0.1)
+    v0_losses, _, _ = CustomDPOTrainer.dpo_loss(
+        v0, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid",
+    )
+    torch.testing.assert_close(actual, v0_losses, rtol=1e-6, atol=1e-6)
+
+    # ---- 1c: multiple beta values (v1 ↔ v0) ----
+    for beta in [0.01, 0.1, 0.5, 1.0]:
+        v0b = _make_mock_v0_dpo(beta=beta)
+        v1b = _make_mock_v1(pref_beta=beta)
+        vl, _, _ = CustomDPOTrainer.dpo_loss(
+            v0b, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid",
+        )
+        v1l = DPOTrainer._sigmoid_dpo_loss(v1b, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED)
+        torch.testing.assert_close(v1l, vl, rtol=1e-6, atol=1e-6)
+
+    # ---- 1d: label_smoothing sweep (v1 ↔ v0) ----
+    for ls in [0.0, 0.1, 0.2, 0.3]:
+        v0s = _make_mock_v0_dpo(beta=0.1, label_smoothing=ls)
+        v1s = _make_mock_v1(pref_beta=0.1, dpo_label_smoothing=ls)
+        vl, _, _ = CustomDPOTrainer.dpo_loss(
+            v0s, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid",
+        )
+        v1l = DPOTrainer._sigmoid_dpo_loss(v1s, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED)
+        torch.testing.assert_close(v1l, vl, rtol=1e-6, atol=1e-6)
+
+    # ---- 1e: label_smoothing=0.5 symmetry (swap chosen↔rejected same loss) ----
+    v1s = _make_mock_v1(pref_beta=0.1, dpo_label_smoothing=0.5)
+    fwd = DPOTrainer._sigmoid_dpo_loss(v1s, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED)
+    swp = DPOTrainer._sigmoid_dpo_loss(v1s, P_REJECTED, P_CHOSEN, R_REJECTED, R_CHOSEN)
+    torch.testing.assert_close(fwd, swp, rtol=1e-6, atol=1e-6)
+
+    # ---- 1f: chosen better → lower loss ----
+    v1c = _make_mock_v1(pref_beta=0.1)
+    loss_good = DPOTrainer._sigmoid_dpo_loss(
+        v1c,
+        torch.tensor([-1.0]), torch.tensor([-10.0]),
+        torch.tensor([-3.0]), torch.tensor([-3.0]),
+    )
+    loss_bad = DPOTrainer._sigmoid_dpo_loss(
+        v1c,
+        torch.tensor([-10.0]), torch.tensor([-1.0]),
+        torch.tensor([-3.0]), torch.tensor([-3.0]),
+    )
+    assert loss_good.item() < loss_bad.item()
+
+    # ---- 1g: policy == ref → loss = log(2) ≈ 0.693 ----
+    logps = torch.tensor([-3.0, -2.0, -4.0])
+    losses = DPOTrainer._sigmoid_dpo_loss(v1c, logps, logps, logps, logps)
+    expected_log2 = torch.full_like(logps, -F.logsigmoid(torch.tensor(0.0)).item())
+    torch.testing.assert_close(losses, expected_log2, rtol=1e-5, atol=1e-5)
+
+    # ---- 1h: non-negative ----
+    assert (actual >= 0).all()
+
+    # ---- 1i: extreme logps stay finite ----
+    v1x = _make_mock_v1(pref_beta=0.1)
+    x = DPOTrainer._sigmoid_dpo_loss(
+        v1x,
+        torch.tensor([-0.1, -50.0, -0.5, -100.0]),
+        torch.tensor([-0.2, -5.0, -30.0, -1.0]),
+        torch.tensor([-0.15, -3.0, -0.6, -2.0]),
+        torch.tensor([-0.25, -4.0, -5.0, -1.5]),
+    )
+    assert torch.isfinite(x).all()
+
+
+# ==============================================================================
+# Test 2 — Random cross-validation & reward equivalence
+# ==============================================================================
+
+def test_cross_validate_and_rewards():
+    """Randomised v0↔v1 cross-validation (50 seeds) + reward-margin check."""
+    torch.manual_seed(42)
+    for _ in range(50):
+        pc = -torch.rand(4) * 10 - 0.01
+        pr = -torch.rand(4) * 15 - 0.01
+        rc = -torch.rand(4) * 10 - 0.01
+        rr = -torch.rand(4) * 12 - 0.01
+        beta = 0.01 + torch.rand(1).item() * 0.5
+        ls = torch.rand(1).item() * 0.3
+
+        v0 = _make_mock_v0_dpo(beta=beta, label_smoothing=ls)
+        v1 = _make_mock_v1(pref_beta=beta, dpo_label_smoothing=ls)
+
+        v0_loss, _, _ = CustomDPOTrainer.dpo_loss(
+            v0, pc, pr, rc, rr, loss_type="sigmoid",
+        )
+        v1_loss = DPOTrainer._sigmoid_dpo_loss(v1, pc, pr, rc, rr)
+        torch.testing.assert_close(v1_loss, v0_loss, rtol=1e-5, atol=1e-5)
+
+        # Reward margin = beta * (chosen_logratio - rejected_logratio)
+        chosen_rewards = beta * (pc - rc)
+        rejected_rewards = beta * (pr - rr)
+        reward_margin = chosen_rewards - rejected_rewards
+        logits = (pc - rc) - (pr - rr)
+        torch.testing.assert_close(reward_margin, beta * logits, rtol=1e-6, atol=1e-6)
+
+    # Fixed-input reward ordering
+    cr = 0.1 * (P_CHOSEN - R_CHOSEN)
+    rr = 0.1 * (P_REJECTED - R_REJECTED)
+    assert (cr > rr).float().mean().item() == 1.0
+
+
+# ==============================================================================
+# Test 3 — End-to-end: log-prob extraction + synthetic batch + LD-DPO
+# ==============================================================================
+
+def _make_batch(num_pairs, seq_len, vocab_size, prompt_len=3, chosen_len=None, rejected_len=None):
+    if chosen_len is None or rejected_len is None:
+        rlen = (seq_len - prompt_len) // 2
+        chosen_len = rlen
+        rejected_len = rlen
+
+    actual = prompt_len + chosen_len + rejected_len
+
+    torch.manual_seed(42)
+    input_ids = torch.randint(0, vocab_size, (num_pairs, actual))
+    labels = input_ids.clone()
+    labels[:, :prompt_len] = IGNORE_INDEX
+
+    token_type_ids = torch.zeros(num_pairs, actual, dtype=torch.long)
+    token_type_ids[:, prompt_len:prompt_len + chosen_len] = 1
+    token_type_ids[:, prompt_len + chosen_len:] = 2
+
+    torch.manual_seed(99)
+    logits = torch.randn(num_pairs, actual, vocab_size)
+    return input_ids, labels, token_type_ids, logits
+
+
+def test_logp_extraction_and_e2e_loss():
+    """Log-prob extraction shapes + e2e sigmoid loss (equal & unequal lengths)."""
+    # --- equal-length batch ---
+    ids, labels, tt_ids, logits = _make_batch(2, 12, 64, prompt_len=2)
+    v1 = _make_mock_v1(pref_beta=0.1)
+
+    c_lp, r_lp, c_avg, r_avg = DPOTrainer._extract_chosen_rejected_logps(v1, logits, labels, tt_ids)
+    assert c_lp.shape == r_lp.shape == c_avg.shape == r_avg.shape == (2,)
+    assert (c_lp <= 1e-6).all() and (r_lp <= 1e-6).all()
+
+    # Create "ref" logits with small noise
+    torch.manual_seed(123)
+    ref_logits = logits + 0.1 * torch.randn_like(logits)
+    rc_lp, rr_lp, _, _ = DPOTrainer._extract_chosen_rejected_logps(v1, ref_logits, labels, tt_ids)
+
+    losses = DPOTrainer._sigmoid_dpo_loss(v1, c_lp, r_lp, rc_lp, rr_lp)
+    assert torch.isfinite(losses).all() and (losses >= 0).all()
+
+    # --- unequal-length (LD-DPO) batch ---
+    ids2, labels2, tt_ids2, logits2 = _make_batch(
+        1, 11, 64, prompt_len=2, chosen_len=6, rejected_len=3,
+    )
+    v1_ld = _make_mock_v1(pref_beta=0.1, ld_alpha=0.5)
+
+    c_lp2, r_lp2, _, _ = DPOTrainer._extract_chosen_rejected_logps(v1_ld, logits2, labels2, tt_ids2)
+
+    torch.manual_seed(123)
+    ref2 = logits2 + 0.1 * torch.randn_like(logits2)
+    rc2, rr2, _, _ = DPOTrainer._extract_chosen_rejected_logps(v1_ld, ref2, labels2, tt_ids2)
+
+    losses2 = DPOTrainer._sigmoid_dpo_loss(v1_ld, c_lp2, r_lp2, rc2, rr2)
+    assert torch.isfinite(losses2).all()

tests_v1/trainers/test_fsdp2_dpo_trainer.py

@@ -0,0 +1,97 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.xfail(reason="CI machines may OOM when heavily loaded.")
+@pytest.mark.runs_on(["cuda", "npu"])
+def test_fsdp2_dpo_trainer(tmp_path: Path):
+    """Test FSDP2 DPO trainer with sigmoid loss by simulating `llamafactory-cli dpo config.yaml`."""
+    config_yaml = """\
+model: Qwen/Qwen3-0.6B
+trust_remote_code: true
+model_class: llm
+
+template: qwen3_nothink
+
+kernel_config:
+    name: auto
+    include_kernels: auto
+
+quant_config: null
+
+dist_config:
+    name: fsdp2
+    dcp_path: null
+
+init_config:
+    name: init_on_meta
+
+# PEFT Configuration
+peft_config:
+    name: lora
+    r: 8
+    lora_alpha: 16
+    lora_dropout: 0.0
+    target_modules: all
+
+### data
+train_dataset: data/v1_dpo_demo.yaml
+
+### training
+output_dir: {output_dir}
+micro_batch_size: 1
+global_batch_size: 1
+cutoff_len: 2048
+learning_rate: 1.0e-4
+bf16: false
+max_steps: 1
+pref_loss: sigmoid
+pref_beta: 0.1
+dpo_label_smoothing: 0.0
+
+### sample
+sample_backend: hf
+max_new_tokens: 128
+"""
+    # Create output directory
+    output_dir = tmp_path / "outputs"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    config_file = tmp_path / "config.yaml"
+    config_file.write_text(config_yaml.format(output_dir=str(output_dir)))
+
+    # Set up environment variables
+    env = os.environ.copy()
+    env["USE_V1"] = "1"  # Use v1 launcher
+    env["FORCE_TORCHRUN"] = "1"  # Force distributed training via torchrun
+
+    # Run the CLI command via subprocess
+    result = subprocess.run(
+        [sys.executable, "-m", "llamafactory.cli", "dpo", str(config_file)],
+        env=env,
+        capture_output=True,
+        cwd=str(Path(__file__).parent.parent.parent),  # LLaMA-Factory root
+    )
+
+    # Decode output with error handling (progress bars may contain non-UTF-8 bytes)
+    stderr = result.stderr.decode("utf-8", errors="replace")
+
+    # Check the result
+    assert result.returncode == 0, f"DPO training failed with return code {result.returncode}\nSTDERR: {stderr}"