Fix vllm cleanup

ProbablyFaiz · Oct 1, 2024 · a2713be · a2713be
1 parent 2247ce4
commit a2713be
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/rl/llm/engines/local.py b/rl/llm/engines/local.py
@@ -380,12 +380,15 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         import torch
         import torch.distributed
-        from vllm.model_executor.parallel_utils.parallel_state import (
+        from vllm.distributed.parallel_state import (
+            destroy_distributed_environment,
             destroy_model_parallel,
         )
 
         LOGGER.info("Unloading VLLM model from GPU memory...")
         destroy_model_parallel()
+        destroy_distributed_environment()
+        del self.vllm.model_executor
         del self.vllm
         gc.collect()
         torch.cuda.empty_cache()