diff --git a/test/srt/rl/test_update_weights_from_distributed.py b/test/srt/rl/test_update_weights_from_distributed.py index f2c6e5131..45d2baf84 100644 --- a/test/srt/rl/test_update_weights_from_distributed.py +++ b/test/srt/rl/test_update_weights_from_distributed.py @@ -187,6 +187,9 @@ def init_process_hf( print(f"[hf] {rank=} {broadcast_time=:.3f}s") param_queue.put(("broadcast_time", broadcast_time)) + # Destroy process group and release related resource + torch.distributed.destroy_process_group(group) + # Delete the huggingface models to free up memory. del hf_instruct_model del hf_base_model