Skip to content

Commit 9c485d9

Browse files
authored
[Core] Free CPU pinned memory on environment cleanup (#10477)
1 parent fa9ee08 commit 9c485d9

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

vllm/distributed/parallel_state.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1183,6 +1183,11 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
11831183
from vllm.platforms import current_platform
11841184
if not current_platform.is_cpu():
11851185
torch.cuda.empty_cache()
1186+
try:
1187+
torch._C._host_emptyCache()
1188+
except AttributeError:
1189+
logger.warning(
1190+
"torch._C._host_emptyCache() only available in Pytorch >=2.5")
11861191

11871192

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

0 commit comments

Comments
 (0)