chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support multi-node DP attention (#2925)

Browse Source 
Co-authored-by: dhou-xai <dhou@x.ai>
This commit is contained in:
Lianmin Zheng
2025-01-16 11:15:00 -08:00
committed by GitHub
parent 58f3f2b840
commit 8b6ce52e92
16 changed files with 287 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -23,6 +23,7 @@ import psutil
import setproctitle
import zmq
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
from sglang.srt.managers.io_struct import (
TokenizedEmbeddingReqInput,
TokenizedGenerateReqInput,
@@ -63,9 +64,10 @@ class DataParallelController:
# Init inter-process communication
self.context = zmq.Context(1 + server_args.dp_size)
self.recv_from_tokenizer = get_zmq_socket(
self.context, zmq.PULL, port_args.scheduler_input_ipc_name
)
if server_args.node_rank == 0:
self.recv_from_tokenizer = get_zmq_socket(
self.context, zmq.PULL, port_args.scheduler_input_ipc_name
)
# Dispatch method
self.round_robin_counter = 0
@@ -75,33 +77,47 @@ class DataParallelController:
}
self.dispatching = dispatch_lookup[self.load_balance_method]
# Start data parallel workers
base_gpu_id = 0
# Launch data parallel workers
self.scheduler_procs = []
self.workers = [None] * server_args.dp_size
if not server_args.enable_dp_attention:
dp_port_args = self.launch_dp_schedulers(server_args, port_args)
else:
dp_port_args = self.launch_dp_attention_schedulers(server_args, port_args)
# Only node rank 0 runs the real data parallel controller that dispatches the requests.
if server_args.node_rank == 0:
for dp_rank in range(server_args.dp_size):
self.workers[dp_rank] = get_zmq_socket(
self.context,
zmq.PUSH,
dp_port_args[dp_rank].scheduler_input_ipc_name,
)
def launch_dp_schedulers(self, server_args, port_args):
base_gpu_id = 0
threads = []
sockets = []
dp_port_args = []
for dp_rank in range(server_args.dp_size):
tmp_port_args = PortArgs.init_new(server_args)
tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
dp_port_args.append(tmp_port_args)
if server_args.enable_dp_attention:
# Data parallelism resues the tensor parallelism group,
# so all dp ranks should use the same nccl port.
tmp_port_args.nccl_port = port_args.nccl_port
else:
# This port is checked free in PortArgs.init_new.
# We hold it first so that the next dp worker gets a different port
sockets.append(bind_port(tmp_port_args.nccl_port))
# This port is checked free in PortArgs.init_new.
# We hold it first so that the next dp worker gets a different port
sockets.append(bind_port(tmp_port_args.nccl_port))
# Create a thread for each worker
thread = threading.Thread(
target=self.launch_worker_func,
target=self.launch_tensor_parallel_group,
args=(server_args, tmp_port_args, base_gpu_id, dp_rank),
)
threads.append(thread)
base_gpu_id += 1 if server_args.enable_dp_attention else server_args.tp_size
base_gpu_id += server_args.tp_size
# Free all sockets before starting the threads to launch TP workers
for sock in sockets:
@@ -113,26 +129,14 @@ class DataParallelController:
for thread in threads:
thread.join()
def launch_worker_func(
self,
server_args: ServerArgs,
port_args: PortArgs,
base_gpu_id: int,
dp_rank: int,
):
logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
return dp_port_args
launch_func_ = (
self.launch_tensor_parallel_process
if server_args.enable_dp_attention
else self.launch_tensor_parallel_group
)
self.workers[dp_rank] = launch_func_(
server_args,
port_args,
base_gpu_id,
dp_rank,
)
def launch_dp_attention_schedulers(self, server_args, port_args):
self.launch_tensor_parallel_group(server_args, port_args, 0, None)
dp_port_args = []
for dp_rank in range(server_args.dp_size):
dp_port_args.append(PortArgs.init_new(server_args, dp_rank))
return dp_port_args
def launch_tensor_parallel_group(
self,
@@ -141,8 +145,10 @@ class DataParallelController:
base_gpu_id: int,
dp_rank: int,
):
if not server_args.enable_dp_attention:
logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
# Launch tensor parallel scheduler processes
scheduler_procs = []
scheduler_pipe_readers = []
tp_size_per_node = server_args.tp_size // server_args.nnodes
tp_rank_range = range(
@@ -150,53 +156,39 @@ class DataParallelController:
tp_size_per_node * (server_args.node_rank + 1),
)
for tp_rank in tp_rank_range:
rank_port_args = port_args
if server_args.enable_dp_attention:
# dp attention has different sharding logic
_, _, dp_rank = compute_dp_attention_world_info(
server_args.enable_dp_attention,
tp_rank,
server_args.tp_size,
server_args.dp_size,
)
# compute zmq ports for this dp rank
rank_port_args = PortArgs.init_new(server_args, dp_rank)
# Data parallelism resues the tensor parallelism group,
# so all dp ranks should use the same nccl port.
rank_port_args.nccl_port = port_args.nccl_port
reader, writer = mp.Pipe(duplex=False)
gpu_id = server_args.base_gpu_id + base_gpu_id + tp_rank % tp_size_per_node
proc = mp.Process(
target=run_scheduler_process,
args=(server_args, port_args, gpu_id, tp_rank, dp_rank, writer),
args=(server_args, rank_port_args, gpu_id, tp_rank, dp_rank, writer),
)
proc.start()
scheduler_procs.append(proc)
self.scheduler_procs.append(proc)
scheduler_pipe_readers.append(reader)
send_to = get_zmq_socket(
self.context, zmq.PUSH, port_args.scheduler_input_ipc_name
)
# Wait for model to finish loading and get max token nums
# Wait for model to finish loading
scheduler_info = []
for i in range(len(scheduler_pipe_readers)):
scheduler_info.append(scheduler_pipe_readers[i].recv())
self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
return send_to
def launch_tensor_parallel_process(
self,
server_args: ServerArgs,
port_args: PortArgs,
base_gpu_id: int,
dp_rank: int,
):
reader, writer = mp.Pipe(duplex=False)
gpu_id = base_gpu_id
tp_rank = dp_rank
proc = mp.Process(
target=run_scheduler_process,
args=(server_args, port_args, gpu_id, tp_rank, dp_rank, writer),
)
proc.start()
send_to = get_zmq_socket(
self.context, zmq.PUSH, port_args.scheduler_input_ipc_name
)
scheduler_info = reader.recv()
self.max_total_num_tokens = scheduler_info["max_total_num_tokens"]
return send_to
def round_robin_scheduler(self, req):
self.workers[self.round_robin_counter].send_pyobj(req)
self.round_robin_counter = (self.round_robin_counter + 1) % len(self.workers)
@@ -221,8 +213,8 @@ class DataParallelController:
):
self.dispatching(recv_req)
else:
# Send other control messages to all workers
for worker in self.workers:
# Send other control messages to first worker of tp group
for worker in self.workers[:: self.server_args.tp_size]:
worker.send_pyobj(recv_req)
@@ -240,7 +232,13 @@ def run_data_parallel_controller_process(
pipe_writer.send(
{"status": "ready", "max_total_num_tokens": controller.max_total_num_tokens}
)
controller.event_loop()
if server_args.node_rank == 0:
controller.event_loop()
for proc in controller.scheduler_procs:
proc.join()
logger.error(
f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}"
)
except Exception:
traceback = get_exception_traceback()
logger.error(f"DataParallelController hit an exception: {traceback}")

7
python/sglang/srt/managers/schedule_batch.py
View File

@@ -1003,6 +1003,11 @@ class ScheduleBatch:
self.req_pool_indices = torch.empty(0, dtype=torch.int32, device=self.device)
self.seq_lens_sum = 0
self.extend_num_tokens = 0
self.sampling_info = SamplingBatchInfo.from_schedule_batch(
self,
self.model_config.vocab_size,
enable_overlap_schedule=self.enable_overlap,
)
def prepare_for_decode(self):
self.forward_mode = ForwardMode.DECODE
@@ -1117,7 +1122,7 @@ class ScheduleBatch:
self.spec_info.merge_batch(other.spec_info)
def get_model_worker_batch(self):
if self.forward_mode.is_decode() or self.forward_mode.is_idle():
if self.forward_mode.is_decode_or_idle():
extend_seq_lens = extend_prefix_lens = extend_logprob_start_lens = None
else:
extend_seq_lens = self.extend_lens

83
python/sglang/srt/managers/scheduler.py
View File

@@ -33,6 +33,7 @@ import zmq
from sglang.global_config import global_config
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.io_struct import (
AbortReq,
@@ -135,7 +136,17 @@ class Scheduler:
# Init inter-process communication
context = zmq.Context(2)
if self.tp_rank == 0 or self.server_args.enable_dp_attention:
self.dp_size = server_args.dp_size
self.attn_tp_rank, self.attn_tp_size, self.dp_rank = (
compute_dp_attention_world_info(
server_args.enable_dp_attention,
self.tp_rank,
self.tp_size,
self.dp_size,
)
)
if self.attn_tp_rank == 0:
self.recv_from_tokenizer = get_zmq_socket(
context, zmq.PULL, port_args.scheduler_input_ipc_name
)
@@ -244,6 +255,7 @@ class Scheduler:
_,
) = self.tp_worker.get_worker_info()
self.tp_cpu_group = self.tp_worker.get_tp_cpu_group()
self.attn_tp_cpu_group = self.tp_worker.get_attention_tp_cpu_group()
self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func()
global_server_args_dict.update(worker_global_server_args_dict)
set_random_seed(self.random_seed)
@@ -447,6 +459,10 @@ class Scheduler:
self.process_input_requests(recv_reqs)
batch = self.get_next_batch_to_run()
if self.server_args.enable_dp_attention: # TODO: simplify this
batch = self.prepare_dp_attn_batch(batch)
self.cur_batch = batch
if batch:
@@ -479,7 +495,7 @@ class Scheduler:
def recv_requests(self) -> List[Req]:
"""Receive results at tp_rank = 0 and broadcast it to all other TP ranks."""
if self.tp_rank == 0 or self.server_args.enable_dp_attention:
if self.attn_tp_rank == 0:
recv_reqs = []
while True:
@@ -491,7 +507,40 @@ class Scheduler:
else:
recv_reqs = None
if self.tp_size != 1 and not self.server_args.enable_dp_attention:
if self.server_args.enable_dp_attention:
if self.attn_tp_rank == 0:
work_reqs = [
req
for req in recv_reqs
if isinstance(
req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)
)
]
control_reqs = [
req
for req in recv_reqs
if not isinstance(
req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)
)
]
else:
work_reqs = None
control_reqs = None
if self.attn_tp_size != 1:
attn_tp_rank_0 = self.dp_rank * self.attn_tp_size
work_reqs = broadcast_pyobj(
work_reqs,
self.attn_tp_rank,
self.attn_tp_cpu_group,
src=attn_tp_rank_0,
)
if self.tp_size != 1:
control_reqs = broadcast_pyobj(
control_reqs, self.tp_rank, self.tp_cpu_group
)
recv_reqs = work_reqs + control_reqs
elif self.tp_size != 1:
recv_reqs = broadcast_pyobj(recv_reqs, self.tp_rank, self.tp_cpu_group)
return recv_reqs
@@ -887,7 +936,7 @@ class Scheduler:
self.being_chunked_req.is_being_chunked += 1
# Print stats
if self.tp_rank == 0:
if self.attn_tp_rank == 0:
self.log_prefill_stats(adder, can_run_list, running_bs, has_being_chunked)
# Create a new batch
@@ -974,7 +1023,7 @@ class Scheduler:
self.forward_ct += 1
if self.is_generation:
if batch.forward_mode.is_decode() or batch.extend_num_tokens != 0:
if batch.forward_mode.is_decode_or_idle() or batch.extend_num_tokens != 0:
if self.spec_algorithm.is_none():
model_worker_batch = batch.get_model_worker_batch()
logits_output, next_token_ids = (
@@ -988,18 +1037,8 @@ class Scheduler:
num_accepted_tokens,
) = self.draft_worker.forward_batch_speculative_generation(batch)
self.num_generated_tokens += num_accepted_tokens
elif batch.forward_mode.is_idle():
model_worker_batch = batch.get_model_worker_batch()
self.tp_worker.forward_batch_idle(model_worker_batch)
return
else:
logits_output = None
if self.skip_tokenizer_init:
next_token_ids = torch.full(
(batch.batch_size(),), self.tokenizer.eos_token_id
)
else:
next_token_ids = torch.full((batch.batch_size(),), 0)
assert False, "batch.extend_num_tokens == 0, this is unexpected!"
batch.output_ids = next_token_ids
ret = logits_output, next_token_ids, model_worker_batch.bid
else: # embedding or reward model
@@ -1016,6 +1055,9 @@ class Scheduler:
self.running_batch = None
elif batch.forward_mode.is_extend():
self.process_batch_result_prefill(batch, result)
elif batch.forward_mode.is_idle():
if self.enable_overlap:
self.tp_worker.resolve_batch_result(result[-1])
elif batch.forward_mode.is_dummy_first():
batch.next_batch_sampling_info.update_regex_vocab_mask()
self.current_stream.synchronize()
@@ -1166,7 +1208,7 @@ class Scheduler:
self.forward_ct_decode = (self.forward_ct_decode + 1) % (1 << 30)
if (
self.tp_rank == 0
self.attn_tp_rank == 0
and self.forward_ct_decode % self.server_args.decode_log_interval == 0
):
self.log_decode_stats()
@@ -1402,12 +1444,7 @@ class Scheduler:
# Check forward mode for cuda graph
if not self.server_args.disable_cuda_graph:
forward_mode_state = torch.tensor(
(
1
if local_batch.forward_mode.is_decode()
or local_batch.forward_mode.is_idle()
else 0
),
(1 if local_batch.forward_mode.is_decode_or_idle() else 0),
dtype=torch.int32,
)
torch.distributed.all_reduce(

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

@@ -101,6 +101,7 @@ class TpModelWorker:
self.max_total_num_tokens // 2
if server_args.max_running_requests is None
else server_args.max_running_requests
// (server_args.dp_size if server_args.enable_dp_attention else 1)
),
self.model_runner.req_to_token_pool.size,
)
@@ -142,16 +143,15 @@ class TpModelWorker:
def get_tp_cpu_group(self):
return self.model_runner.tp_group.cpu_group
def get_attention_tp_cpu_group(self):
return self.model_runner.attention_tp_group.cpu_group
def get_memory_pool(self):
return (
self.model_runner.req_to_token_pool,
self.model_runner.token_to_kv_pool,
)
def forward_batch_idle(self, model_worker_batch: ModelWorkerBatch):
forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
self.model_runner.forward(forward_batch)
def forward_batch_generation(
self,
model_worker_batch: ModelWorkerBatch,

3
python/sglang/srt/managers/tp_worker_overlap_thread.py
View File

@@ -92,6 +92,9 @@ class TpModelWorkerClient:
def get_tp_cpu_group(self):
return self.worker.get_tp_cpu_group()
def get_attention_tp_cpu_group(self):
return self.worker.get_attention_tp_cpu_group()
def get_memory_pool(self):
return (
self.worker.model_runner.req_to_token_pool,