diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 468fbf9..445101c 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -487,4 +487,6 @@ if __name__ == "__main__": else: rl_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang")) - if dist.is_initialized(): dist.destroy_process_group() + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 8e8714c..085b2a7 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -243,4 +243,6 @@ if __name__ == "__main__": train_epoch(epoch, loader, len(loader), teacher_model, lm_config_student, 0, wandb, args.alpha, args.temperature) # ========== 9. 清理分布进程 ========== - if dist.is_initialized(): dist.destroy_process_group() \ No newline at end of file + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() \ No newline at end of file diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index d541c43..479eb85 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -223,4 +223,6 @@ if __name__ == "__main__": train_epoch(epoch, loader, len(loader), ref_model, lm_config, 0, wandb, args.beta) # ========== 9. 清理分布进程 ========== - if dist.is_initialized(): dist.destroy_process_group() \ No newline at end of file + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() \ No newline at end of file diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 4cef014..7584565 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -168,4 +168,6 @@ if __name__ == "__main__": train_epoch(epoch, loader, len(loader), 0, wandb) # ========== 9. 清理分布进程 ========== - if dist.is_initialized(): dist.destroy_process_group() \ No newline at end of file + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() \ No newline at end of file diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index d39f0d6..200adb6 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -329,4 +329,6 @@ if __name__ == "__main__": grpo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang")) # ========== 9. 清理分布进程 ========== - if dist.is_initialized(): dist.destroy_process_group() \ No newline at end of file + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() \ No newline at end of file diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 9f4c64b..19a1499 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -181,4 +181,6 @@ if __name__ == "__main__": train_epoch(epoch, loader, len(loader), lora_params, 0, wandb) # ========== 10. 清理分布进程 ========== - if dist.is_initialized(): dist.destroy_process_group() \ No newline at end of file + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() \ No newline at end of file diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 07ba5ac..ee71590 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -432,4 +432,6 @@ if __name__ == "__main__": 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")) # ========== 9. 清理分布进程 ========== - if dist.is_initialized(): dist.destroy_process_group() \ No newline at end of file + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() \ No newline at end of file diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 13f67b0..50ba878 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -167,4 +167,6 @@ if __name__ == "__main__": train_epoch(epoch, loader, len(loader), 0, wandb) # ========== 9. 清理分布进程 ========== - if dist.is_initialized(): dist.destroy_process_group() \ No newline at end of file + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() \ No newline at end of file