diff --git a/tests/experimental/agent_loop/test_sglang_sampling_seed_on_cpu.py b/tests/experimental/agent_loop/test_sglang_sampling_seed_on_cpu.py new file mode 100644 index 00000000000..d702ed71424 --- /dev/null +++ b/tests/experimental/agent_loop/test_sglang_sampling_seed_on_cpu.py @@ -0,0 +1,149 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ast +import asyncio +from pathlib import Path + +import pytest + +AGENT_LOOP_PATH = Path(__file__).parents[3] / "verl" / "experimental" / "agent_loop" / "agent_loop.py" +SGLANG_SERVER_PATH = ( + Path(__file__).parents[3] / "verl" / "workers" / "rollout" / "sglang_rollout" / "async_sglang_server.py" +) + + +def _load_agent_loop_functions(*names): + module = ast.parse(AGENT_LOOP_PATH.read_text()) + selected_nodes = [] + for node in module.body: + if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef) and node.name in names: + selected_nodes.append(node) + + found_names = {node.name for node in selected_nodes} + missing_names = set(names) - found_names + if missing_names: + raise AssertionError(f"agent_loop helper missing: {sorted(missing_names)}") + + ns = {"Any": object, "hashlib": __import__("hashlib")} + exec(compile(ast.Module(body=selected_nodes, type_ignores=[]), str(AGENT_LOOP_PATH), "exec"), ns) + return tuple(ns[name] for name in names) + + +def _load_sglang_server_functions(*names): + module = ast.parse(SGLANG_SERVER_PATH.read_text()) + selected_nodes = [] + for node in module.body: + if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef) and node.name in names: + selected_nodes.append(node) + + found_names = {node.name for node in selected_nodes} + missing_names = set(names) - found_names + if missing_names: + raise AssertionError(f"SGLang server helper missing: {sorted(missing_names)}") + + ns = {"Any": object} + exec(compile(ast.Module(body=selected_nodes, type_ignores=[]), str(SGLANG_SERVER_PATH), "exec"), ns) + return tuple(ns[name] for name in names) + + +def test_stable_sglang_sampling_seed_is_positive_int31_on_cpu(): + (stable_seed,) = _load_agent_loop_functions("_stable_sglang_sampling_seed") + + seed = stable_seed(sample_index=0, rollout_n=0, step=0, base_seed=42) + + assert 0 < seed <= 0x7FFFFFFF + assert seed == stable_seed(sample_index=0, rollout_n=0, step=0, base_seed=42) + assert seed != stable_seed(sample_index=0, rollout_n=1, step=0, base_seed=42) + assert seed != stable_seed(sample_index=0, rollout_n=0, step=1, base_seed=42) + + +def test_sglang_sampling_seed_base_defaults_for_deterministic_inference(): + (seed_base,) = _load_agent_loop_functions("_get_sglang_sampling_seed_base") + + assert seed_base({"enable_deterministic_inference": True}) == 42 + assert seed_base({"enable_deterministic_inference": True, "random_seed": 7}) == 7 + assert seed_base({"random_seed": 7}) is None + assert seed_base({}) is None + + +def test_global_trajectory_chunks_avoid_duplicate_sglang_seeds_across_workers_on_cpu(): + get_trajectory_info, split_trajectory_info, stable_seed = _load_agent_loop_functions( + "get_trajectory_info", + "_split_trajectory_info_for_chunks", + "_stable_sglang_sampling_seed", + ) + + repeated_prompt_indices = [1, 1, 1, 1] + chunk_sizes = [2, 2] + + buggy_chunks = [ + asyncio.run(get_trajectory_info(step=10, index=repeated_prompt_indices[:2], validate=False)), + asyncio.run(get_trajectory_info(step=10, index=repeated_prompt_indices[2:], validate=False)), + ] + assert [[item["rollout_n"] for item in chunk] for chunk in buggy_chunks] == [[0, 1], [0, 1]] + buggy_seeds = [ + stable_seed(sample_index=item["sample_index"], rollout_n=item["rollout_n"], step=item["step"], base_seed=42) + for chunk in buggy_chunks + for item in chunk + ] + assert len(set(buggy_seeds)) < len(buggy_seeds) + + trajectory_info = asyncio.run(get_trajectory_info(step=10, index=repeated_prompt_indices, validate=False)) + chunks = split_trajectory_info(trajectory_info, chunk_sizes=chunk_sizes) + + assert [[item["rollout_n"] for item in chunk] for chunk in chunks] == [[0, 1], [2, 3]] + + seeds = [ + stable_seed(sample_index=item["sample_index"], rollout_n=item["rollout_n"], step=item["step"], base_seed=42) + for chunk in chunks + for item in chunk + ] + assert len(set(seeds)) == len(seeds) + + +def test_sglang_server_normalizes_custom_engine_kwargs_on_cpu(): + (normalize,) = _load_sglang_server_functions("_normalize_sglang_engine_kwargs") + + assert normalize({}) is False + assert normalize({"enable_deterministic_inference": True}) is True + + fsdp_kwargs = {"rl_on_policy_target": "fsdp"} + assert normalize(fsdp_kwargs) is True + assert fsdp_kwargs == {} + + with pytest.raises(ValueError, match="rl_on_policy_target"): + normalize({"rl_on_policy_target": True}) + + +def test_sglang_server_copies_engine_kwargs_before_mutation_on_cpu(): + module = ast.parse(SGLANG_SERVER_PATH.read_text()) + server_class = next( + node for node in module.body if isinstance(node, ast.ClassDef) and node.name == "SGLangHttpServer" + ) + launch_fn = next( + node for node in server_class.body if isinstance(node, ast.AsyncFunctionDef) and node.name == "launch_server" + ) + + engine_kwargs_assignment = next( + node + for node in launch_fn.body + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "engine_kwargs" for target in node.targets) + ) + + assert ( + ast.unparse(engine_kwargs_assignment.value) + == "dict((self.config.get('engine_kwargs', {}) or {}).get('sglang', {}) or {})" + ) diff --git a/verl/experimental/agent_loop/agent_loop.py b/verl/experimental/agent_loop/agent_loop.py index 9b618f89d2f..9f8c33fbac4 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/verl/experimental/agent_loop/agent_loop.py @@ -28,6 +28,7 @@ """ import asyncio +import hashlib import logging import os import random @@ -78,6 +79,26 @@ DEFAULT_ROUTING_CACHE_SIZE = 10000 +def _stable_sglang_sampling_seed(sample_index: Any, rollout_n: int, step: Any, base_seed: Any) -> int: + key = f"{step}|{sample_index}|{rollout_n}|{base_seed}" + seed = int(hashlib.sha256(key.encode()).hexdigest()[:8], 16) & 0x7FFFFFFF + return seed or 1 + + +def _get_sglang_sampling_seed_base(sglang_kwargs: dict[str, Any]) -> Any: + base_seed = sglang_kwargs.get("random_seed") + rl_on_policy_target = sglang_kwargs.get("rl_on_policy_target", None) + if rl_on_policy_target not in (None, False, "fsdp"): + raise ValueError("SGLang rl_on_policy_target must be None, False, or 'fsdp'.") + + deterministic_inference = sglang_kwargs.get("enable_deterministic_inference", False) or ( + rl_on_policy_target == "fsdp" + ) + if deterministic_inference: + return 42 if base_seed is None else base_seed + return None + + class AgentLoopMetrics(BaseModel): """Agent loop performance metrics.""" @@ -558,7 +579,9 @@ def _get_mm_processor_kwargs(self, audio_data: Optional[list[Any]] = None) -> di mm_processor_kwargs["sampling_rate"] = int(sampling_rate) return mm_processor_kwargs - async def generate_sequences(self, batch: DataProto) -> DataProto: + async def generate_sequences( + self, batch: DataProto, trajectory_info: list[dict[str, Any]] | None = None + ) -> DataProto: """Generate sequences from agent loop. Args: @@ -626,9 +649,12 @@ def apply_greedy_sampling_params(params: dict[str, Any]) -> None: else: traced_indices = set(range(len(batch))) - trajectory_info = await get_trajectory_info( - batch.meta_info.get("global_steps", -1), index.tolist(), batch.meta_info.get("validate", False) - ) + if trajectory_info is None: + trajectory_info = await get_trajectory_info( + batch.meta_info.get("global_steps", -1), index.tolist(), batch.meta_info.get("validate", False) + ) + else: + assert len(trajectory_info) == len(batch), "trajectory_info length must match batch size" # NOTE: __do_sample__ is an internal per-sample override used by REMAX combined rollout. # Do not forward it to concrete agent loops, which may reject unknown kwargs. @@ -661,6 +687,16 @@ async def _run_agent_loop( trace: bool = True, **kwargs, ) -> _InternalAgentLoopOutput: + seed_base = _get_sglang_sampling_seed_base((self.rollout_config.engine_kwargs or {}).get("sglang", {}) or {}) + if seed_base is not None and "sampling_seed" not in sampling_params: + sampling_params = dict(sampling_params) + sampling_params["sampling_seed"] = _stable_sglang_sampling_seed( + sample_index=trajectory.get("sample_index", ""), + rollout_n=trajectory.get("rollout_n", 0), + step=trajectory.get("step", 0), + base_seed=seed_base, + ) + with rollout_trace_attr( step=trajectory["step"], sample_index=trajectory["sample_index"], @@ -1138,6 +1174,19 @@ async def get_trajectory_info(step, index, validate): return trajectory_info +def _split_trajectory_info_for_chunks( + trajectory_info: list[dict[str, Any]], chunk_sizes: list[int] +) -> list[list[dict[str, Any]]]: + chunks = [] + offset = 0 + for chunk_size in chunk_sizes: + next_offset = offset + chunk_size + chunks.append(trajectory_info[offset:next_offset]) + offset = next_offset + assert offset == len(trajectory_info), "chunk sizes must cover trajectory_info" + return chunks + + class AgentLoopManager: """Agent loop manager that manages a group of agent loop workers. @@ -1212,11 +1261,22 @@ async def generate_sequences(self, prompts: DataProto) -> DataProto: if "priority" not in prompts.non_tensor_batch: prompts.non_tensor_batch["priority"] = np.arange(len(prompts), dtype=np.int64) + if "index" in prompts.non_tensor_batch: + index = prompts.non_tensor_batch["index"] + else: + index = np.arange(len(prompts)) + full_trajectory_info = await get_trajectory_info( + prompts.meta_info.get("global_steps", -1), + index.tolist(), + prompts.meta_info.get("validate", False), + ) + chunkes = prompts.chunk(len(self.agent_loop_workers)) + traj_chunks = _split_trajectory_info_for_chunks(full_trajectory_info, [len(chunk) for chunk in chunkes]) outputs = await asyncio.gather( *[ - worker.generate_sequences.remote(chunk) - for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True) + worker.generate_sequences.remote(chunk, traj_chunk) + for worker, chunk, traj_chunk in zip(self.agent_loop_workers, chunkes, traj_chunks, strict=True) ] ) output = DataProto.concat(outputs) diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index 171a9b3f2e7..eb93db4fa55 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -59,6 +59,19 @@ visible_devices_keyword = get_visible_devices_keyword() +def _normalize_sglang_engine_kwargs(engine_kwargs: dict[str, Any]) -> bool: + deterministic_inference = bool(engine_kwargs.pop("enable_deterministic_inference", False)) + + rl_on_policy_target = engine_kwargs.get("rl_on_policy_target", None) + if rl_on_policy_target is None or rl_on_policy_target is False: + engine_kwargs.pop("rl_on_policy_target", None) + return deterministic_inference + if rl_on_policy_target == "fsdp": + engine_kwargs.pop("rl_on_policy_target", None) + return True + raise ValueError("SGLang rl_on_policy_target must be None, False, or 'fsdp'.") + + def _extract_prompt_logprobs_sglang( meta_info: dict, num_prompt_logprobs: int, @@ -253,9 +266,15 @@ async def launch_server(self, master_address: str = None, master_port: int = Non self._master_address = master_address self._master_port = master_port - engine_kwargs = self.config.get("engine_kwargs", {}).get("sglang", {}) or {} + engine_kwargs = dict((self.config.get("engine_kwargs", {}) or {}).get("sglang", {}) or {}) attention_backend = engine_kwargs.pop("attention_backend", None) mm_attention_backend = engine_kwargs.pop("mm_attention_backend", None) + deterministic_sampling = _normalize_sglang_engine_kwargs(engine_kwargs) + if deterministic_sampling: + sampling_backend = engine_kwargs.get("sampling_backend") or "pytorch" + if sampling_backend != "pytorch": + raise ValueError("SGLang deterministic sampling_seed requires sampling_backend='pytorch'.") + engine_kwargs["sampling_backend"] = "pytorch" if attention_backend is None: if torch.version.hip is not None: attention_backend = "aiter"