Set CUDA_VISIBLE_DEVICES to achieve one GPU per process (#9170)

Co-authored-by: SangBin Cho <rkooo567@gmail.com>
Co-authored-by: Cheng Wan <cwan@x.ai>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
This commit is contained in:
Lianmin Zheng
2025-10-17 17:30:06 -07:00
committed by GitHub
parent 69fe3c9726
commit 9eefe2c0b7
7 changed files with 69 additions and 51 deletions

View File

@@ -75,6 +75,7 @@ from sglang.srt.utils import (
is_cuda,
kill_process_tree,
launch_dummy_health_check_server,
maybe_reindex_device_id,
prepare_model_and_tokenizer,
set_prometheus_multiproc_dir,
set_ulimit,
@@ -782,22 +783,24 @@ def _launch_subprocesses(
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
)
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)
proc = mp.Process(
target=run_scheduler_process,
args=(
server_args,
port_args,
gpu_id,
tp_rank,
moe_ep_rank,
pp_rank,
None,
writer,
),
)
with memory_saver_adapter.configure_subprocess():
proc.start()
with maybe_reindex_device_id(gpu_id) as gpu_id:
proc = mp.Process(
target=run_scheduler_process,
args=(
server_args,
port_args,
gpu_id,
tp_rank,
moe_ep_rank,
pp_rank,
None,
writer,
),
)
with memory_saver_adapter.configure_subprocess():
proc.start()
scheduler_procs.append(proc)
scheduler_pipe_readers.append(reader)
else: