From a6ffddc0ea01313d61a016fb7b0a7f35bc6bfb22 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Mon, 21 Sep 2026 23:24:55 -0700 Subject: [PATCH] fix: update residual gradients before saving final checkpoints --- trainer/train_agent.py | 8 +------- trainer/train_distillation.py | 11 +---------- trainer/train_dpo.py | 11 +---------- trainer/train_full_sft.py | 11 +---------- trainer/train_lora.py | 11 +---------- trainer/train_pretrain.py | 11 +---------- 6 files changed, 6 insertions(+), 57 deletions(-) diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 7ef6ed1..4619dbb 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -240,12 +240,10 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa # ================================ 工具与 Reward = End ================================ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model=None, start_step=0, wandb=None, use_sglang=False): - last_step = start_step for step, batch in enumerate(loader, start=start_step + 1): messages_batch = batch['messages'] tools_batch = batch['tools'] gt_batch = batch['gt'] - last_step = step with torch.no_grad(): completions, contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch = rollout_batch(rollout_engine, tokenizer, messages_batch, tools_batch, args.num_generations, max_turns=3, max_new_tokens=args.max_gen_len, thinking_ratio=args.thinking_ratio, device=args.device) @@ -331,7 +329,7 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model loss = (policy_loss + aux_loss) / args.accumulation_steps loss.backward() - if step % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0 or step == iters: if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step(); scheduler.step(); optimizer.zero_grad() @@ -365,10 +363,6 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model del per_token_logps, ref_per_token_logps del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask - if last_step > start_step and last_step % args.accumulation_steps != 0: - if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step(); scheduler.step(); optimizer.zero_grad() - if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind Agent RL") diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 085b2a7..dce5346 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -38,14 +38,12 @@ def distillation_loss(student_logits, teacher_logits, temperature=1.0, reduction def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_step=0, wandb=None, alpha=0.0, temperature=1.0): start_time = time.time() - last_step = start_step if teacher_model is not None: teacher_model.eval() teacher_model.requires_grad_(False) for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): - last_step = step input_ids = input_ids.to(args.device) labels = labels.to(args.device) loss_mask = (labels[..., 1:] != -100).float() @@ -94,7 +92,7 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st scaler.scale(loss).backward() - if step % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0 or step == iters: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) scaler.step(optimizer) @@ -135,13 +133,6 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st del input_ids, labels, loss_mask, res, student_logits, ce_loss, distill_loss, loss - if last_step > start_step and last_step % args.accumulation_steps != 0: - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad(set_to_none=True) - if __name__ == "__main__": # 模拟用moe模型蒸馏dense模型，也可以用更大teacher_hidden_size模型蒸馏更小student_hidden_size的 diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 479eb85..69fa773 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -52,10 +52,8 @@ def dpo_loss(ref_log_probs, policy_log_probs, mask, beta): 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 for step, batch in enumerate(loader, start=start_step + 1): - last_step = step x_chosen = batch['x_chosen'].to(args.device) x_rejected = batch['x_rejected'].to(args.device) y_chosen = batch['y_chosen'].to(args.device) @@ -86,7 +84,7 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= scaler.scale(loss).backward() - if step % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0 or step == iters: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) scaler.step(optimizer) @@ -120,13 +118,6 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= 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 - if last_step > start_step and last_step % args.accumulation_steps != 0: - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad(set_to_none=True) - if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind DPO (Direct Preference Optimization)") diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index c22eb80..2610d1a 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -23,11 +23,9 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): start_time = time.time() - last_step = start_step for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): input_ids = input_ids.to(args.device) labels = labels.to(args.device) - last_step = step lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr @@ -39,7 +37,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): scaler.scale(loss).backward() - if step % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0 or step == iters: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) @@ -73,13 +71,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): del input_ids, labels, res, loss - if last_step > start_step and last_step % args.accumulation_steps != 0: - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad(set_to_none=True) - if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind Full SFT") diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 19a1499..4717292 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -24,11 +24,9 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None): start_time = time.time() - last_step = start_step for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): input_ids = input_ids.to(args.device) labels = labels.to(args.device) - last_step = step lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr @@ -40,7 +38,7 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None): scaler.scale(loss).backward() - if step % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0 or step == iters: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip) scaler.step(optimizer) @@ -68,13 +66,6 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None): del input_ids, labels, res, loss - if last_step > start_step and last_step % args.accumulation_steps != 0: - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad(set_to_none=True) - if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning") parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录") diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 83a0382..a5d2fa2 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -23,11 +23,9 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): start_time = time.time() - last_step = start_step for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): input_ids = input_ids.to(args.device) labels = labels.to(args.device) - last_step = step lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr @@ -39,7 +37,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): scaler.scale(loss).backward() - if step % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0 or step == iters: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) @@ -72,13 +70,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): del input_ids, labels, res, loss - if last_step > start_step and last_step % args.accumulation_steps != 0: - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad(set_to_none=True) - if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind Pretraining")