modelscope--ms-swift
wehub-resource-sync
95 行
4.2 KiB
Python
95 行
4.2 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import torch
|
|
from collections import namedtuple
|
|
from functools import partial
|
|
from megatron.core import mpu
|
|
from torch.distributed.nn import all_reduce
|
|
|
|
from swift.rlhf_trainers import DPOTrainer
|
|
from swift.utils import get_current_device, get_logger
|
|
from .rlhf_mixin import MegatronRLHFTrainer
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class DummyDPOTrainer(DPOTrainer):
|
|
# For reusing the dpo_loss function implemented in Swift's DPOTrainer.
|
|
def __init__(self, args):
|
|
self.accelerator = namedtuple('Accelerator', ['device'])(device=get_current_device())
|
|
self.f_alpha_divergence_coef = 1.
|
|
self.f_divergence_params = {'alpha_divergence_coef': self.f_alpha_divergence_coef}
|
|
self.reference_free = args.reference_free
|
|
self.label_smoothing = args.label_smoothing
|
|
self.f_divergence_type = args.f_divergence_type
|
|
self.loss_type = args.loss_type
|
|
self.beta = args.beta
|
|
|
|
|
|
class MegatronDPOTrainer(MegatronRLHFTrainer):
|
|
|
|
def __init__(self, args, template):
|
|
super().__init__(args, template)
|
|
self.dummy_dpo_trainer = DummyDPOTrainer(args)
|
|
|
|
def loss_func(self, output_tensor: torch.Tensor, *, labels: torch.Tensor, packed_seq_params):
|
|
ref_output_tensor = output_tensor[:output_tensor.shape[0] // 2].detach()
|
|
output_tensor = output_tensor[output_tensor.shape[0] // 2:]
|
|
args = self.args
|
|
num_samples = labels.shape[0] if packed_seq_params is None else packed_seq_params.seq_lens.shape[0]
|
|
|
|
logps = self.get_logps(output_tensor, labels, packed_seq_params)
|
|
ref_logps = self.get_logps(ref_output_tensor, labels, packed_seq_params)
|
|
loss, chosen_rewards, rejected_rewards = self.dummy_dpo_trainer.dpo_loss(
|
|
logps[:num_samples // 2],
|
|
logps[num_samples // 2:],
|
|
ref_logps[:num_samples // 2],
|
|
ref_logps[num_samples // 2:],
|
|
)
|
|
if args.rpo_alpha:
|
|
loss_mask = labels != -100
|
|
if args.padding_free:
|
|
num_tokens = packed_seq_params.cu_seqlens_q[num_samples // 2] // args.context_parallel_size
|
|
loss_mask[:, num_tokens:] = 0
|
|
else:
|
|
loss_mask[num_samples // 2:] = 0
|
|
nll_loss = torch.concat([torch.sum(output_tensor * loss_mask)[None], loss_mask.sum()[None]])
|
|
if args.context_parallel_size > 1:
|
|
nll_loss = all_reduce(nll_loss, group=mpu.get_context_parallel_group())
|
|
nll_loss = nll_loss[0] / nll_loss[1]
|
|
loss = loss + args.rpo_alpha * nll_loss
|
|
loss = loss.mean()
|
|
metric = {
|
|
'loss': loss.detach().clone(),
|
|
'logps/chosen': logps[:num_samples // 2].mean(),
|
|
'logps/rejected': logps[num_samples // 2:].mean(),
|
|
'rewards/chosen': chosen_rewards.mean(),
|
|
'rewards/rejected': rejected_rewards.mean(),
|
|
'rewards/accuracies': (chosen_rewards > rejected_rewards).float().mean(),
|
|
'rewards/margins': (chosen_rewards - rejected_rewards).mean(),
|
|
}
|
|
if args.rpo_alpha:
|
|
metric['nll_loss'] = nll_loss.detach()
|
|
metric = self._all_reduce_metric(metric)
|
|
# fix megatron-lm bug
|
|
loss = loss / mpu.get_context_parallel_world_size()
|
|
return loss, metric
|
|
|
|
def forward_step(self, data_iterator, model):
|
|
unwrapped_model = model.module.module
|
|
input_tensor = unwrapped_model.get_input_tensor()
|
|
vp_stage = unwrapped_model.vp_stage
|
|
data = self.get_batch(data_iterator, vp_stage)
|
|
data.pop('loss_scale', None)
|
|
# ref_model
|
|
with torch.no_grad(), self.null_ref_context() as ref_models:
|
|
ref_model = ref_models[vp_stage or 0]
|
|
if input_tensor is not None:
|
|
ref_model.set_input_tensor(input_tensor[:input_tensor.shape[0] // 2].detach())
|
|
ref_output_tensor = ref_model(**data)
|
|
|
|
if input_tensor is not None:
|
|
unwrapped_model.set_input_tensor(input_tensor[input_tensor.shape[0] // 2:])
|
|
output_tensor = model(**data)
|
|
return torch.concat([ref_output_tensor, output_tensor], dim=0), partial(
|
|
self.loss_func, labels=data.get('labels'), packed_seq_params=data.get('packed_seq_params'))
|