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

[refactor] PPO cleanup

pull/825/head
jingyaogong 2026-07-23 11:11:05 +08:00
parent 185baaa022
commit f07588a31b
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
trainer/train_ppo.py
View File

@ -12,7 +12,6 @@ import warnings
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import torch.nn.functional as F import torch.nn.functional as F
from transformers import AutoTokenizer
from contextlib import nullcontext from contextlib import nullcontext
from torch import optim, nn from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
@ -76,7 +75,7 @@ def calculate_rewards(prompts, responses, reward_model):
return rewards return rewards
def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step=0, wandb=None, use_sglang=False): def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step=0, wandb=None):
actor_model.train() actor_model.train()
critic_model.train() critic_model.train()
grad_accum_step = 0 grad_accum_step = 0
@ -442,9 +441,9 @@ if __name__ == "__main__":
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang")) ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb)
else: else:
ppo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang")) ppo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, 0, wandb)
# ========== 9. 清理分布进程 ========== # ========== 9. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():