Skip to content

GitLab

Unverified Commit 3589aa79 authored by Stefan He's avatar Stefan He Committed by GitHub
Browse files

[RL] Fix illegal memory for _import_static_state (#7733) 


Co-authored-by: default avatarnanjiangwill <willjiang2018@gmail.com>
parent e00715eb
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
python/sglang/srt/managers/scheduler.py
View file @ 3589aa79
......@@ -2346,6 +2346,7 @@ class Scheduler(
self.stashed_model_static_state = _export_static_state(
self.tp_worker.worker.model_runner.model
)
torch.distributed.barrier(self.tp_cpu_group)
self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
return ReleaseMemoryOccupationReqOutput()
......@@ -2357,6 +2358,7 @@ class Scheduler(
if GPU_MEMORY_TYPE_WEIGHTS in tags:
self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
torch.distributed.barrier(self.tp_cpu_group)
_import_static_state(
self.tp_worker.worker.model_runner.model,
self.stashed_model_static_state,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment