From 5929cb85aa17bc15a9415185c207a4c34f8dbcd3 Mon Sep 17 00:00:00 2001 From: Ziyi <1034337098@qq.com> Date: Fri, 21 Aug 2026 02:08:41 +0000 Subject: [PATCH 1/2] fix(opd): keep teacher weights across sleep/wake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Preserve colocate managed teacher weights across memory-saver sleep/wake - Managed OPD/SDPO teacher enables enable_memory_saver in colocate mode but did not enable enable_weights_cpu_backup; every release/resume cycle re-allocates its weight pages uninitialized, so the teacher serves a uniform distribution (all logprobs = -ln(vocab_size)) - Default enable_weights_cpu_backup=True when memory saver is on; explicit --teacher-sglang-enable-weights-cpu-backup still overrides; dedicated-PG teachers are unaffected --- relax/utils/opd/opd_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/relax/utils/opd/opd_utils.py b/relax/utils/opd/opd_utils.py index ce052263a..c69a5375b 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 From c859c3be093e0cb80516b123ce0f8d4f3b62b5b1 Mon Sep 17 00:00:00 2001 From: Ziyi <1034337098@qq.com> Date: Wed, 26 Aug 2026 10:17:14 +0000 Subject: [PATCH 2/2] fix(opd): auto-enable rollout weights cpu backup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Default --sglang-enable-weights-cpu-backup on for OPD - Add maybe_enable_sglang_weights_cpu_backup: when --use-opd and offload_rollout resolves True, default the student rollout SGLang engines to enable_weights_cpu_backup so sleep/wake cannot leave weight pages uninitialized (same failure mode as the teacher-side fix) - Call it in slime_validate_args right after offload_rollout is fully resolved, so colocate runs (incl. managed-teacher split bundles) get it automatically while hybrid / debug_rollout_only / non-offloaded dedicated rollout are left untouched; an explicit flag or per-engine --sglang-config override still wins --- # 🔩 Chore ## Drop redundant flags from OPD launch scripts - Remove manually-added --sglang-enable-weights-cpu-backup from the 9 OPD example scripts (vision_opd / mopd / math_opd / agentic webshop+alfworld); the GRPO baseline scripts keep theirs since they do not pass --use-opd --- # ✅ Tests ## Cover maybe_enable_sglang_weights_cpu_backup - Enables when use_opd + offload_rollout - No-op without OPD or without offload - Idempotent when explicitly set --- .../run-alfworld-opd-qwen35-35B-A3B-8xgpu.sh | 1 - .../run-webshop-opd-qwen35-35B-A3B-8xgpu.sh | 1 - .../run-opd-qwen35-35B-A3B-8xgpu-colocate.sh | 1 - .../run-mopd-qwen3-vl-2b-8xgpu-colocate.sh | 1 - .../run-mopd-qwen35-35ba3b-16xgpu-colocate.sh | 1 - .../mopd/run-mopd-qwen35-9b-8xgpu-colocate.sh | 1 - ...wen3.5_35ba3b-122ba10b-128xgpu-colocate.sh | 1 - ...en3.5_9b-35ba3b-8xgpu-2teacher-colocate.sh | 1 - ...on-opd-qwen3.5_9b-35ba3b-8xgpu-colocate.sh | 1 - relax/utils/arguments.py | 2 ++ relax/utils/opd/opd_utils.py | 12 +++++++ .../ray/test_teacher_sglang_overrides.py | 33 +++++++++++++++++++ 12 files changed, 47 insertions(+), 9 deletions(-) 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 c69a5375b..6ab08d0c8 100644 --- a/relax/utils/opd/opd_utils.py +++ b/relax/utils/opd/opd_utils.py @@ -463,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