EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[Fix] Check the device backend before calling empty_cache function (#4212)

Browse Source
This commit is contained in:
Chen Shengzhi
2025-03-13 12:28:48 +08:00
committed by GitHub
parent 0a59a4657a
commit f1cf6eefbe
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/distributed/parallel_state.py
View File

@@ -1228,7 +1228,16 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
ray.shutdown()
gc.collect()
if not current_platform.is_cpu():
torch.cuda.empty_cache()
if hasattr(torch, "cuda") and torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch._C, "_host_emptyCache"):
torch._C._host_emptyCache()
else:
logger.warning(
"torch._C._host_emptyCache() only available in Pytorch >=2.5"
)
elif hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]: