feat: configure Ray GPU scheduling fractions

docs/en/get_started/usage.md

@@ -27,6 +27,7 @@ With the default configuration, we use these parameters to allocate `actor_num_n
 For co-located training and inference, you also need to configure:
 
   - `--colocate`: Enables co-located training and inference. By default, this makes the number of GPUs for training and inference equal. You can explicitly set a different positive `--rollout-num-gpus`, for example to use more rollout GPUs than actor GPUs; the extra GPUs are used as rollout-only resources. If `--rollout-num-gpus 0` is set explicitly, slime launches only the router and no local SGLang servers.
+  - `--ray-train-gpu-fraction` and `--ray-rollout-gpu-fraction`: Fractional Ray resource claims used to place the training and rollout actors on each GPU. These are scheduling values, not limits on CUDA utilization. In colocated mode their sum must not exceed 1.
 
 Additionally, slime supports Prefill and Decode disaggregation (PD Disaggregation). You can set the number of servers used for Prefill by setting the `--prefill-num-servers` argument.
 

docs/zh/get_started/usage.md

@@ -30,6 +30,7 @@
 当需要训推一体的时候，还需要配置上：
 
 - `--colocate`：开启训推一体。开启后默认会让训练和推理的卡数相等；也可以显式设置一个不同的正数，例如让 rollout 卡数多于 actor，多出的 GPU 会作为 rollout-only 资源使用。如果显式设置 `--rollout-num-gpus 0`，则只启动 router，不启动本地 SGLang server。
+- `--ray-train-gpu-fraction` 和 `--ray-rollout-gpu-fraction`：Ray 在每张 GPU 上放置训练与 rollout actor 时使用的分数资源声明。这些值只影响调度，并不会限制 CUDA 利用率；在训推一体模式下，两者之和不能超过 1。
 
 此外，slime 支持 Prefill 和 Decode 的分离部署 (PD Disaggregation)，可以通过设置 `--prefill-num-servers` 参数来指定用于 Prefill 的服务器数量。
 

slime/ray/placement_group.py

@@ -152,7 +152,7 @@ def allocate_train_group(
         num_nodes=num_nodes,
         num_gpus_per_node=num_gpus_per_node,
         pg=pg,
-        num_gpus_per_actor=0.4,
+        num_gpus_per_actor=args.ray_train_gpu_fraction,
         role=role,
         with_ref=with_ref,
         with_opd_teacher=with_opd_teacher,

slime/ray/rollout.py

@@ -224,7 +224,7 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
                 continue
 
             global_rank = self.rank_offset + i
-            num_gpus = 0.2
+            num_gpus = self.args.ray_rollout_gpu_fraction
             num_cpus = num_gpus
 
             # Get the base GPU ID from placement group using gpu_offset.

slime/utils/arguments.py

@@ -76,6 +76,24 @@ def add_cluster_arguments(parser):
                     "Turning this on will also set --offload to true."
                 ),
             )
+            parser.add_argument(
+                "--ray-train-gpu-fraction",
+                type=float,
+                default=0.4,
+                help=(
+                    "Fractional GPU resource claimed by each Ray training actor for placement. "
+                    "This is a scheduling claim and does not cap CUDA utilization."
+                ),
+            )
+            parser.add_argument(
+                "--ray-rollout-gpu-fraction",
+                type=float,
+                default=0.2,
+                help=(
+                    "Fractional GPU resource claimed by each Ray rollout actor for placement. "
+                    "This is a scheduling claim and does not cap CUDA utilization."
+                ),
+            )
             parser.add_argument(
                 "--offload",
                 action="store_true",
@@ -1767,6 +1785,19 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
 def slime_validate_args(args):
     args.eval_datasets = _resolve_eval_datasets(args)
 
+    for name in ("ray_train_gpu_fraction", "ray_rollout_gpu_fraction"):
+        value = getattr(args, name)
+        if not 0 < value <= 1:
+            raise ValueError(f"--{name.replace('_', '-')} must be in (0, 1]")
+    if (
+        args.colocate
+        and not args.debug_train_only
+        and not args.debug_rollout_only
+        and args.rollout_num_gpus != 0
+        and args.ray_train_gpu_fraction + args.ray_rollout_gpu_fraction > 1
+    ):
+        raise ValueError("colocated --ray-train-gpu-fraction and --ray-rollout-gpu-fraction must sum to at most 1")
+
     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.")

tests/test_megatron_argument_validation.py

@@ -227,6 +227,8 @@ def make_slime_validate_args(**overrides):
         debug_rollout_only=False,
         colocate=False,
         rollout_num_gpus=8,
+        ray_train_gpu_fraction=0.4,
+        ray_rollout_gpu_fraction=0.2,
         eval_function_path=None,
         rollout_function_path="custom.rollout",
         num_steps_per_rollout=None,
@@ -342,6 +344,37 @@ def test_slime_validate_args_preserves_zero_rollout_gpus_without_colocate(monkey
     assert args.offload_rollout is False
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("name", "value"),
+    [
+        ("ray_train_gpu_fraction", 0),
+        ("ray_train_gpu_fraction", 1.1),
+        ("ray_rollout_gpu_fraction", -0.1),
+        ("ray_rollout_gpu_fraction", 1.1),
+    ],
+)
+def test_ray_gpu_fractions_must_be_valid(monkeypatch, name, value):
+    module = load_slime_arguments_module(monkeypatch)
+    args = make_slime_validate_args(**{name: value})
+
+    with pytest.raises(ValueError, match=name.replace("_", "-")):
+        module.slime_validate_args(args)
+
+
+@pytest.mark.unit
+def test_colocated_ray_gpu_fractions_must_fit_one_gpu(monkeypatch):
+    module = load_slime_arguments_module(monkeypatch)
+    args = make_slime_validate_args(
+        colocate=True,
+        ray_train_gpu_fraction=0.6,
+        ray_rollout_gpu_fraction=0.5,
+    )
+
+    with pytest.raises(ValueError, match="must sum to at most 1"):
+        module.slime_validate_args(args)
+
+
 @pytest.mark.unit
 def test_update_weight_delta_requires_disk_transport(monkeypatch):
     module = load_slime_arguments_module(monkeypatch)

tests/test_placement_group.py

@@ -8,6 +8,7 @@
 if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
+from slime.ray import placement_group
 from slime.ray.placement_group import _create_placement_group, _get_placement_group_layout
 
 NUM_GPUS = 0
@@ -50,5 +51,21 @@ def test_create_zero_gpu_placement_group_is_empty():
     assert _create_placement_group(0) == (None, [], [])
 
 
+def test_allocate_train_group_uses_configured_ray_gpu_fraction(monkeypatch):
+    captured = {}
+
+    def fake_train_group(**kwargs):
+        captured.update(kwargs)
+        return "group"
+
+    monkeypatch.setattr(placement_group, "RayTrainGroup", fake_train_group)
+    args = Namespace(ray_train_gpu_fraction=0.55)
+
+    result = placement_group.allocate_train_group(args, 1, 1, pg="placement")
+
+    assert result == "group"
+    assert captured["num_gpus_per_actor"] == 0.55
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__]))