Publish elastic-EP active ranks from a dedicated step (#25721)
This commit is contained in:
@@ -2923,10 +2923,16 @@ class Scheduler(
|
|||||||
pooled_hidden_states=pooler_output.pooled_hidden_states,
|
pooled_hidden_states=pooler_output.pooled_hidden_states,
|
||||||
)
|
)
|
||||||
|
|
||||||
if (
|
self._maybe_report_active_ranks()
|
||||||
|
|
||||||
|
return ret
|
||||||
|
|
||||||
|
def _maybe_report_active_ranks(self) -> None:
|
||||||
|
if not (
|
||||||
self.server_args.enable_dp_attention
|
self.server_args.enable_dp_attention
|
||||||
and self.server_args.elastic_ep_backend is not None
|
and self.server_args.elastic_ep_backend is not None
|
||||||
):
|
):
|
||||||
|
return
|
||||||
# Get the tensors indicating rank activeness
|
# Get the tensors indicating rank activeness
|
||||||
tp_active_ranks = self.tp_group.active_ranks.detach().cpu().numpy()
|
tp_active_ranks = self.tp_group.active_ranks.detach().cpu().numpy()
|
||||||
tp_active_ranks_cpu = self.tp_group.active_ranks_cpu.detach().numpy()
|
tp_active_ranks_cpu = self.tp_group.active_ranks_cpu.detach().numpy()
|
||||||
@@ -2936,8 +2942,6 @@ class Scheduler(
|
|||||||
ActiveRanksOutput(status=dp_active_ranks.tolist())
|
ActiveRanksOutput(status=dp_active_ranks.tolist())
|
||||||
)
|
)
|
||||||
|
|
||||||
return ret
|
|
||||||
|
|
||||||
def launch_batch_sample_if_needed(
|
def launch_batch_sample_if_needed(
|
||||||
self, batch_result: GenerationBatchResult
|
self, batch_result: GenerationBatchResult
|
||||||
) -> Union[GenerationBatchResult]:
|
) -> Union[GenerationBatchResult]:
|
||||||
|
|||||||
Reference in New Issue
Block a user