mirror of
https://github.com/jingyaogong/minimind.git
synced 2026-09-28 05:47:23 +00:00
[fix] residual step order
This commit is contained in:
@@ -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 = (policy_loss + aux_loss) / args.accumulation_steps # scalar
|
||||||
loss.backward()
|
loss.backward()
|
||||||
|
|
||||||
if step % args.accumulation_steps == 0:
|
if step % args.accumulation_steps == 0 or step == iters:
|
||||||
if args.grad_clip > 0:
|
if args.grad_clip > 0:
|
||||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||||
optimizer.step()
|
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 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
|
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__":
|
if __name__ == "__main__":
|
||||||
parser = argparse.ArgumentParser(description="MiniMind GRPO (Group Relative Policy Optimization)")
|
parser = argparse.ArgumentParser(description="MiniMind GRPO (Group Relative Policy Optimization)")
|
||||||
|
|||||||
Reference in New Issue
Block a user