[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:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user