[RL] support only do cpu backup on draft model (#13318)

This commit is contained in:
Zilin Zhu
2025-11-15 16:31:57 +08:00
committed by GitHub
parent 2a96e302cb
commit 56fc483073
2 changed files with 11 additions and 2 deletions

View File

@@ -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,

View File

@@ -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",