From 3c56aca1c1811001343c2b0662714a33c01a692d Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Thu, 1 Jan 2026 22:37:49 +0800 Subject: [PATCH] [update] aux loss --- model/model_minimind.py | 9 ++------- trainer/train_distill_reason.py | 10 ++++++---- trainer/train_distillation.py | 17 ++++++++++------- trainer/train_dpo.py | 9 ++++++--- trainer/train_full_sft.py | 10 ++++++---- trainer/train_grpo.py | 19 +++++++++++++------ trainer/train_lora.py | 10 ++++++---- trainer/train_ppo.py | 18 ++++++++++++------ trainer/train_pretrain.py | 10 ++++++---- trainer/train_spo.py | 19 +++++++++++++------ 10 files changed, 80 insertions(+), 51 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index fd40738..830c78d 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -284,7 +284,7 @@ class MoEGate(nn.Module): fi = ce * self.n_routed_experts aux_loss = (Pi * fi).sum() * self.alpha else: - aux_loss = 0 + aux_loss = scores.new_zeros(1).squeeze() return topk_idx, topk_weight, aux_loss @@ -423,12 +423,7 @@ class MiniMindModel(nn.Module): hidden_states = self.norm(hidden_states) - aux_loss = sum( - layer.mlp.aux_loss - for layer in self.layers - if isinstance(layer.mlp, MOEFeedForward) - ) - + aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze()) return hidden_states, presents, aux_loss diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index be69738..3a10847 100644 --- a/trainer/train_distill_reason.py +++ b/trainer/train_distill_reason.py @@ -53,8 +53,8 @@ def train_epoch(epoch, loader, iters, tokenizer, lm_config, start_step=0, wandb= loss_mask_sum = loss_mask.sum() loss_mask[sp_ids] = 10 # 对思考标签增加10倍权重 loss_mask = loss_mask.view(Y.size()) - loss = (loss * loss_mask).sum() / loss_mask_sum - loss += res.aux_loss + logits_loss = (loss * loss_mask).sum() / loss_mask_sum + loss = logits_loss + res.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -69,12 +69,14 @@ def train_epoch(epoch, loader, iters, tokenizer, lm_config, start_step=0, wandb= if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_logits_loss = logits_loss.item() + current_aux_loss = res.aux_loss.item() current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_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, "lr": current_lr, "epoch_Time": eta_min}) + if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 712bac6..24450a3 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -71,9 +71,9 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st ignore_index=0, reduction='none' ) - ce_loss = torch.sum(ce_loss * loss_mask_flat) / loss_mask_flat.sum() - if lm_config_student.use_moe: - ce_loss += res.aux_loss + ce_loss_raw = torch.sum(ce_loss * loss_mask_flat) / loss_mask_flat.sum() + if lm_config_student.use_moe: ce_loss = ce_loss_raw + res.aux_loss + else: ce_loss = ce_loss_raw # 2) Distillation Loss if teacher_model is not None: @@ -100,18 +100,21 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_ce_loss = ce_loss_raw.item() + current_aux_loss = res.aux_loss.item() if lm_config_student.use_moe else 0.0 current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} ce:{ce_loss.item():.4f} distill:{distill_loss.item():.4f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, ce: {current_ce_loss:.4f}, aux_loss: {current_aux_loss:.4f}, distill: {distill_loss.item():.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min') if wandb: wandb.log({ "loss": current_loss, - "ce_loss": ce_loss.item(), + "ce_loss": current_ce_loss, + "aux_loss": current_aux_loss, "distill_loss": distill_loss.item() if teacher_model is not None else 0.0, - "lr": current_lr, - "epoch_Time": eta_min + "learning_rate": current_lr, + "epoch_time": eta_min }) if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index b20b53f..eccfaff 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -79,7 +79,8 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= logits = outputs.logits policy_log_probs = logits_to_log_probs(logits, y) - loss = dpo_loss(ref_log_probs, policy_log_probs, mask, beta=beta) + dpo_loss_val = dpo_loss(ref_log_probs, policy_log_probs, mask, beta=beta) + loss = dpo_loss_val + outputs.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -94,12 +95,14 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_dpo_loss = dpo_loss_val.item() + current_aux_loss = outputs.aux_loss.item() current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + 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') - if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_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 (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index cea2ffb..faec3f5 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -38,8 +38,8 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): Y.view(-1) ).view(Y.size()) - loss = (loss * loss_mask).sum() / loss_mask.sum() - loss += res.aux_loss + logits_loss = (loss * loss_mask).sum() / loss_mask.sum() + loss = logits_loss + res.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -56,12 +56,14 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_logits_loss = logits_loss.item() + current_aux_loss = res.aux_loss.item() current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_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, "lr": current_lr, "epoch_Time": eta_min}) + if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 897d9a8..d5b30ed 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -119,7 +119,11 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token per_token_logps.append(torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1)) return torch.stack(per_token_logps) - per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1)) # [B*num_gen, R] + with autocast_ctx: + per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1)) # [B*num_gen, R] + res = model(outputs) if lm_config.use_moe else None + aux_loss = res.aux_loss if res is not None else torch.tensor(0.0, device=args.device) + with torch.no_grad(): ref_per_token_logps = get_per_token_logps(ref_model, outputs, completion_ids.size(1)) # [B*num_gen, R] @@ -140,7 +144,8 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token kl_div = ref_per_token_logps - per_token_logps per_token_kl = torch.exp(kl_div) - kl_div - 1 # [B*num_gen, R] per_token_loss = -(torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) - args.beta * per_token_kl) # [B*num_gen, R] - loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() / args.accumulation_steps # scalar + policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar loss.backward() if (step + 1) % args.accumulation_steps == 0: @@ -151,18 +156,20 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token optimizer.zero_grad() if step % args.log_interval == 0 or step == iters: - policy_loss_val = loss.item() + policy_loss_val = loss.item() * args.accumulation_steps + current_aux_loss = aux_loss.item() avg_reward_val = rewards.mean().item() avg_len_val = completion_mask.sum(dim=1).float().mean().item() current_lr = optimizer.param_groups[0]['lr'] - Logger(f'Epoch: {epoch+1}, Step: {step}/{iters}, ' - f'Actor Loss: {policy_loss_val:.6f}, Reward: {avg_reward_val:.6f}, ' - f'Avg Response Len: {avg_len_val:.2f}, LR: {current_lr:.2e}') + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), ' + f'Actor Loss: {policy_loss_val:.4f}, Aux Loss: {current_aux_loss:.4f}, Reward: {avg_reward_val:.4f}, ' + f'Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}') if wandb and is_main_process(): wandb.log({ "policy_loss": policy_loss_val, + "aux_loss": current_aux_loss, "reward": avg_reward_val, "avg_response_len": avg_len_val, "advantages_mean": advantages.mean().item(), diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 89cb7a9..dff0f19 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -39,8 +39,8 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None): Y.view(-1) ).view(Y.size()) - loss = (loss * loss_mask).sum() / loss_mask.sum() - loss += res.aux_loss + logits_loss = (loss * loss_mask).sum() / loss_mask.sum() + loss = logits_loss + res.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -57,12 +57,14 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None): if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_logits_loss = logits_loss.item() + current_aux_loss = res.aux_loss.item() current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_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, "lr": current_lr, "epoch_Time": eta_min}) + if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index cb0ec38..fd094b3 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -143,7 +143,11 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche values = values_seq[torch.arange(values_seq.size(0), device=values_seq.device), last_indices] # [B] advantages = rewards - values.detach() # [B] - logits = actor_model(input_ids=gen_out, attention_mask=full_mask).logits # [B, P+R, V] + with autocast_ctx: + res = actor_model(input_ids=gen_out, attention_mask=full_mask) + logits = res.logits # [B, P+R, V] + aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device) + labels = gen_out[:, 1:].clone() # [B, P+R-1] logp_tokens = F.log_softmax(logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1) # [B, P+R-1] seq_len = gen_out.size(1) - 1 @@ -167,7 +171,7 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche surr2 = torch.clamp(ratio, 1.0 - args.clip_epsilon, 1.0 + args.clip_epsilon) * advantages # [B] policy_loss = -torch.min(surr1, surr2).mean() # scalar value_loss = F.mse_loss(values, rewards) # scalar - loss = policy_loss + args.vf_coef * value_loss + args.kl_coef * kl_ref # scalar + loss = (policy_loss + args.vf_coef * value_loss + args.kl_coef * kl_ref + aux_loss) / args.accumulation_steps # scalar loss.backward() if (step + 1) % args.accumulation_steps == 0: @@ -190,6 +194,7 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche actor_loss_val = policy_loss.item() critic_loss_val = value_loss.item() + current_aux_loss = aux_loss.item() reward_val = rewards.mean().item() kl_val = kl.item() kl_ref_val = kl_ref.item() @@ -201,6 +206,7 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche wandb.log({ "actor_loss": actor_loss_val, "critic_loss": critic_loss_val, + "aux_loss": current_aux_loss, "reward": reward_val, "kl": kl_val, "kl_ref": kl_ref_val, @@ -208,10 +214,10 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche "actor_lr": actor_lr, }) - Logger(f"Epoch: {epoch+1}, Step: {step}/{iters}, " - f"Actor Loss: {actor_loss_val:.6f}, Critic Loss: {critic_loss_val:.6f}, " - f"Reward: {reward_val:.6f}, KL: {kl_val:.6f}, KL_ref: {kl_ref_val:.6f}, " - f"Avg Response Len: {avg_len_val:.2f}, Actor LR: {actor_lr:.2e}, Critic LR: {critic_lr:.2e}") + Logger(f"Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), " + f"Actor Loss: {actor_loss_val:.4f}, Critic Loss: {critic_loss_val:.4f}, Aux Loss: {current_aux_loss:.4f}, " + f"Reward: {reward_val:.4f}, KL: {kl_val:.4f}, KL_ref: {kl_ref_val:.4f}, " + f"Avg Response Len: {avg_len_val:.2f}, Actor LR: {actor_lr:.8f}, Critic LR: {critic_lr:.8f}") if (step + 1) % args.update_old_actor_freq == 0: state_dict = actor_model.module.state_dict() if isinstance(actor_model, DistributedDataParallel) else actor_model.state_dict() diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 5f05341..4596041 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -38,8 +38,8 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): Y.view(-1) ).view(Y.size()) - loss = (loss * loss_mask).sum() / loss_mask.sum() - loss += res.aux_loss + logits_loss = (loss * loss_mask).sum() / loss_mask.sum() + loss = logits_loss + res.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -56,12 +56,14 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_logits_loss = logits_loss.item() + current_aux_loss = res.aux_loss.item() current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_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, "lr": current_lr, "epoch_Time": eta_min}) + if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() diff --git a/trainer/train_spo.py b/trainer/train_spo.py index 37493e4..a15a2ac 100755 --- a/trainer/train_spo.py +++ b/trainer/train_spo.py @@ -155,7 +155,11 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni per_token_logps.append(torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1)) return torch.stack(per_token_logps) - per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1)) # [B, R] + with autocast_ctx: + per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1)) # [B, R] + res = model(outputs) if lm_config.use_moe else None + aux_loss = res.aux_loss if res is not None else torch.tensor(0.0, device=args.device) + with torch.no_grad(): ref_per_token_logps = get_per_token_logps(ref_model, outputs, completion_ids.size(1)) # [B, R] @@ -180,7 +184,8 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni kl_div = ref_per_token_logps - per_token_logps # [B, R] per_token_kl = torch.exp(kl_div) - kl_div - 1 # [B, R] per_token_loss = -per_token_logps * advantages.unsqueeze(1) + args.beta * per_token_kl # [B, R] - loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() / args.accumulation_steps # scalar + policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar loss.backward() response_masks = completion_mask.float() # [B, R] @@ -194,21 +199,23 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni optimizer.zero_grad() if step % args.log_interval == 0 or step == iters: - policy_loss_val = loss.item() + policy_loss_val = loss.item() * args.accumulation_steps + current_aux_loss = aux_loss.item() avg_reward_val = rewards.mean().item() avg_len_val = completion_mask.sum(dim=1).float().mean().item() kl_val = ((per_token_kl * completion_mask).sum() / (completion_mask.sum() + 1e-8)).item() avg_baseline_val = baselines.mean().item() current_lr = optimizer.param_groups[0]['lr'] - Logger(f'Epoch: {epoch+1}, Step: {step}/{iters}, ' - f'Actor Loss: {policy_loss_val:.6f}, Reward: {avg_reward_val:.6f}, ' + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), ' + f'Actor Loss: {policy_loss_val:.4f}, Aux Loss: {current_aux_loss:.4f}, Reward: {avg_reward_val:.4f}, ' f'Baseline: {avg_baseline_val:.4f}, KL: {kl_val:.4f}, Rho: {rho:.4f}, ' - f'Avg Response Len: {avg_len_val:.2f}, LR: {current_lr:.2e}') + f'Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}') if wandb and is_main_process(): wandb.log({ "policy_loss": policy_loss_val, + "aux_loss": current_aux_loss, "reward": avg_reward_val, "kl": kl_val, "rho": float(rho),