From f4516645041137d4ee826fdda7efc05b8fba4819 Mon Sep 17 00:00:00 2001 From: Haodi Lei <152671156+bingyang-lei@users.noreply.github.com> Date: Sun, 1 Mar 2026 04:09:55 +0800 Subject: [PATCH] [Fix] Add --disable-draft-model-update to control draft model updates(especially in RL) (#15726) Co-authored-by: leihaodi --- python/sglang/srt/managers/io_struct.py | 2 ++ python/sglang/srt/managers/scheduler_update_weights_mixin.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 9326c691c..89fd89496 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py index 293a84350..ba038877b 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py @@ -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: