xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-03-26 22:35:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jani Monoses 2025-01-21 21:56:41 +02:00 committed by GitHub
parent fa9ee08121
commit 9c485d9e25
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/distributed/parallel_state.py
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],