[Core] Free CPU pinned memory on environment cleanup (#10477)

9c485d9e · Jani Monoses · GitHub · fa9ee081 · 9c485d9e
Unverified Commit 9c485d9e authored Jan 21, 2025 by Jani Monoses Committed by GitHub Jan 21, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

vllm/distributed/parallel_state.py vllm/distributed/parallel_state.py +5 -0

No files found.
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -1183,6 +1183,11 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
    from vllm.platforms import current_platform
    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")


 def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],