diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 3ba5d352b..412414924 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -756,9 +756,12 @@ class ModelRunner: # Remove monkey_patch when linear.py quant remove dependencies with vllm monkey_patch_vllm_parallel_state() + enable_cpu_backup = self.server_args.enable_weights_cpu_backup or ( + self.is_draft_worker and self.server_args.enable_draft_weights_cpu_backup + ) with self.memory_saver_adapter.region( GPU_MEMORY_TYPE_WEIGHTS, - enable_cpu_backup=self.server_args.enable_weights_cpu_backup, + enable_cpu_backup=enable_cpu_backup, ): self.model = get_model( model_config=self.model_config, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index fe7fc1a6b..45663eeee 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -511,6 +511,7 @@ class ServerArgs: delete_ckpt_after_loading: bool = False enable_memory_saver: bool = False enable_weights_cpu_backup: bool = False + enable_draft_weights_cpu_backup: bool = False allow_auto_truncate: bool = False enable_custom_logit_processor: bool = False flashinfer_mla_disable_ragged: bool = False @@ -3409,7 +3410,12 @@ class ServerArgs: parser.add_argument( "--enable-weights-cpu-backup", action="store_true", - help="Save model weights to CPU memory during release_weights_occupation and resume_weights_occupation", + help="Save model weights (both main model and draft model, if any) to CPU memory during release_weights_occupation and resume_weights_occupation", + ) + parser.add_argument( + "--enable-draft-weights-cpu-backup", + action="store_true", + help="Save draft model weights to CPU memory during release_weights_occupation and resume_weights_occupation", ) parser.add_argument( "--allow-auto-truncate",