[sgl] provide an option to send control req to all dp ranks rank0 (#22758)
This commit is contained in:
@@ -163,7 +163,13 @@ class DataParallelController:
|
|||||||
|
|
||||||
if server_args.enable_dp_attention:
|
if server_args.enable_dp_attention:
|
||||||
self.launch_dp_attention_schedulers(server_args, port_args)
|
self.launch_dp_attention_schedulers(server_args, port_args)
|
||||||
self.control_message_step = server_args.tp_size
|
# When local control broadcast is enabled, send control messages to
|
||||||
|
# every DP group leader (attn_tp_rank=0) so each leader broadcasts
|
||||||
|
# within its own attn_tp_group instead of the full tp_group.
|
||||||
|
# Otherwise fall back to the original behaviour: send to only the
|
||||||
|
# first leader, which then broadcasts over the full tp_group.
|
||||||
|
local_ctrl = server_args.enable_dp_attention_local_control_broadcast
|
||||||
|
self.control_message_step = 1 if local_ctrl else server_args.tp_size
|
||||||
else:
|
else:
|
||||||
self.launch_dp_schedulers(server_args, port_args)
|
self.launch_dp_schedulers(server_args, port_args)
|
||||||
self.control_message_step = 1
|
self.control_message_step = 1
|
||||||
|
|||||||
@@ -1575,7 +1575,28 @@ class Scheduler(
|
|||||||
src=self.attn_cp_group.ranks[0],
|
src=self.attn_cp_group.ranks[0],
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.tp_size != 1:
|
# When dp_attention_local_control_broadcast is enabled, each DP
|
||||||
|
# group leader already receives control messages from the DP
|
||||||
|
# controller, so we broadcast within attn_tp_group + attn_cp_group
|
||||||
|
# instead of the full tp_group. This avoids an expensive
|
||||||
|
# all-ranks gloo sync.
|
||||||
|
_local_ctrl = self.server_args.enable_dp_attention_local_control_broadcast
|
||||||
|
if _local_ctrl:
|
||||||
|
if self.attn_tp_size != 1:
|
||||||
|
control_reqs = broadcast_pyobj(
|
||||||
|
control_reqs,
|
||||||
|
self.attn_tp_group.rank,
|
||||||
|
self.attn_tp_cpu_group,
|
||||||
|
src=self.attn_tp_group.ranks[0],
|
||||||
|
)
|
||||||
|
if self.attn_cp_size != 1:
|
||||||
|
control_reqs = broadcast_pyobj(
|
||||||
|
control_reqs,
|
||||||
|
self.attn_cp_group.rank,
|
||||||
|
self.attn_cp_cpu_group,
|
||||||
|
src=self.attn_cp_group.ranks[0],
|
||||||
|
)
|
||||||
|
elif self.tp_size != 1:
|
||||||
control_reqs = broadcast_pyobj(
|
control_reqs = broadcast_pyobj(
|
||||||
control_reqs,
|
control_reqs,
|
||||||
self.tp_group.rank,
|
self.tp_group.rank,
|
||||||
|
|||||||
@@ -639,6 +639,7 @@ class ServerArgs:
|
|||||||
disable_overlap_schedule: bool = False
|
disable_overlap_schedule: bool = False
|
||||||
enable_mixed_chunk: bool = False
|
enable_mixed_chunk: bool = False
|
||||||
enable_dp_attention: bool = False
|
enable_dp_attention: bool = False
|
||||||
|
enable_dp_attention_local_control_broadcast: bool = False
|
||||||
enable_dp_lm_head: bool = False
|
enable_dp_lm_head: bool = False
|
||||||
enable_two_batch_overlap: bool = False
|
enable_two_batch_overlap: bool = False
|
||||||
enable_single_batch_overlap: bool = False
|
enable_single_batch_overlap: bool = False
|
||||||
@@ -5830,6 +5831,13 @@ class ServerArgs:
|
|||||||
action="store_true",
|
action="store_true",
|
||||||
help="Enabling data parallelism for attention and tensor parallelism for FFN. The dp size should be equal to the tp size. Currently DeepSeek-V2 and Qwen 2/3 MoE models are supported.",
|
help="Enabling data parallelism for attention and tensor parallelism for FFN. The dp size should be equal to the tp size. Currently DeepSeek-V2 and Qwen 2/3 MoE models are supported.",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--enable-dp-attention-local-control-broadcast",
|
||||||
|
action="store_true",
|
||||||
|
help="With DP-attention, send control messages to every DP group leader "
|
||||||
|
"and broadcast within attn_tp_group instead of the full tp_group. "
|
||||||
|
"Eliminates a costly all-ranks gloo sync on every scheduler iteration.",
|
||||||
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--enable-dp-lm-head",
|
"--enable-dp-lm-head",
|
||||||
action="store_true",
|
action="store_true",
|
||||||
|
|||||||
Reference in New Issue
Block a user