diff --git a/examples/v1/train_lora/train_lora_dpo.yaml b/examples/v1/train_lora/train_lora_dpo.yaml new file mode 100644 index 000000000..32beec105 --- /dev/null +++ b/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 diff --git a/src/llamafactory/v1/config/training_args.py b/src/llamafactory/v1/config/training_args.py index 200938b9c..a69060f69 100644 --- a/src/llamafactory/v1/config/training_args.py +++ b/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) diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index c2eb2bebd..d869877ea 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/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: diff --git a/src/llamafactory/v1/launcher.py b/src/llamafactory/v1/launcher.py index f0481d4e7..b20988b24 100644 --- a/src/llamafactory/v1/launcher.py +++ b/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 diff --git a/src/llamafactory/v1/trainers/dpo_trainer.py b/src/llamafactory/v1/trainers/dpo_trainer.py index e69de29bb..d8f289b34 100644 --- a/src/llamafactory/v1/trainers/dpo_trainer.py +++ b/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() diff --git a/tests_v1/trainers/test_dpo_loss_precision.py b/tests_v1/trainers/test_dpo_loss_precision.py new file mode 100644 index 000000000..3009406e8 --- /dev/null +++ b/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() diff --git a/tests_v1/trainers/test_fsdp2_dpo_trainer.py b/tests_v1/trainers/test_fsdp2_dpo_trainer.py new file mode 100644 index 000000000..62cefdb38 --- /dev/null +++ b/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}"