obscura/minimind
1
0
Fork 0
mirror of https://github.com/jingyaogong/minimind.git synced 2026-09-26 04:47:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #868 from DaoyuanLi2816/fix/final-accumulation-checkpoint

Browse Source 
[fix] save final checkpoints after the residual accumulation step
This commit is contained in:
jingyaogong
2026-09-22 23:39:01 +08:00
committed by GitHub
parent 7da0f92ff6 a6ffddc0ea
commit 6fff76fe04
6 changed files with 6 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
trainer/train_agent.py
+1 -7
View File
@@ -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")
trainer/train_distillation.py
+1 -10
View File
@@ -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的
trainer/train_dpo.py
+1 -10
View File
@@ -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)")
trainer/train_full_sft.py
+1 -10
View File
@@ -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")
trainer/train_lora.py
+1 -10
View File
@@ -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="模型保存目录")
trainer/train_pretrain.py
+1 -10
View File
@@ -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")