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

This commit is contained in:
Jani Monoses
2025-01-21 21:56:41 +02:00
committed by GitHub
parent fa9ee08121
commit 9c485d9e25

View File

@ -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],