[fix] sampler-ddp

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)

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)

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)

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)

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)

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)

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:

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)

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)