From 1a18e651f642f6f2c51a18f1db2090ebdbb172e4 Mon Sep 17 00:00:00 2001 From: QZL-max Date: Wed, 12 Aug 2026 22:53:12 +0800 Subject: [PATCH] Add CE regularization term to DPO loss --- trainer/train_dpo.py | 31 ++++++++++++++++++++++++++----- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 479eb85..730b019 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -50,6 +50,19 @@ def dpo_loss(ref_log_probs, policy_log_probs, mask, beta): return loss.mean() +def sft_ce_loss(logits, labels, mask): + """SFT正则项：对chosen回答部分计算标准语言建模交叉熵（按掩码位置取平均）， + 与SFT阶段的损失口径一致，用于防止偏好优化中遗忘语言能力。""" + vocab_size = logits.size(-1) + ce = F.cross_entropy( + logits.reshape(-1, vocab_size), + labels.reshape(-1), + reduction='none' + ) + mask_flat = mask.reshape(-1).float() + return (ce * mask_flat).sum() / mask_flat.sum().clamp(min=1.0) + + def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=None, beta=0.1): start_time = time.time() last_step = start_step @@ -81,7 +94,13 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= policy_log_probs = logits_to_log_probs(logits, y) dpo_loss_val = dpo_loss(ref_log_probs, policy_log_probs, mask, beta=beta) - loss = dpo_loss_val + outputs.aux_loss + # SFT/CE正则项：只对chosen回答部分计算交叉熵，权重由 --ce_alpha 控制 + chosen_num = x.size(0) // 2 + if args.ce_alpha > 0: + ce_loss_val = sft_ce_loss(logits[:chosen_num], y[:chosen_num], mask[:chosen_num]) + else: + ce_loss_val = logits.new_zeros(()) + loss = dpo_loss_val + args.ce_alpha * ce_loss_val + outputs.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -97,13 +116,14 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps current_dpo_loss = dpo_loss_val.item() + current_ce_loss = ce_loss_val.item() current_aux_loss = outputs.aux_loss.item() current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 - Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, dpo_loss: {current_dpo_loss:.4f}, aux_loss: {current_aux_loss:.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min') + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, dpo_loss: {current_dpo_loss:.4f}, ce_loss: {current_ce_loss:.4f}, aux_loss: {current_aux_loss:.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min') - if wandb: wandb.log({"loss": current_loss, "dpo_loss": current_dpo_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) + if wandb: wandb.log({"loss": current_loss, "dpo_loss": current_dpo_loss, "ce_loss": current_ce_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters) and is_main_process(): model.eval() @@ -118,7 +138,7 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= del state_dict del x_chosen, x_rejected, y_chosen, y_rejected, mask_chosen, mask_rejected, x, y, mask - del ref_outputs, ref_logits, ref_log_probs, outputs, logits, policy_log_probs, loss + del ref_outputs, ref_logits, ref_log_probs, outputs, logits, policy_log_probs, ce_loss_val, loss if last_step > start_step and last_step % args.accumulation_steps != 0: scaler.unscale_(optimizer) @@ -150,6 +170,7 @@ if __name__ == "__main__": parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练") parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）") parser.add_argument('--beta', default=0.15, type=float, help="DPO中的beta参数") + parser.add_argument('--ce_alpha', default=0.1, type=float, help="SFT/CE正则项权重（0表示关闭，建议0.05~0.5；防止DPO训练中遗忘语言能力）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-DPO", help="wandb项目名") parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）") @@ -225,4 +246,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group()