Add GPU placement validation before starting rollout engines

slime/ray/rollout.py

@@ -27,6 +27,7 @@
 from slime.utils.types import Sample
 
 from ..utils.metric_utils import has_repetition
+from .rollout_validation import validate_server_group_gpu_indices
 from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock
 
 logging.getLogger("httpx").setLevel(logging.WARNING)
@@ -87,6 +88,16 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
         num_gpu_per_engine = min(self.num_gpus_per_engine, self.args.num_gpus_per_node)
 
         pg, reordered_bundle_indices, reordered_gpu_ids = self.pg
+        validate_server_group_gpu_indices(
+            worker_type=self.worker_type,
+            gpu_offset=self.gpu_offset,
+            num_gpus_per_engine=self.num_gpus_per_engine,
+            num_gpu_per_engine=num_gpu_per_engine,
+            num_engines=len(self.all_engines),
+            num_available_gpus=len(reordered_gpu_ids),
+            rollout_num_gpus=self.args.rollout_num_gpus,
+            rollout_num_gpus_per_engine=self.args.rollout_num_gpus_per_engine,
+        )
 
         RolloutRayActor = ray.remote(SGLangEngine)
 

slime/ray/rollout_validation.py

@@ -0,0 +1,32 @@
+def validate_server_group_gpu_indices(
+    *,
+    worker_type: str,
+    gpu_offset: int,
+    num_gpus_per_engine: int,
+    num_gpu_per_engine: int,
+    num_engines: int,
+    num_available_gpus: int,
+    rollout_num_gpus: int,
+    rollout_num_gpus_per_engine: int,
+) -> None:
+    if num_engines == 0:
+        return
+
+    required_gpu_slots = gpu_offset + num_engines * num_gpu_per_engine
+    if gpu_offset >= 0 and num_gpu_per_engine > 0 and required_gpu_slots <= num_available_gpus:
+        return
+
+    raise ValueError(
+        "Invalid rollout server group GPU placement: "
+        f"worker_type={worker_type}, "
+        f"gpu_offset={gpu_offset}, "
+        f"num_gpus_per_engine={num_gpus_per_engine}, "
+        f"num_gpu_per_engine_on_node={num_gpu_per_engine}, "
+        f"num_engines={num_engines}, "
+        f"required_gpu_slots={required_gpu_slots}, "
+        f"len(reordered_gpu_ids)={num_available_gpus}, "
+        f"rollout_num_gpus={rollout_num_gpus}, "
+        f"rollout_num_gpus_per_engine={rollout_num_gpus_per_engine}. "
+        "Please align --rollout-num-gpus, --rollout-num-gpus-per-engine, "
+        "and --sglang-config server_groups."
+    )

tests/test_rollout_validation.py

@@ -0,0 +1,56 @@
+import pytest
+
+from slime.ray.rollout_validation import validate_server_group_gpu_indices
+
+
+@pytest.mark.unit
+def test_validate_server_group_gpu_indices_accepts_valid_config():
+    validate_server_group_gpu_indices(
+        worker_type="regular",
+        gpu_offset=2,
+        num_gpus_per_engine=1,
+        num_gpu_per_engine=1,
+        num_engines=2,
+        num_available_gpus=4,
+        rollout_num_gpus=4,
+        rollout_num_gpus_per_engine=1,
+    )
+
+
+@pytest.mark.unit
+def test_validate_server_group_gpu_indices_allows_empty_group():
+    validate_server_group_gpu_indices(
+        worker_type="placeholder",
+        gpu_offset=4,
+        num_gpus_per_engine=1,
+        num_gpu_per_engine=1,
+        num_engines=0,
+        num_available_gpus=4,
+        rollout_num_gpus=4,
+        rollout_num_gpus_per_engine=1,
+    )
+
+
+@pytest.mark.unit
+def test_validate_server_group_gpu_indices_reports_config_context():
+    with pytest.raises(ValueError) as exc_info:
+        validate_server_group_gpu_indices(
+            worker_type="regular",
+            gpu_offset=3,
+            num_gpus_per_engine=2,
+            num_gpu_per_engine=2,
+            num_engines=1,
+            num_available_gpus=4,
+            rollout_num_gpus=4,
+            rollout_num_gpus_per_engine=2,
+        )
+
+    message = str(exc_info.value)
+    assert "worker_type=regular" in message
+    assert "gpu_offset=3" in message
+    assert "num_gpus_per_engine=2" in message
+    assert "num_engines=1" in message
+    assert "required_gpu_slots=5" in message
+    assert "len(reordered_gpu_ids)=4" in message
+    assert "rollout_num_gpus=4" in message
+    assert "rollout_num_gpus_per_engine=2" in message