# Copyright (c) ModelScope Contributors. All rights reserved. import torch import torch.nn as nn import torch.nn.functional as F import warnings from accelerate.utils import gather_object from contextlib import contextmanager, nullcontext from peft import PeftModel from transformers import PreTrainedModel from transformers.utils.versions import require_version from trl import DPOTrainer as HFDPOTrainer from trl.trainer.dpo_config import DPOConfig from typing import Any, Dict, List, Literal, Optional, Tuple, Union from swift.trainers import DataLoaderMixin, SwiftMixin from swift.utils import get_logger, to_device from .rlhf_mixin import RLHFTrainerMixin try: from trl.trainer.utils import RunningMoments except ImportError: # trl >= 0.29 from trl.experimental.bco.bco_trainer import RunningMoments _ALPHA_DIVERGENCE_COEF_KEY = 'alpha_divergence_coef' _ALPHA_DIVERGENCE_COEF_DEFAULT = 1.0 del HFDPOTrainer.__init__ logger = get_logger() def _get_exp_cap(value, decimal=4): vdtype_max = torch.zeros([1]).to(value.dtype) + torch.finfo(value.dtype).max vdtype_log_max = torch.log(vdtype_max).to(value.device) return torch.floor(vdtype_log_max * 10**decimal) / 10**decimal if decimal > 0 else vdtype_log_max def _cap_exp(value, cap=-1): cap = _get_exp_cap(value) if cap < 0 else cap return torch.exp(torch.clamp(value, max=cap)) def new_gather_function(tensor): tensor_list = gather_object([tensor]) tensor_list = [t[None] if t.ndim == 0 else t for t in tensor_list] return torch.concat(to_device(tensor_list, tensor.device), dim=0) class DPOTrainer(RLHFTrainerMixin, SwiftMixin, DataLoaderMixin, HFDPOTrainer): def __init__(self, model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, *_args, **kwargs): args = kwargs['args'] self.label_smoothing = args.label_smoothing if 'loss_weights' in DPOConfig.__dict__: # trl >= 0.20 self.loss_type = args.loss_type if isinstance(args.loss_type, list) else [args.loss_type] self.loss_weights = args.loss_weights else: self.loss_type = args.loss_type loss_types = self.loss_type if isinstance(self.loss_type, list) else [self.loss_type] for loss_type in loss_types: if (loss_type in ['hinge', 'ipo', 'bco_pair', 'sppo_hard', 'nca_pair', 'apo_zero', 'apo_down'] and args.label_smoothing > 0): warnings.warn( f'You are using the {loss_type} loss type that does not support label smoothing. The ' '`label_smoothing` parameter will be ignored. ' 'Set `label_smoothing` to `0.0` to remove this warning.', UserWarning, ) if loss_type == 'kto_pair': raise ValueError('Support for kto_pair has been removed in DPOTrainer. Please use KTOTrainer.') self.precompute_ref_log_probs = args.precompute_ref_log_probs self.f_divergence_type = getattr(args, 'f_divergence_type', 'reverse_kl') self.f_alpha_divergence_coef = getattr(args, 'f_alpha_divergence_coef', 0.5) self.f_divergence_params = {_ALPHA_DIVERGENCE_COEF_KEY: self.f_alpha_divergence_coef} self.ref_adapter_name = getattr(args, 'ref_adapter_name', None) self.model_adapter_name = None self.reference_free = getattr(args, 'reference_free', None) or False self.use_weighting = False super().__init__(model, ref_model, *_args, **kwargs) if 'bco_pair' in loss_types: self.running = RunningMoments(self.accelerator) if self.args.ld_alpha is not None: require_version('trl>=0.18', '`ld_alpha` requires that "trl>=0.18".') if self.template.packing: self.accelerator.gather_for_metrics = new_gather_function def concatenated_forward( self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]], is_ref_model: bool = False ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: batch = batch.copy() use_logits_to_keep = self.get_use_logits_to_keep(self.template.sequence_parallel_size == 1) if use_logits_to_keep: self.prepare_logits_to_keep(batch) if self.aux_loss_enabled: batch['output_router_logits'] = True labels = batch.pop('labels', None) if self.is_encoder_decoder: batch['labels'] = labels text_position_ids = batch.pop('text_position_ids', None) if text_position_ids is None: text_position_ids = batch.get('position_ids') outputs = model(**batch, use_cache=False) all_logits = outputs.logits if all_logits.shape[1] != labels.shape[1]: # for llava, the model returns logits for the entire sequence, including the image tokens # (placed before the text tokens) all_logits = all_logits[:, -labels.shape[1]:] if not self.is_encoder_decoder and self.template.sequence_parallel_size == 1: # Shift so that tokens < n predict n labels = torch.roll(labels, shifts=-1, dims=1) per_token_logps, mean_all_logits, loss_mask = self.get_per_token_logps( all_logits, labels, label_pad_token_id=self.label_pad_token_id) origin_per_token_logps = per_token_logps loss_types = self.loss_type if isinstance(self.loss_type, list) else [self.loss_type] if 'ipo' in loss_types: size_completion = loss_mask.sum(dim=-1) per_token_logps = per_token_logps / size_completion output = {} if self.template.padding_free: cu_seqlens = self.get_cu_seqlens(text_position_ids, batch.get('logits_to_keep')) num_examples = (cu_seqlens.shape[0] - 1) // 2 all_logps = per_token_logps.new_zeros((num_examples * 2, )) completion_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]) chosen_lengths = completion_lengths[:num_examples] rejected_lengths = completion_lengths[num_examples:] public_lengths = torch.min(chosen_lengths, rejected_lengths) # l_p in the paper for i in range(cu_seqlens.shape[0] - 1): start, end = cu_seqlens[i], cu_seqlens[i + 1] length = end - start public_length = public_lengths[i % num_examples] if self.args.ld_alpha is not None and not is_ref_model and length > public_length: front_logps = per_token_logps[:, start:start + public_length].sum() rear_logps = per_token_logps[:, start + public_length:end].sum() all_logps[i] = front_logps + self.args.ld_alpha * rear_logps else: all_logps[i] = per_token_logps[:, start:end].sum() num_tokens = cu_seqlens[num_examples] if not is_ref_model: output['nll_loss'] = -origin_per_token_logps[:, :num_tokens][loss_mask[:, :num_tokens]].mean() output['chosen_logps'] = all_logps[:num_examples] output['rejected_logps'] = all_logps[num_examples:] output['mean_chosen_logits'] = mean_all_logits[:, :num_tokens][loss_mask[:, :num_tokens]].mean() output['mean_rejected_logits'] = mean_all_logits[:, num_tokens:][loss_mask[:, num_tokens:]].mean() else: num_examples = labels.shape[0] // 2 if not is_ref_model: output['nll_loss'] = -origin_per_token_logps[:num_examples][loss_mask[:num_examples]].mean() if self.args.ld_alpha is not None and not is_ref_model: completion_lengths = loss_mask.sum(dim=1) chosen_lengths = completion_lengths[:num_examples] rejected_lengths = completion_lengths[num_examples:] public_lengths = torch.min(chosen_lengths, rejected_lengths) # l_p in the paper public_lengths = torch.cat([public_lengths, public_lengths], dim=0) # Use loss_mask to compute position within completion # cumsum gives position within completion (1-indexed), subtract 1 to get 0-indexed completion_position_ids = (loss_mask.cumsum(dim=1) - 1) * loss_mask ld_mask = completion_position_ids < public_lengths.unsqueeze(1) # front_mask: positions within public_lengths (shared prefix) # rear_mask: positions beyond public_lengths (length-dependent suffix) front_mask = (ld_mask & loss_mask).float() rear_mask = (~ld_mask & loss_mask).float() front_logps = (per_token_logps * front_mask).sum(dim=1) rear_logps = (per_token_logps * rear_mask).sum(dim=1) all_logps = front_logps + self.args.ld_alpha * rear_logps else: all_logps = per_token_logps.sum(-1) output['chosen_logps'] = all_logps[:num_examples] output['rejected_logps'] = all_logps[num_examples:] output['mean_chosen_logits'] = mean_all_logits[:num_examples][loss_mask[:num_examples]].mean() output['mean_rejected_logits'] = mean_all_logits[num_examples:][loss_mask[num_examples:]].mean() if self.aux_loss_enabled: output['aux_loss'] = outputs.aux_loss return output def compute_ref_log_probs(self, batch): compute_ref_context_manager = ( torch.autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext()) with torch.no_grad(), compute_ref_context_manager: if self.ref_model is None: with self.null_ref_context(): ref_model_output = self.concatenated_forward(self.model, batch, is_ref_model=True) else: ref_model_output = self.concatenated_forward(self.ref_model, batch, is_ref_model=True) return ref_model_output['chosen_logps'], ref_model_output['rejected_logps'] def dpo_loss( self, chosen_logps: torch.FloatTensor, rejected_logps: torch.FloatTensor, ref_chosen_logps: torch.FloatTensor, ref_rejected_logps: torch.FloatTensor, loss_type: str = 'sigmoid', model_output: Optional[Dict[str, torch.FloatTensor]] = None, ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: device = self.accelerator.device chosen_logratios = chosen_logps.to(device) - (not self.reference_free) * ref_chosen_logps.to(device) rejected_logratios = rejected_logps.to(device) - (not self.reference_free) * ref_rejected_logps.to(device) if self.f_divergence_type == 'alpha_divergence': alpha_coef = _ALPHA_DIVERGENCE_COEF_DEFAULT if self.f_divergence_params and _ALPHA_DIVERGENCE_COEF_KEY in self.f_divergence_params: alpha_coef = float(self.f_divergence_params[_ALPHA_DIVERGENCE_COEF_KEY]) logits = (_cap_exp(rejected_logratios * -alpha_coef) - _cap_exp(chosen_logratios * -alpha_coef)) / alpha_coef else: logratios = chosen_logps - rejected_logps if self.reference_free: ref_logratios = torch.tensor([0], dtype=logratios.dtype, device=logratios.device) else: ref_logratios = ref_chosen_logps - ref_rejected_logps logratios = logratios.to(device) ref_logratios = ref_logratios.to(device) logits = logratios - ref_logratios if self.f_divergence_type == 'js_divergence': logits -= F.softplus(chosen_logratios) - F.softplus(rejected_logratios) if loss_type == 'sigmoid': losses = (-F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) - F.logsigmoid(-self.beta * logits) * self.label_smoothing) elif loss_type == 'robust': losses = (-F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + F.logsigmoid(-self.beta * logits) * self.label_smoothing) / (1 - 2 * self.label_smoothing) elif loss_type == 'exo_pair': import math if self.label_smoothing == 0: self.label_smoothing = 1e-3 losses = (self.beta * logits).sigmoid() * (F.logsigmoid( self.beta * logits) - math.log(1 - self.label_smoothing)) + (-self.beta * logits).sigmoid() * ( F.logsigmoid(-self.beta * logits) - math.log(self.label_smoothing)) elif loss_type == 'hinge': losses = torch.relu(1 - self.beta * logits) elif loss_type == 'ipo': losses = (logits - 1 / (2 * self.beta))**2 elif loss_type == 'bco_pair': chosen_logratios = chosen_logps - ref_chosen_logps rejected_logratios = rejected_logps - ref_rejected_logps chosen_rewards = self.beta * chosen_logratios rejected_rewards = self.beta * rejected_logratios rewards = torch.cat((chosen_rewards, rejected_rewards), 0).mean().detach() self.running.update(rewards) delta = self.running.mean losses = -F.logsigmoid( (self.beta * chosen_logratios) - delta) - F.logsigmoid(-(self.beta * rejected_logratios - delta)) elif loss_type == 'sppo_hard': a = chosen_logps - ref_chosen_logps b = rejected_logps - ref_rejected_logps losses = (a - 0.5 / self.beta)**2 + (b + 0.5 / self.beta)**2 elif loss_type == 'nca_pair': chosen_rewards = (chosen_logps - ref_chosen_logps) * self.beta rejected_rewards = (rejected_logps - ref_rejected_logps) * self.beta losses = (-F.logsigmoid(chosen_rewards) - 0.5 * F.logsigmoid(-chosen_rewards) - 0.5 * F.logsigmoid(-rejected_rewards)) elif loss_type == 'aot_unpaired': chosen_logratios = chosen_logps - ref_chosen_logps rejected_logratios = rejected_logps - ref_rejected_logps chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0) rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0) delta = chosen_logratios_sorted - rejected_logratios_sorted losses = (-F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) - F.logsigmoid(-self.beta * delta) * self.label_smoothing) elif loss_type == 'aot': logratios = chosen_logps - rejected_logps ref_logratios = ref_chosen_logps - ref_rejected_logps logratios_sorted, _ = torch.sort(logratios, dim=0) ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0) delta = logratios_sorted - ref_logratios_sorted losses = (-F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) - F.logsigmoid(-self.beta * delta) * self.label_smoothing) elif loss_type == 'apo_zero': losses_chosen = 1 - F.sigmoid(self.beta * chosen_logratios) losses_rejected = F.sigmoid(self.beta * rejected_logratios) losses = losses_chosen + losses_rejected elif loss_type == 'apo_down': losses_chosen = F.sigmoid(self.beta * chosen_logratios) losses_rejected = 1 - F.sigmoid(self.beta * (chosen_logratios - rejected_logratios)) losses = losses_chosen + losses_rejected elif loss_type == 'discopop': logratios = chosen_logps - rejected_logps ref_logratios = ref_chosen_logps - ref_rejected_logps logits = logratios - ref_logratios logits = logits * self.beta log_ratio_modulation = torch.sigmoid(logits / self.args.discopop_tau) logistic_component = -F.logsigmoid(logits) exp_component = torch.exp(-logits) losses = logistic_component * (1 - log_ratio_modulation) + exp_component * log_ratio_modulation elif loss_type == 'sft': sft_loss = model_output['nll_loss'] batch_size = chosen_logps.shape[0] losses = sft_loss.expand(batch_size) chosen_rewards = torch.zeros_like(chosen_logps) rejected_rewards = torch.zeros_like(rejected_logps) else: raise ValueError( f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', " "'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_unpaired', 'discopop', 'apo_zero', " "'apo_down', 'sft']") chosen_rewards = self.beta * (chosen_logps.to(device) - ref_chosen_logps.to(device)).detach() rejected_rewards = self.beta * (rejected_logps.to(device) - ref_rejected_logps.to(device)).detach() return losses, chosen_rewards, rejected_rewards def get_batch_loss_metrics( self, model: Union[PreTrainedModel, nn.Module], batch: Dict[str, Union[List, torch.LongTensor]], train_eval: Literal['train', 'eval'] = 'train', ) -> Tuple[torch.Tensor, Dict[str, float]]: metrics = {} model_output = self.concatenated_forward(model, batch) if 'ref_chosen_logps' in batch and 'ref_rejected_logps' in batch: ref_chosen_logps = batch['ref_chosen_logps'] ref_rejected_logps = batch['ref_rejected_logps'] else: ref_chosen_logps, ref_rejected_logps = self.compute_ref_log_probs(batch) losses = 0 chosen_rewards = 0 rejected_rewards = 0 loss_types = self.loss_type if isinstance(self.loss_type, list) else [self.loss_type] loss_weights = self.loss_weights if hasattr(self, 'loss_weights') and self.loss_weights else None for idx, loss_type in enumerate(loss_types): _losses, _chosen_rewards, _rejected_rewards = self.dpo_loss( model_output['chosen_logps'], model_output['rejected_logps'], ref_chosen_logps, ref_rejected_logps, loss_type, model_output, ) weight = loss_weights[idx] if loss_weights else 1.0 losses = losses + _losses * weight chosen_rewards = chosen_rewards + _chosen_rewards * weight rejected_rewards = rejected_rewards + _rejected_rewards * weight reward_accuracies = (chosen_rewards > rejected_rewards).float() if self.args.rpo_alpha is not None: losses = losses + self.args.rpo_alpha * model_output['nll_loss'] if self.use_weighting: losses = losses * model_output['policy_weights'] if self.aux_loss_enabled: losses = losses + self.aux_loss_coef * model_output['aux_loss'] prefix = 'eval_' if train_eval == 'eval' else '' metrics[f'{prefix}rewards/chosen'] = self.accelerator.gather_for_metrics(chosen_rewards).mean().item() metrics[f'{prefix}rewards/rejected'] = self.accelerator.gather_for_metrics(rejected_rewards).mean().item() metrics[f'{prefix}rewards/accuracies'] = self.accelerator.gather_for_metrics(reward_accuracies).mean().item() metrics[f'{prefix}rewards/margins'] = ( self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards).mean().item()) metrics[f'{prefix}logps/chosen'] = ( self.accelerator.gather_for_metrics(model_output['chosen_logps']).detach().mean().item()) metrics[f'{prefix}logps/rejected'] = ( self.accelerator.gather_for_metrics(model_output['rejected_logps']).detach().mean().item()) metrics[f'{prefix}logits/chosen'] = ( self.accelerator.gather_for_metrics(model_output['mean_chosen_logits']).detach().mean().item()) metrics[f'{prefix}logits/rejected'] = ( self.accelerator.gather_for_metrics(model_output['mean_rejected_logits']).detach().mean().item()) if self.args.rpo_alpha is not None or 'sft' in loss_types: metrics[f'{prefix}nll_loss'] = ( self.accelerator.gather_for_metrics(model_output['nll_loss']).detach().mean().item()) if self.aux_loss_enabled: metrics[f'{prefix}aux_loss'] = ( self.accelerator.gather_for_metrics(model_output['aux_loss']).detach().mean().item()) return losses.mean(), metrics def store_metrics(self, metrics, train_eval='train'): for key, value in metrics.items(): self._stored_metrics[train_eval][key].append(value) def log(self, logs, start_time=None): from transformers import Trainer train_eval = 'train' if 'loss' in logs else 'eval' for key, metrics in self._stored_metrics[train_eval].items(): logs[key] = torch.tensor(metrics).mean().item() del self._stored_metrics[train_eval] mode = 'train' if self.model.training else 'eval' custom_metrics = self.custom_metrics[mode] prefix = 'eval_' if mode == 'eval' else '' logs.update(self.compute_custom_metrics(custom_metrics, prefix)) return Trainer.log(self, logs, start_time) def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): compute_loss_context_manager = ( torch.autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext()) with compute_loss_context_manager: loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval='train') loss = loss.to(self.args.device) self.store_metrics(metrics, train_eval='train') if num_items_in_batch is not None and self.model_accepts_loss_kwargs: loss = loss / self.args.gradient_accumulation_steps if return_outputs: return loss, metrics return loss def training_step(self, model, inputs, *args, **kwargs): with self.template.forward_context(self.model, inputs): return super().training_step(model, inputs, *args, **kwargs) def prediction_step(self, model, inputs, prediction_loss_only=False, *args, **kwargs): with self.template.forward_context(self.model, inputs): inputs = self._prepare_inputs(inputs) with torch.no_grad(): loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval='eval') self.store_metrics(metrics, train_eval='eval') if prediction_loss_only: return loss.detach(), None, None logits_dict = { 'eval_logits/chosen': metrics['eval_logits/chosen'], 'eval_logits/rejected': metrics['eval_logits/rejected'], } logits = torch.tensor(list(logits_dict.values()), device=self.accelerator.device) labels = torch.zeros(logits.shape[0], device=self.accelerator.device) return (loss.detach(), logits, labels)