diff --git a/docs/en/get_started/usage.md b/docs/en/get_started/usage.md index 2ecebaba65..308fe62c79 100644 --- a/docs/en/get_started/usage.md +++ b/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. diff --git a/docs/zh/get_started/usage.md b/docs/zh/get_started/usage.md index 533729c40a..66b2c6b51a 100644 --- a/docs/zh/get_started/usage.md +++ b/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 的服务器数量。 diff --git a/slime/ray/placement_group.py b/slime/ray/placement_group.py index eb652c3af9..864f9b108c 100644 --- a/slime/ray/placement_group.py +++ b/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, diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index abd8e35d70..1bd6a43281 100644 --- a/slime/ray/rollout.py +++ b/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. diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 2ef7dad59b..a1c3e9d34e 100644 --- a/slime/utils/arguments.py +++ b/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.") diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index ba3304d79c..dbbf2febac 100644 --- a/tests/test_megatron_argument_validation.py +++ b/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) diff --git a/tests/test_placement_group.py b/tests/test_placement_group.py index 537b883d7a..e1545c8487 100644 --- a/tests/test_placement_group.py +++ b/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__]))