github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[fix] residual step order

master
jingyaogong 2026-09-23 00:07:30 +08:00
parent 6fff76fe04
commit f659b55761
1 changed files with 1 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
trainer/train_grpo.py
View File

@ -144,7 +144,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar
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()
@ -195,13 +195,6 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask, completion_pad_mask, prompt_lens, logp_pos
if step > start_step and 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 GRPO (Group Relative Policy Optimization)")