support non-disturbing remote-instance-weight-loader (#13125)
Signed-off-by: Anqi Shen <amy.saq@antgroup.com>
This commit is contained in:
@@ -2573,6 +2573,9 @@ class Scheduler(
|
||||
self.send_to_detokenizer.send_output(recv_req, recv_req)
|
||||
return None
|
||||
|
||||
def get_remote_instance_transfer_engine_info(self):
|
||||
return self.tp_worker.get_remote_instance_transfer_engine_info()
|
||||
|
||||
|
||||
class IdleSleeper:
|
||||
"""
|
||||
@@ -2686,13 +2689,29 @@ def run_scheduler_process(
|
||||
pp_rank,
|
||||
dp_rank,
|
||||
)
|
||||
pipe_writer.send(
|
||||
{
|
||||
"status": "ready",
|
||||
"max_total_num_tokens": scheduler.max_total_num_tokens,
|
||||
"max_req_input_len": scheduler.max_req_input_len,
|
||||
}
|
||||
)
|
||||
if server_args.remote_instance_weight_loader_support_transfer_engine:
|
||||
(
|
||||
remote_instance_transfer_engine_session_id,
|
||||
remote_instance_transfer_engine_weights_info_dict,
|
||||
) = scheduler.get_remote_instance_transfer_engine_info()
|
||||
pipe_writer.send(
|
||||
{
|
||||
"status": "ready",
|
||||
"max_total_num_tokens": scheduler.max_total_num_tokens,
|
||||
"max_req_input_len": scheduler.max_req_input_len,
|
||||
"tp_rank": tp_rank,
|
||||
"remote_instance_transfer_engine_session_id": remote_instance_transfer_engine_session_id,
|
||||
"remote_instance_transfer_engine_weights_info_dict": remote_instance_transfer_engine_weights_info_dict,
|
||||
}
|
||||
)
|
||||
else:
|
||||
pipe_writer.send(
|
||||
{
|
||||
"status": "ready",
|
||||
"max_total_num_tokens": scheduler.max_total_num_tokens,
|
||||
"max_req_input_len": scheduler.max_req_input_len,
|
||||
}
|
||||
)
|
||||
|
||||
disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode
|
||||
if disaggregation_mode == DisaggregationMode.NULL:
|
||||
|
||||
@@ -366,6 +366,12 @@ class TpModelWorker(BaseTpWorker):
|
||||
can_run_cuda_graph=can_run_cuda_graph,
|
||||
)
|
||||
|
||||
def get_remote_instance_transfer_engine_info(self):
|
||||
return (
|
||||
self.model_runner.remote_instance_transfer_engine_session_id,
|
||||
self.model_runner.remote_instance_transfer_engine_weight_info,
|
||||
)
|
||||
|
||||
def forward_batch_generation(
|
||||
self,
|
||||
model_worker_batch: ModelWorkerBatch,
|
||||
|
||||
Reference in New Issue
Block a user