From 1e4f464553b59626121cb7a5e0e5d4d2284a4296 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Thu, 23 Oct 2025 15:03:19 +0800 Subject: [PATCH] [fix] sampler-ddp --- trainer/train_distill_reason.py | 1 + trainer/train_distillation.py | 1 + trainer/train_dpo.py | 1 + trainer/train_full_sft.py | 1 + trainer/train_grpo.py | 1 + trainer/train_lora.py | 1 + trainer/train_ppo.py | 1 + trainer/train_pretrain.py | 1 + trainer/train_spo.py | 1 + 9 files changed, 9 insertions(+) diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index ab7866c..9c1c9eb 100644 --- a/trainer/train_distill_reason.py +++ b/trainer/train_distill_reason.py @@ -213,4 +213,5 @@ if __name__ == "__main__": iter_per_epoch = len(train_loader) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) train_epoch(epoch, wandb) diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index f4e2d41..9b6c0a7 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -262,4 +262,5 @@ if __name__ == "__main__": iter_per_epoch = len(train_loader) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) train_epoch(epoch, wandb) diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index c4acf36..cdb9d92 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -245,4 +245,5 @@ if __name__ == "__main__": iter_per_epoch = len(train_loader) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) train_epoch(epoch, wandb) diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 58f87aa..bae0dff 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -199,4 +199,5 @@ if __name__ == "__main__": iter_per_epoch = len(train_loader) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) train_epoch(epoch, wandb) diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 0402f41..8aeb349 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -313,4 +313,5 @@ if __name__ == "__main__": model = DistributedDataParallel(model, device_ids=[ddp_local_rank]) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) grpo_train_epoch(epoch, wandb) diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 5da1607..4c73126 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -205,4 +205,5 @@ if __name__ == "__main__": iter_per_epoch = len(train_loader) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) train_epoch(epoch, wandb) diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 1132da9..f8a8d38 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -367,6 +367,7 @@ if __name__ == "__main__": old_actor_model.to(args.device) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) ppo_train_epoch(epoch, wandb, old_actor_model, ref_model, actor_scheduler, critic_scheduler) if ddp: diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 4166db5..ab6f68f 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -197,4 +197,5 @@ if __name__ == "__main__": iter_per_epoch = len(train_loader) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) train_epoch(epoch, wandb) diff --git a/trainer/train_spo.py b/trainer/train_spo.py index 8db67b9..7ea15bf 100755 --- a/trainer/train_spo.py +++ b/trainer/train_spo.py @@ -364,4 +364,5 @@ if __name__ == "__main__": value_tracker = AutoAdaptiveValueTracker(rho_mode='kl', rho_const=0.9, D_half=0.06, clip_lower=0.5, clip_upper=0.96) for epoch in range(args.epochs): + train_sampler and train_sampler.set_epoch(epoch) spo_train_epoch(epoch, wandb, value_tracker)