diff --git a/examples/on_policy_distillation/agentic_opd/alfworld/run-alfworld-opd-qwen35-35B-A3B-8xgpu.sh b/examples/on_policy_distillation/agentic_opd/alfworld/run-alfworld-opd-qwen35-35B-A3B-8xgpu.sh index c03272149..d61aa9e5b 100644 --- a/examples/on_policy_distillation/agentic_opd/alfworld/run-alfworld-opd-qwen35-35B-A3B-8xgpu.sh +++ b/examples/on_policy_distillation/agentic_opd/alfworld/run-alfworld-opd-qwen35-35B-A3B-8xgpu.sh @@ -157,7 +157,6 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine ${ROLLOUT_NUM_GPUS_PER_ENGINE:-4} --sglang-mem-fraction-static ${STUDENT_MEM_FRACTION:-0.7} --sglang-load-format dummy - --sglang-enable-weights-cpu-backup --sglang-max-running-requests 64 ) diff --git a/examples/on_policy_distillation/agentic_opd/webshop/run-webshop-opd-qwen35-35B-A3B-8xgpu.sh b/examples/on_policy_distillation/agentic_opd/webshop/run-webshop-opd-qwen35-35B-A3B-8xgpu.sh index 2e1ffc7f6..f04dae9cc 100644 --- a/examples/on_policy_distillation/agentic_opd/webshop/run-webshop-opd-qwen35-35B-A3B-8xgpu.sh +++ b/examples/on_policy_distillation/agentic_opd/webshop/run-webshop-opd-qwen35-35B-A3B-8xgpu.sh @@ -165,7 +165,6 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine ${ROLLOUT_NUM_GPUS_PER_ENGINE:-4} --sglang-mem-fraction-static ${STUDENT_MEM_FRACTION:-0.7} --sglang-load-format dummy - --sglang-enable-weights-cpu-backup --sglang-max-running-requests 64 ) diff --git a/examples/on_policy_distillation/math_opd/run-opd-qwen35-35B-A3B-8xgpu-colocate.sh b/examples/on_policy_distillation/math_opd/run-opd-qwen35-35B-A3B-8xgpu-colocate.sh index 90af97a62..759ee8527 100644 --- a/examples/on_policy_distillation/math_opd/run-opd-qwen35-35B-A3B-8xgpu-colocate.sh +++ b/examples/on_policy_distillation/math_opd/run-opd-qwen35-35B-A3B-8xgpu-colocate.sh @@ -133,7 +133,6 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine 4 --sglang-mem-fraction-static 0.6 --sglang-load-format dummy - --sglang-enable-weights-cpu-backup --sglang-max-running-requests 64 ) diff --git a/examples/on_policy_distillation/mopd/run-mopd-qwen3-vl-2b-8xgpu-colocate.sh b/examples/on_policy_distillation/mopd/run-mopd-qwen3-vl-2b-8xgpu-colocate.sh index e114d6756..ffd68f3b1 100755 --- a/examples/on_policy_distillation/mopd/run-mopd-qwen3-vl-2b-8xgpu-colocate.sh +++ b/examples/on_policy_distillation/mopd/run-mopd-qwen3-vl-2b-8xgpu-colocate.sh @@ -163,7 +163,6 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.8 --sglang-load-format dummy - --sglang-enable-weights-cpu-backup ) RESOURCE_JSON="{\"actor\": [1, ${ACTOR_GPUS}], \"rollout\": [1, ${ROLLOUT_GPUS}], \"teacher\": [1, ${TEACHER_GPUS}]}" diff --git a/examples/on_policy_distillation/mopd/run-mopd-qwen35-35ba3b-16xgpu-colocate.sh b/examples/on_policy_distillation/mopd/run-mopd-qwen35-35ba3b-16xgpu-colocate.sh index 5e9ff2f53..91cd83616 100644 --- a/examples/on_policy_distillation/mopd/run-mopd-qwen35-35ba3b-16xgpu-colocate.sh +++ b/examples/on_policy_distillation/mopd/run-mopd-qwen35-35ba3b-16xgpu-colocate.sh @@ -176,7 +176,6 @@ SGLANG_ARGS=( --sglang-mem-fraction-static 0.6 --sglang-max-running-requests 128 --sglang-load-format dummy - --sglang-enable-weights-cpu-backup ) PARTIAL_ROLLOUT_ARGS=( diff --git a/examples/on_policy_distillation/mopd/run-mopd-qwen35-9b-8xgpu-colocate.sh b/examples/on_policy_distillation/mopd/run-mopd-qwen35-9b-8xgpu-colocate.sh index 660a9ff02..f2648a60d 100755 --- a/examples/on_policy_distillation/mopd/run-mopd-qwen35-9b-8xgpu-colocate.sh +++ b/examples/on_policy_distillation/mopd/run-mopd-qwen35-9b-8xgpu-colocate.sh @@ -163,7 +163,6 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine 2 --sglang-mem-fraction-static 0.7 --sglang-load-format dummy - --sglang-enable-weights-cpu-backup ) RESOURCE_JSON="{\"actor\": [1, ${ACTOR_GPUS}], \"rollout\": [1, ${ROLLOUT_GPUS}], \"teacher\": [1, ${TEACHER_GPUS}]}" diff --git a/examples/on_policy_distillation/vision_opd/run-opd-qwen3.5_35ba3b-122ba10b-128xgpu-colocate.sh b/examples/on_policy_distillation/vision_opd/run-opd-qwen3.5_35ba3b-122ba10b-128xgpu-colocate.sh index 3164c65de..73e4c3a7f 100644 --- a/examples/on_policy_distillation/vision_opd/run-opd-qwen3.5_35ba3b-122ba10b-128xgpu-colocate.sh +++ b/examples/on_policy_distillation/vision_opd/run-opd-qwen3.5_35ba3b-122ba10b-128xgpu-colocate.sh @@ -199,7 +199,6 @@ SGLANG_ARGS=( --sglang-mem-fraction-static ${STUDENT_MEM_FRACTION:-0.7} --sglang-load-format dummy --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256) - --sglang-enable-weights-cpu-backup ) WANDB_ARGS=( diff --git a/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-2teacher-colocate.sh b/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-2teacher-colocate.sh index 8e7afbe08..6d19bff7a 100755 --- a/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-2teacher-colocate.sh +++ b/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-2teacher-colocate.sh @@ -129,7 +129,6 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine "${ROLLOUT_NUM_GPUS_PER_ENGINE:-2}" --sglang-mem-fraction-static "${STUDENT_MEM_FRACTION:-0.8}" --sglang-load-format dummy - --sglang-enable-weights-cpu-backup --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256) ) diff --git a/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-colocate.sh b/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-colocate.sh index 50de1e445..ec3b7f174 100755 --- a/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-colocate.sh +++ b/examples/on_policy_distillation/vision_opd/run-vision-opd-qwen3.5_9b-35ba3b-8xgpu-colocate.sh @@ -133,7 +133,6 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine "${ROLLOUT_NUM_GPUS_PER_ENGINE:-2}" --sglang-mem-fraction-static "${STUDENT_MEM_FRACTION:-0.8}" --sglang-load-format dummy - --sglang-enable-weights-cpu-backup --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256) ) diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index f7d1e710e..7b084569e 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -18,6 +18,7 @@ from relax.utils.opd.opd_utils import ( add_opd_arguments, is_managed_opd_teacher_enabled, + maybe_enable_sglang_weights_cpu_backup, teacher_sglang_parse_args, validate_managed_opd_teacher_colocate_args, validate_opd_args, @@ -3455,6 +3456,7 @@ def slime_validate_args(args): args.offload_train = False if args.offload_rollout is None: args.offload_rollout = False + maybe_enable_sglang_weights_cpu_backup(args) if args.use_critic: args.offload_train = True diff --git a/relax/utils/opd/opd_utils.py b/relax/utils/opd/opd_utils.py index ce052263a..6ab08d0c8 100644 --- a/relax/utils/opd/opd_utils.py +++ b/relax/utils/opd/opd_utils.py @@ -143,6 +143,9 @@ def build_teacher_overrides(args: Any, colocate_sync: bool = False) -> dict[str, overrides["model_path"] = args.teacher_hf_checkpoint overrides.setdefault("load_format", "auto") overrides.setdefault("enable_memory_saver", colocate_sync) + if overrides.get("enable_memory_saver"): + # Release/resume without a CPU backup leaves weight pages uninitialized (uniform output). + overrides.setdefault("enable_weights_cpu_backup", True) return overrides @@ -460,6 +463,18 @@ def validate_managed_opd_teacher_colocate_args(args: Any) -> None: ) +def maybe_enable_sglang_weights_cpu_backup(args: Any) -> None: + # Rollout sleep/wake without a CPU backup leaves weight pages uninitialized. + if not getattr(args, "use_opd", False): + return + if not getattr(args, "offload_rollout", False): + return + if getattr(args, "sglang_enable_weights_cpu_backup", False): + return + logger.info("OPD with offloaded rollout: auto-enabling --sglang-enable-weights-cpu-backup.") + args.sglang_enable_weights_cpu_backup = True + + def add_opd_arguments(parser: Any) -> Any: parser.add_argument( "--use-opd", diff --git a/tests/distributed/ray/test_teacher_sglang_overrides.py b/tests/distributed/ray/test_teacher_sglang_overrides.py index ef2ff6a8d..d07a5c376 100644 --- a/tests/distributed/ray/test_teacher_sglang_overrides.py +++ b/tests/distributed/ray/test_teacher_sglang_overrides.py @@ -6,6 +6,7 @@ from relax.utils.opd.opd_utils import ( build_teacher_engine_args, build_teacher_overrides, + maybe_enable_sglang_weights_cpu_backup, teacher_sglang_parse_args, ) @@ -52,3 +53,35 @@ def add_cli_args(parser): assert args.teacher_sglang_mem_fraction_static == 0.73 assert not hasattr(args, "teacher_sglang_model_path") + + +def test_opd_with_offloaded_rollout_enables_weights_cpu_backup(): + args = SimpleNamespace(use_opd=True, offload_rollout=True) + + maybe_enable_sglang_weights_cpu_backup(args) + + assert args.sglang_enable_weights_cpu_backup is True + + +def test_opd_without_offload_keeps_weights_cpu_backup_untouched(): + args = SimpleNamespace(use_opd=True, offload_rollout=False) + + maybe_enable_sglang_weights_cpu_backup(args) + + assert not hasattr(args, "sglang_enable_weights_cpu_backup") + + +def test_non_opd_keeps_weights_cpu_backup_untouched(): + args = SimpleNamespace(offload_rollout=True) + + maybe_enable_sglang_weights_cpu_backup(args) + + assert not hasattr(args, "sglang_enable_weights_cpu_backup") + + +def test_explicit_weights_cpu_backup_is_preserved(): + args = SimpleNamespace(use_opd=True, offload_rollout=True, sglang_enable_weights_cpu_backup=True) + + maybe_enable_sglang_weights_cpu_backup(args) + + assert args.sglang_enable_weights_cpu_backup is True