diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 4e175e7..0cf844d 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -12,7 +12,6 @@ import warnings import torch import torch.distributed as dist import torch.nn.functional as F -from transformers import AutoTokenizer from contextlib import nullcontext from torch import optim, nn from torch.nn.parallel import DistributedDataParallel @@ -76,7 +75,7 @@ def calculate_rewards(prompts, responses, reward_model): 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() critic_model.train() 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) if skip > 0: 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: - 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. 清理分布进程 ========== if dist.is_initialized():