support offloading non-updatable server

.github/workflows/pr-test.yml

@@ -52,7 +52,69 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_short.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config_distributed.py"}]
+        info: [{"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}]
+    defaults:
+      run:
+        working-directory: ${{ github.workspace }}
+    env:
+      GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }}
+      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+      SLIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }}
+      SLIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }}
+      SLIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }}
+      SLIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }}
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+
+      - name: Install
+        shell: bash
+        run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages
+
+
+      - name: Execute
+        shell: bash
+        run: |
+          TEST_PATH="${{ matrix.info.test_file }}"
+          if [[ "$TEST_PATH" != tests/* ]]; then
+            TEST_PATH="tests/$TEST_PATH"
+          fi
+          if [ "${{ matrix.info.num_gpus }}" = "0" ]; then
+            python "$TEST_PATH"
+          else
+            python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH"
+          fi
+
+  e2e-test-sglang-config:
+
+    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-sglang-config'))
+
+
+    runs-on: self-hosted
+    container:
+      image: slimerl/slime:latest
+      options: >
+        --gpus all
+        --ipc=host
+        --shm-size=16g
+        --ulimit memlock=-1
+        --ulimit stack=67108864
+        --memory=0
+        --memory-swap=0
+        -e http_proxy=$http_proxy
+        -e https_proxy=$https_proxy
+        -e HTTP_PROXY=$HTTP_PROXY
+        -e HTTPS_PROXY=$HTTPS_PROXY
+        -v /mnt/nvme0n1/slime_ci:/data/slime_ci
+        -v /mnt/nvme0n1/slime_ci/models:/root/models
+        -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets
+
+    strategy:
+      fail-fast: false
+      matrix:
+        info: [{"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config_distributed.py"}, {"num_gpus": 8, "test_file": "test_sglang_config_mixed_offload.py"}, {"num_gpus": 8, "test_file": "test_sglang_config_mixed_offload_ft.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}

.github/workflows/pr-test.yml.j2

@@ -4,8 +4,16 @@
       'tests': [
         {'test_file': 'test_qwen2.5_0.5B_gsm8k_async_short.py', 'num_gpus': 4},
         {'test_file': 'test_qwen2.5_0.5B_gsm8k_short.py', 'num_gpus': 4},
+        {'test_file': 'test_qwen2.5_0.5B_ppo_critic_only_short.py', 'num_gpus': 4},
+      ],
+    },
+    'e2e-test-sglang-config': {
+      'label': 'run-ci-sglang-config',
+      'tests': [
         {'test_file': 'test_qwen2.5_0.5B_sglang_config.py', 'num_gpus': 8},
         {'test_file': 'test_qwen2.5_0.5B_sglang_config_distributed.py', 'num_gpus': 8},
+        {'test_file': 'test_sglang_config_mixed_offload.py', 'num_gpus': 8},
+        {'test_file': 'test_sglang_config_mixed_offload_ft.py', 'num_gpus': 8},
       ],
     },
     'e2e-test-megatron': {

slime/backends/megatron_utils/model.py

@@ -655,7 +655,7 @@ def train(
                     # TODO: figure out why KL is not exactly zero when using PPO loss with KL clipping, and whether this is expected behavior or a bug.
                     assert log_dict["train/ppo_kl"] < 1e-8, f"{log_dict=}"
                 if accumulated_step_id == 0 and "train/kl_loss" in log_dict:
-                    assert log_dict["train/kl_loss"] == 0.0, f"{log_dict=}"
+                    assert log_dict["train/kl_loss"] < 1e-8, f"{log_dict=}"
 
             logger.info(f"{role_tag}step {accumulated_step_id}: {log_dict}")
 

slime/backends/sglang_utils/sglang_engine.py

@@ -90,21 +90,6 @@ def _wait_server_healthy(base_url, api_key, is_process_alive):
 
             time.sleep(2)
 
-        # use flush_cache to make sure the working queue is empty, so that we can do offload
-        while True:
-            try:
-                response = session.get(f"{base_url}/flush_cache", headers=headers)
-                if response.status_code == 200:
-                    break
-
-            except requests.RequestException:
-                pass
-
-            if not is_process_alive():
-                raise Exception("Server process terminated unexpectedly.")
-
-            time.sleep(2)
-
 
 class SGLangEngine(RayActor):
     def __init__(
@@ -370,6 +355,17 @@ def resume_memory_occupation(self, tags: list[str] = None):
     def check_weights(self, action: str):
         return self._make_request("weights_checker", {"action": action})
 
+    def update_weights_from_disk(self, model_path: str, load_format: str | None = None):
+        """Reload weights from *model_path* without restarting the engine.
+
+        Used for non-updatable (frozen) models that overlap with megatron:
+        after offload, weights are restored from disk instead of CPU cache.
+        """
+        payload = {"model_path": model_path}
+        if load_format is not None:
+            payload["load_format"] = load_format
+        return self._make_request("update_weights_from_disk", payload)
+
     def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend):
         return self._make_request(
             "init_weights_update_group",

slime/ray/rollout.py

@@ -53,6 +53,8 @@ class ServerGroup:
     rank_offset: int = 0  # cumulative engine count before this group
     gpu_offset: int = 0  # cumulative GPU count before this group
     sglang_overrides: dict = dataclasses.field(default_factory=dict)
+    needs_offload: bool = False  # True when this group's GPUs overlap with megatron
+    model_path: str | None = None  # checkpoint path for update_weights_from_disk
     router_ip: str | None = None
     router_port: int | None = None
 
@@ -174,17 +176,35 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
     def offload(self):
         """Fire release_memory_occupation on all engines (non-blocking).
 
-        Returns a list of Ray ObjectRefs.
+        Returns a list of Ray ObjectRefs.  Skipped for groups that do not
+        overlap with megatron GPUs (``needs_offload=False``).
         """
+        if not self.needs_offload:
+            return []
         return [engine.release_memory_occupation.remote() for engine in self.engines if engine is not None]
 
     def onload(self, tags: list[str] | None = None):
         """Fire resume_memory_occupation on all engines (non-blocking).
 
-        Returns a list of Ray ObjectRefs.
+        Returns a list of Ray ObjectRefs.  Skipped for groups that do not
+        overlap with megatron GPUs (``needs_offload=False``).
         """
+        if not self.needs_offload:
+            return []
         return [engine.resume_memory_occupation.remote(tags=tags) for engine in self.engines if engine is not None]
 
+    def onload_weights_from_disk(self):
+        """Reload weights from ``model_path`` for non-updatable groups.
+
+        Used instead of ``resume_memory_occupation(tags=[WEIGHTS])`` so that
+        CPU memory is not consumed by offloaded weight copies.
+        """
+        if not self.needs_offload or not self.model_path:
+            return []
+        return [
+            engine.update_weights_from_disk.remote(self.model_path) for engine in self.engines if engine is not None
+        ]
+
 
 @dataclasses.dataclass
 class RolloutServer:
@@ -261,20 +281,32 @@ def recover(self):
 
         # Post-recovery: offload then onload weights for newly created engines.
         release_handles = []
-        new_engines_all = []
-        for g, dead_indices in zip(self.server_groups, dead_per_group, strict=True):
+        updatable_new_engines = []
+        non_updatable_groups_engines: list[tuple[str, list]] = []
+        for g, dead_indices in zip(self.engine_groups, dead_per_group, strict=True):
             logger.info(f"Recovered {g.num_new_engines} dead rollout engines (worker_type={g.worker_type})")
             assert g.num_new_engines == len(dead_indices), "num_new_engines does not match dead_indices length"
-            if g.args.offload_rollout and dead_indices:
+            if g.needs_offload and dead_indices:
                 new_engines = [g.all_engines[i] for i in dead_indices]
                 release_handles.extend(engine.release_memory_occupation.remote() for engine in new_engines)
-                new_engines_all.extend(new_engines)
+                if self.update_weights:
+                    updatable_new_engines.extend(new_engines)
+                elif g.model_path:
+                    non_updatable_groups_engines.append((g.model_path, new_engines))
 
         if release_handles:
             ray.get(release_handles)
-            ray.get(
-                [engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) for engine in new_engines_all]
-            )
+            # Resume GPU memory for all engines that need offload.
+            all_resume_engines = updatable_new_engines[:]
+            for _model_path, engines in non_updatable_groups_engines:
+                all_resume_engines.extend(engines)
+            if all_resume_engines:
+                ray.get(
+                    [
+                        engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS])
+                        for engine in all_resume_engines
+                    ]
+                )
 
     def offload(self):
         """Release memory occupation across all groups (concurrent)."""
@@ -290,6 +322,28 @@ def onload(self, tags: list[str] | None = None):
             handles.extend(g.onload(tags))
         return ray.get(handles) if handles else []
 
+    def onload_weights(self):
+        """Restore weights for offloaded groups.
+
+        All groups resume from CPU cache via ``resume_memory_occupation``.
+        For updatable servers, weights will be overwritten by
+        ``update_weights`` shortly after.  For non-updatable servers the
+        CPU backup already contains the correct (unchanged) weights.
+        """
+        handles = []
+        for g in self.engine_groups:
+            if not g.needs_offload:
+                continue
+            handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
+        return ray.get(handles) if handles else []
+
+    def onload_kv(self):
+        """Resume KV cache and CUDA graphs for offloaded groups."""
+        handles = []
+        for g in self.engine_groups:
+            handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH]))
+        return ray.get(handles) if handles else []
+
 
 @ray.remote
 class RolloutManager:
@@ -444,10 +498,12 @@ def onload(self, tags: list[str] | None = None):
             srv.onload(tags)
 
     def onload_weights(self):
-        self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS])
+        for srv in self.servers.values():
+            srv.onload_weights()
 
     def onload_kv(self):
-        self.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH])
+        for srv in self.servers.values():
+            srv.onload_kv()
 
     def recover_updatable_engines(self):
         """Restart any dead rollout engines and update num_new_engines for update_weights detection.
@@ -887,6 +943,30 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool
     return router_ip, router_port
 
 
+def _compute_rollout_offset(args) -> int:
+    """Offset (in PG bundle slots) where rollout GPUs start."""
+    if args.debug_train_only or args.debug_rollout_only or args.colocate:
+        return 0
+    if args.critic_train_only:
+        return args.critic_num_nodes * args.critic_num_gpus_per_node
+    offset = args.actor_num_nodes * args.actor_num_gpus_per_node
+    if args.use_critic:
+        offset += args.critic_num_nodes * args.critic_num_gpus_per_node
+    return offset
+
+
+def _compute_megatron_num_gpus(args) -> int:
+    """Total number of megatron (actor + critic) GPU slots in the placement group."""
+    if args.debug_rollout_only:
+        return 0
+    if args.critic_train_only:
+        return args.critic_num_nodes * args.critic_num_gpus_per_node
+    num = args.actor_num_nodes * args.actor_num_gpus_per_node
+    if args.use_critic:
+        num += args.critic_num_nodes * args.critic_num_gpus_per_node
+    return num
+
+
 def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
     """Start rollout servers: one per model, each with its own router.
 
@@ -906,6 +986,10 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
     gpu_offset = 0
     engine_offset = 0
 
+    # Compute megatron GPU range for per-group offload decisions.
+    rollout_pg_offset = _compute_rollout_offset(args)
+    megatron_num_gpus = _compute_megatron_num_gpus(args)
+
     for model_idx, model_cfg in enumerate(config.models):
         model_cfg.resolve(args)
 
@@ -926,6 +1010,17 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
             num_gpu_per_engine_local = min(gpus_per_engine, args.num_gpus_per_node)
             num_engines = group_cfg.num_gpus // num_gpu_per_engine_local
 
+            # Only offload groups whose GPUs overlap with megatron.
+            group_abs_start = rollout_pg_offset + gpu_offset
+            needs_offload = args.offload_rollout and group_abs_start < megatron_num_gpus
+            overrides = dict(group_cfg.overrides)
+            if args.offload_rollout and not needs_offload:
+                overrides.setdefault("enable_memory_saver", False)
+            logger.info(
+                f"Engine group '{group_cfg.worker_type}' gpu_offset={gpu_offset} "
+                f"(abs={group_abs_start}): needs_offload={needs_offload}"
+            )
+
             group = ServerGroup(
                 args=args,
                 pg=pg,
@@ -935,7 +1030,9 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
                 worker_type=group_cfg.worker_type,
                 rank_offset=engine_offset,
                 gpu_offset=gpu_offset,
-                sglang_overrides=group_cfg.overrides,
+                sglang_overrides=overrides,
+                needs_offload=needs_offload,
+                model_path=overrides.get("model_path", args.hf_checkpoint),
                 router_ip=router_ip,
                 router_port=router_port,
             )

slime/utils/arguments.py

@@ -1521,6 +1521,15 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
 def slime_validate_args(args):
     args.eval_datasets = _resolve_eval_datasets(args)
 
+    if args.critic_train_only:
+        if not args.use_critic:
+            raise ValueError("--critic-train-only requires --use-critic (or --advantage-estimator ppo).")
+        if args.actor_num_nodes != 0 or args.actor_num_gpus_per_node != 0:
+            raise ValueError(
+                "--critic-train-only requires --actor-num-nodes 0 --actor-num-gpus-per-node 0, "
+                f"but got actor_num_nodes={args.actor_num_nodes}, actor_num_gpus_per_node={args.actor_num_gpus_per_node}."
+            )
+
     if args.kl_coef != 0 or args.use_kl_loss:
         if not os.path.exists(args.ref_load):
             raise FileNotFoundError(f"ref_load {args.ref_load} does not exist, please check the path.")