[Fix] Add --disable-draft-model-update to control draft model updates(especially in RL) (#15726)

Co-authored-by: leihaodi <haodilei@gmail.com>
This commit is contained in:
Haodi Lei
2026-03-01 04:09:55 +08:00
committed by GitHub
parent 9c81ce4707
commit f451664504
2 changed files with 6 additions and 1 deletions

View File

@@ -1337,6 +1337,8 @@ class UpdateWeightsFromTensorReqInput(BaseReq):
abort_all_requests: bool = False
# Optional: Update weight version along with weights
weight_version: Optional[str] = None
# Optional: Determine whether to disable updating the draft model
disable_draft_model: Optional[bool] = None
@dataclass

View File

@@ -88,7 +88,10 @@ class SchedulerUpdateWeightsMixin:
self: Scheduler, recv_req: UpdateWeightsFromTensorReqInput
):
"""Update the online model parameter from tensors."""
worker = self.draft_worker or self.tp_worker
if recv_req.disable_draft_model:
worker = self.tp_worker
else:
worker = self.draft_worker or self.tp_worker
success, message = worker.update_weights_from_tensor(recv_req)
# TODO extract common code b/t update_weights_from_distributed and update_weights_from_tensor later
if success: