[RL] support only do cpu backup on draft model (#13318)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user