From 508ff1f5c5aeeaa0c4101c1b939a685e4877ccda Mon Sep 17 00:00:00 2001
From: Jani Monoses <jani.monoses@gmail.com>
Date: Wed, 20 Nov 2024 08:14:38 +0200
Subject: [PATCH] Free CPU pinned memory on environment cleanup

---
 vllm/distributed/parallel_state.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index 87ade377266a2..9c5b363fdf426 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -1194,6 +1194,11 @@ 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")
 
 
 def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]: