Free CPU pinned memory on environment cleanup

vllm-project · Nov 20, 2024 · 9d6dd78 · 9d6dd78
1 parent 343041c
commit 9d6dd78
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
@@ -1194,6 +1194,12 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     gc.collect()
     if not current_platform.is_cpu():
         torch.cuda.empty_cache()
+    try:
+        torch._C._host_emptyCache()
+    except AttributeError:
+        logger.warning(
+            "torch._C._host_emptyCache() only available in Pytorch >=2.5"
+        )  #noqa: E501
 
 
 def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]: