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

Fix torch memory saver not enabled in DP scenario (#5560)

Browse Source
This commit is contained in:
fzyzcjy
2025-04-21 05:20:52 +08:00
committed by GitHub
parent 463d4b7400
commit fa2f677e18
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/managers/data_parallel_controller.py
View File

@@ -30,6 +30,7 @@ from sglang.srt.managers.io_struct import (
) )
from sglang.srt.managers.scheduler import run_scheduler_process from sglang.srt.managers.scheduler import run_scheduler_process
from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils import bind_port, configure_logger, get_zmq_socket from sglang.srt.utils import bind_port, configure_logger, get_zmq_socket
from sglang.utils import get_exception_traceback from sglang.utils import get_exception_traceback
@@ -174,6 +175,10 @@ class DataParallelController:
if not server_args.enable_dp_attention: if not server_args.enable_dp_attention:
logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.") logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=server_args.enable_memory_saver
)
# Launch tensor parallel scheduler processes # Launch tensor parallel scheduler processes
scheduler_pipe_readers = [] scheduler_pipe_readers = []
tp_size_per_node = server_args.tp_size // server_args.nnodes tp_size_per_node = server_args.tp_size // server_args.nnodes
@@ -208,7 +213,8 @@ class DataParallelController:
target=run_scheduler_process, target=run_scheduler_process,
args=(server_args, rank_port_args, gpu_id, tp_rank, dp_rank, writer), args=(server_args, rank_port_args, gpu_id, tp_rank, dp_rank, writer),
) )
proc.start() with memory_saver_adapter.configure_subprocess():
proc.start()
self.scheduler_procs.append(proc) self.scheduler_procs.append(proc)
scheduler_pipe_readers.append(reader) scheduler_pipe_readers.append(reader)