diff --git a/pyproject.toml b/pyproject.toml index 8adfbd9fe2..707646c856 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,9 @@ where = [""] # tightening the first pattern to "torchtitan.*" cannot silently drop it. include = ["torchtitan*", "torchtitan_recipes*"] +[tool.setuptools.package-data] +"torchtitan.experiments.rl.examples.verifiers" = ["verifiers_env.toml"] + [tool.pytest.ini_options] addopts = ["--showlocals"] # show local variables in tracebacks testpaths = ["tests"] diff --git a/torchtitan/experiments/__init__.py b/torchtitan/experiments/__init__.py index 7066523fb9..2149234269 100644 --- a/torchtitan/experiments/__init__.py +++ b/torchtitan/experiments/__init__.py @@ -20,5 +20,6 @@ "alphabet_sort", "dapo_math", "search_r1", + "verifiers", ] ) diff --git a/torchtitan/experiments/rl/README.md b/torchtitan/experiments/rl/README.md index 5df4088888..0419f2b524 100644 --- a/torchtitan/experiments/rl/README.md +++ b/torchtitan/experiments/rl/README.md @@ -10,7 +10,7 @@ Together, the unified model, batch-invariant mode, and single training stack pro Note: Unified-model performance varies by model, input shape, and parallelism: it can trail native vLLM in inference-only workloads but outperform it end to end in some RL configurations. Batch invariance trades throughput for exact numerics and can be used for debugging or controlled on-policy studies. -[Architecture](#architecture) · [Write an experiment](#write-an-experiment) · [DAPO Math](./examples/dapo_math) · [Observability](#observability) · [Quick Start](#quick-start) +[Architecture](#architecture) · [Write an experiment](#write-an-experiment) · [DAPO Math](./examples/dapo_math) · [Verifiers](./examples/verifiers) · [Observability](#observability) · [Quick Start](#quick-start) > **Note:** TitanRL is under active development. APIs and configurations may change. @@ -94,6 +94,12 @@ Train on verifiable math with DAPO loss and Math-Verify rewards. [Run DAPO Math](./examples/dapo_math) +### Verifiers: optional integration example + +Run the DAPO Math workload with Verifiers managing the local rollout environment. + +[Run Verifiers](./examples/verifiers) + ### Search-R1: multi-turn tool use Train a model to issue search queries, consume tool responses, and answer with an exact-match reward. diff --git a/torchtitan/experiments/rl/examples/verifiers/README.md b/torchtitan/experiments/rl/examples/verifiers/README.md new file mode 100644 index 0000000000..f32026a8bd --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/README.md @@ -0,0 +1,31 @@ +# Verifiers + +This example keeps the existing [DAPO Math](../dapo_math) recipe unchanged and replaces only its rollout path with [Verifiers](https://github.com/PrimeIntellect-ai/verifiers). Training still uses the filtered DAPO-Math dataset, AIME 2025 validation, DAPO loss, and the Qwen3-4B-Base model. + +Verifiers runs a single-turn math task with its `null` harness. The runtime is a local subprocess; there is no Docker or remote sandbox and no tools are exposed. Do not use this configuration for untrusted code execution. + +Verifiers is optional, and all Verifiers integration code lives in this example. +Other TitanRL recipes do not require it. + +## Setup + +Follow the [TitanRL setup](../../README.md), then install this example's dependencies: + +```bash +pip install -r torchtitan/experiments/rl/examples/verifiers/requirements.txt + +python scripts/download_hf_assets.py \ + --repo_id Qwen/Qwen3-4B-Base \ + --local_dir torchtitan/experiments/rl/example_checkpoint \ + --all +``` + +## Run + +```bash +python -m torchtitan.experiments.rl.train \ + --module verifiers \ + --config rl_dapo_qwen3_4b_verifiers_8k +``` + +Use `rl_dapo_qwen3_4b_verifiers_32k` for the 32K response variant. diff --git a/torchtitan/experiments/rl/examples/verifiers/__init__.py b/torchtitan/experiments/rl/examples/verifiers/__init__.py new file mode 100644 index 0000000000..a797feeedf --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from torchtitan.experiments.rl.examples.verifiers.rollouter import ( + VerifiersMathRollouter, +) + +__all__ = ["VerifiersMathRollouter"] diff --git a/torchtitan/experiments/rl/examples/verifiers/components/__init__.py b/torchtitan/experiments/rl/examples/verifiers/components/__init__.py new file mode 100644 index 0000000000..e45b4ea5d3 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/components/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from torchtitan.experiments.rl.examples.verifiers.components.dataset import ( + VerifiersTaskDataset, + VerifiersTaskSample, +) +from torchtitan.experiments.rl.examples.verifiers.components.env_server import ( + VerifiersEnvServer, +) +from torchtitan.experiments.rl.examples.verifiers.components.rollouter import ( + VerifiersRewardFn, + VerifiersRollouter, +) + +__all__ = [ + "VerifiersEnvServer", + "VerifiersRewardFn", + "VerifiersRollouter", + "VerifiersTaskDataset", + "VerifiersTaskSample", +] diff --git a/torchtitan/experiments/rl/examples/verifiers/components/dataset.py b/torchtitan/experiments/rl/examples/verifiers/components/dataset.py new file mode 100644 index 0000000000..2d03d023b3 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/components/dataset.py @@ -0,0 +1,121 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import importlib +import random +import sys +from collections.abc import Iterator +from dataclasses import dataclass, field +from typing import Any + +from torchtitan.config import Configurable + + +def register_local_taskset_alias(taskset_id: str) -> str: + """Register a dotted local taskset under an ID Verifiers 0.3.0 can import.""" + if "." not in taskset_id or "/" in taskset_id: + return taskset_id + + module = importlib.import_module(taskset_id) + alias = taskset_id.replace(".", "_").lower() + existing = sys.modules.get(alias) + if existing is not None and existing is not module: + raise ValueError(f"taskset alias {alias!r} is already registered") + sys.modules[alias] = module + return alias + + +@dataclass(frozen=True, kw_only=True, slots=True) +class VerifiersTaskSample: + """Serialized task data dispatched to a stateless Verifiers EnvServer.""" + + task_data: dict[str, Any] + + +class VerifiersTaskDataset(Configurable): + """Load a Verifiers taskset into TorchTitan's resumable dataset contract.""" + + @dataclass(kw_only=True, slots=True) + class Config(Configurable.Config): + # Importable Verifiers taskset plugin ID or local dotted module path. + taskset_id: str + # Keyword arguments used to construct the taskset's config. + taskset_args: dict[str, Any] = field(default_factory=dict) + # Optional task cap; required when the taskset is infinite. + num_tasks: int | None = None + # Seed used to produce a reproducible task order. + seed: int = 42 + # Whether to reshuffle the task order at initialization and each epoch. + shuffle: bool = True + + def __post_init__(self) -> None: + if self.num_tasks is not None and self.num_tasks <= 0: + raise ValueError("num_tasks must be positive") + + def __init__(self, config: Config) -> None: + # Verifiers is an example-only dependency. Keep the import local so the + # rest of TorchTitan RL does not require it. + from verifiers.v1.utils.loaders import load_taskset, taskset_config_type + + taskset_id = register_local_taskset_alias(config.taskset_id) + taskset_config = taskset_config_type(taskset_id).model_validate( + {"id": taskset_id, **config.taskset_args} + ) + taskset = load_taskset(taskset_config) + if config.num_tasks is None and taskset.INFINITE: + raise ValueError( + f"Verifiers taskset {config.taskset_id!r} is infinite; " + "num_tasks is required" + ) + tasks = list( + taskset if config.num_tasks is None else taskset.head(config.num_tasks) + ) + if not tasks: + raise ValueError( + f"Verifiers taskset {config.taskset_id!r} yielded no tasks" + ) + if config.num_tasks is not None and len(tasks) != config.num_tasks: + raise ValueError( + f"Verifiers taskset {config.taskset_id!r} yielded {len(tasks)} " + f"tasks, expected {config.num_tasks}" + ) + + self._samples = [ + VerifiersTaskSample(task_data=task.data.model_dump(mode="json")) + for task in tasks + ] + self._rng = random.Random(config.seed) + self._shuffle = config.shuffle + self._order = list(range(len(self._samples))) + self._position = 0 + if self._shuffle: + self._rng.shuffle(self._order) + + def __iter__(self) -> Iterator[VerifiersTaskSample]: + return self + + def __next__(self) -> VerifiersTaskSample: + if self._position == len(self._order): + self._position = 0 + if self._shuffle: + self._rng.shuffle(self._order) + sample = self._samples[self._order[self._position]] + self._position += 1 + return sample + + def state_dict(self) -> dict: + return { + "rng_state": self._rng.getstate(), + "order": list(self._order), + "position": self._position, + } + + def load_state_dict(self, state_dict: dict) -> None: + self._rng.setstate(state_dict["rng_state"]) + self._order = list(state_dict["order"]) + self._position = int(state_dict["position"]) diff --git a/torchtitan/experiments/rl/examples/verifiers/components/env_server.py b/torchtitan/experiments/rl/examples/verifiers/components/env_server.py new file mode 100644 index 0000000000..0bfdf35014 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/components/env_server.py @@ -0,0 +1,179 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import asyncio +import contextlib +import multiprocessing +import os +from dataclasses import dataclass +from pathlib import Path +from queue import Empty +from typing import Any + +from torchtitan.config import Configurable +from torchtitan.experiments.rl.examples.verifiers.components.dataset import ( + register_local_taskset_alias, +) + + +def _run_env_server_process( + config_path: str, + address: str, + address_queue: Any, + death_pipe: Any, +) -> None: + """Run a Verifiers EnvServer from a TOML file in a spawned process.""" + import tomllib + from functools import partial + + from verifiers.v1.configs.serve import pool_serve_kwargs, ServeConfig + from verifiers.v1.serve import env_config_data, serve_env + from verifiers.v1.utils.loaders import resolve_env_config + from verifiers.v1.utils.logging import setup_logging + + no_proxy = os.environ.get("no_proxy", "") + no_proxy = ",".join(filter(None, (no_proxy, "127.0.0.1", "localhost"))) + os.environ["no_proxy"] = no_proxy + os.environ["NO_PROXY"] = no_proxy + + with open(config_path, "rb") as file: + data = tomllib.load(file) + taskset = data.get("env", {}).get("taskset", {}) + if taskset_id := taskset.get("id"): + taskset["id"] = register_local_taskset_alias(taskset_id) + env_config = resolve_env_config(data.get("env")) + serve_config = ServeConfig.model_validate(data.get("serve", {})) + serve_env( + **pool_serve_kwargs(serve_config.pool), + legacy=False, + address=address, + address_queue=address_queue, + death_pipe=death_pipe, + log_setup=partial(setup_logging, "INFO"), + config_data=env_config_data(env_config), + max_concurrent=serve_config.max_concurrent, + ) + + +class VerifiersEnvServer(Configurable): + """Locally managed Verifiers EnvServer process.""" + + @dataclass(kw_only=True, slots=True) + class Config(Configurable.Config): + # TOML file defining the Verifiers environment and server pool. + config_path: str + # ZMQ address for the local server; port 0 requests an ephemeral port. + bind_address: str = "tcp://127.0.0.1:0" + # Maximum time to wait for the server process to publish its address. + startup_timeout_sec: float = 120.0 + + def __post_init__(self) -> None: + if not Path(self.config_path).is_file(): + raise ValueError( + f"Verifiers EnvServer config does not exist: {self.config_path}" + ) + if self.startup_timeout_sec <= 0: + raise ValueError("startup_timeout_sec must be positive") + + def __init__(self, config: Config) -> None: + self.config = config + self.process: Any = None + self.address_queue: Any = None + self.parent_conn: Any = None + self.address: str | None = None + + async def start(self) -> str: + """Start the server and return its resolved ZMQ address.""" + if self.address is not None: + return self.address + + context = multiprocessing.get_context("spawn") + address_queue = context.Queue() + parent_conn, child_conn = context.Pipe() + process = context.Process( + target=_run_env_server_process, + args=( + self.config.config_path, + self.config.bind_address, + address_queue, + child_conn, + ), + daemon=False, + ) + process.start() + child_conn.close() + + deadline = asyncio.get_running_loop().time() + self.config.startup_timeout_sec + while True: + try: + address = address_queue.get_nowait() + break + except Empty: + if not process.is_alive(): + exit_code = process.exitcode + await self._close_server_process_resources( + process=process, + address_queue=address_queue, + parent_conn=parent_conn, + ) + raise RuntimeError( + f"Verifiers EnvServer exited with code {exit_code}" + ) from None + if asyncio.get_running_loop().time() >= deadline: + await self._close_server_process_resources( + process=process, + address_queue=address_queue, + parent_conn=parent_conn, + ) + raise TimeoutError( + "Verifiers EnvServer did not publish its address within " + f"{self.config.startup_timeout_sec} seconds" + ) from None + await asyncio.sleep(0.1) + + self.process = process + self.address_queue = address_queue + self.parent_conn = parent_conn + self.address = address + return address + + async def close(self) -> None: + """Stop the EnvServer and release its multiprocessing resources.""" + process = self.process + address_queue = self.address_queue + parent_conn = self.parent_conn + self.process = None + self.address_queue = None + self.parent_conn = None + self.address = None + await self._close_server_process_resources( + process=process, + address_queue=address_queue, + parent_conn=parent_conn, + ) + + @staticmethod + async def _close_server_process_resources( + *, + process: Any, + address_queue: Any, + parent_conn: Any, + ) -> None: + if process is not None: + process.terminate() + await asyncio.to_thread(process.join, 10) + if process.is_alive(): + process.kill() + await asyncio.to_thread(process.join, 5) + process.close() + if parent_conn is not None: + with contextlib.suppress(Exception): + parent_conn.close() + if address_queue is not None: + address_queue.close() + address_queue.join_thread() diff --git a/torchtitan/experiments/rl/examples/verifiers/components/model_adapter.py b/torchtitan/experiments/rl/examples/verifiers/components/model_adapter.py new file mode 100644 index 0000000000..552bc53f80 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/components/model_adapter.py @@ -0,0 +1,255 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass + +from aiohttp import web + +from torchtitan.experiments.rl.rollout.types import GenerateFn + +logger = logging.getLogger(__name__) + +_SESSION_ID_HEADER = "X-Session-ID" + + +@dataclass(frozen=True, slots=True) +class GenerationMetadata: + """Metadata for one successful model call not stored in its Verifiers trace.""" + + min_policy_version: int + max_policy_version: int + metrics: list + + +class GeneratorModelAdapter: + """Expose a TorchTitan ``GenerateFn`` through Verifiers' token API.""" + + def __init__( + self, + *, + host: str, + port: int, + model: str, + max_model_len: int, + ) -> None: + self.host = host + self.requested_port = port + self.model = model + self.max_model_len = max_model_len + self.generate_fn: GenerateFn | None = None + self.runner: web.AppRunner | None = None + self.bound_port: int | None = None + self.turn_counts: dict[str, int] = {} + self.generation_metadata: dict[str, list[GenerationMetadata]] = {} + + @property + def port(self) -> int: + if self.bound_port is None: + raise RuntimeError("GeneratorModelAdapter has not started") + return self.bound_port + + def set_generate_fn(self, generate_fn: GenerateFn) -> None: + self.generate_fn = generate_fn + + async def start(self) -> None: + if self.runner is not None: + return + app = web.Application() + app.router.add_get("/healthz", self._handle_health_request) + app.router.add_get("/v1/models", self._handle_models_request) + app.router.add_post("/inference/v1/generate", self._handle_generate_request) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, self.host, self.requested_port) + await site.start() + sockets = getattr(site._server, "sockets", None) + if not sockets: + await runner.cleanup() + raise RuntimeError("model adapter did not bind a listening socket") + self.runner = runner + self.bound_port = int(sockets[0].getsockname()[1]) + logger.info( + "Verifiers model adapter listening on http://%s:%d", + self.host, + self.bound_port, + ) + + async def close(self) -> None: + runner = self.runner + if runner is not None: + await runner.cleanup() + self.runner = None + self.bound_port = None + self.turn_counts.clear() + self.generation_metadata.clear() + + def pop_generation_metadata(self, session_id: str) -> list[GenerationMetadata]: + """Remove and return generation metadata recorded for one rollout.""" + self.turn_counts.pop(session_id, None) + return self.generation_metadata.pop(session_id, []) + + async def _handle_health_request(self, request: web.Request) -> web.Response: + del request + return web.json_response({"status": "ok"}) + + async def _handle_models_request(self, request: web.Request) -> web.Response: + del request + return web.json_response( + { + "object": "list", + "data": [ + { + "id": self.model, + "object": "model", + "owned_by": "torchtitan", + "max_model_len": self.max_model_len, + } + ], + } + ) + + async def _handle_generate_request(self, request: web.Request) -> web.Response: + if self.generate_fn is None: + return web.json_response( + {"error": "TorchTitan GenerateFn is not ready"}, status=503 + ) + session_id = request.headers.get(_SESSION_ID_HEADER) + if not session_id: + return web.json_response( + {"error": f"missing {_SESSION_ID_HEADER} header"}, status=400 + ) + + try: + body = await request.json() + prompt_token_ids = _validate_token_ids( + body.get("token_ids"), field_name="token_ids" + ) + sampling = _parse_sampling_config(body.get("sampling_params")) + if body.get("features") is not None: + raise ValueError("multimodal features are not supported") + except (TypeError, ValueError) as error: + return web.json_response({"error": str(error)}, status=400) + + turn_id = self.turn_counts.get(session_id, 0) + self.turn_counts[session_id] = turn_id + 1 + request_id = f"{session_id}/turn={turn_id}" + try: + completion = await self.generate_fn( + prompt_token_ids, + request_id=request_id, + routing_session_id=session_id, + sampling_config=sampling, + ) + except asyncio.CancelledError: + raise + except Exception as error: + logger.exception("TorchTitan generation failed for %s", request_id) + return web.json_response({"error": str(error)}, status=500) + + if completion is None: + return web.json_response( + {"error": f"generation returned no completion for {request_id}"}, + status=502, + ) + if len(completion.token_ids) != len(completion.token_logprobs): + return web.json_response( + {"error": "completion token IDs and logprobs have different lengths"}, + status=500, + ) + if completion.finish_reason not in ("stop", "length"): + return web.json_response( + { + "error": "generation finished without a usable completion: " + f"{completion.finish_reason}" + }, + status=502, + ) + + self.generation_metadata.setdefault(session_id, []).append( + GenerationMetadata( + min_policy_version=completion.min_policy_version, + max_policy_version=completion.max_policy_version, + metrics=list(completion.metrics), + ) + ) + return web.json_response( + { + "request_id": completion.request_id, + "choices": [ + { + "index": 0, + "token_ids": completion.token_ids, + "logprobs": { + "content": [ + { + "token": f"token_id:{token_id}", + "logprob": logprob, + } + for token_id, logprob in zip( + completion.token_ids, + completion.token_logprobs, + strict=True, + ) + ] + }, + "finish_reason": completion.finish_reason, + } + ], + "prompt_logprobs": None, + "kv_transfer_params": None, + } + ) + + +def _validate_token_ids(value: object, *, field_name: str) -> list[int]: + """Validate an untyped JSON value as integer token IDs and return a copy.""" + if not isinstance(value, list) or any( + isinstance(token_id, bool) or not isinstance(token_id, int) + for token_id in value + ): + raise ValueError(f"{field_name} must be a list of integer token IDs") + return list(value) + + +def _parse_sampling_config(value: object): + """Convert Verifiers' vLLM sampling payload to TorchTitan config.""" + from torchtitan.experiments.rl.actors.generator import SamplingConfig + + if not isinstance(value, dict): + raise ValueError("sampling_params must be an object") + supported = { + "temperature", + "top_p", + "max_tokens", + "seed", + "stop_token_ids", + } + protocol_fields = { + "logprobs", + "skip_special_tokens", + "routed_experts_prompt_start", + } + unsupported = set(value) - supported - protocol_fields + if unsupported: + raise ValueError(f"unsupported sampling parameters: {sorted(unsupported)}") + + stop_token_ids = value.get("stop_token_ids") + if stop_token_ids is not None: + stop_token_ids = _validate_token_ids( + stop_token_ids, + field_name="stop_token_ids", + ) + return SamplingConfig( + temperature=float(value.get("temperature", 0.8)), + top_p=float(value.get("top_p", 0.95)), + max_tokens=int(value.get("max_tokens", 100)), + seed=value.get("seed"), + stop_token_ids=stop_token_ids, + ) diff --git a/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py new file mode 100644 index 0000000000..5665b00eba --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py @@ -0,0 +1,392 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Rollouter backed by a Verifiers environment service.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field, replace +from typing import Any, TYPE_CHECKING + +from torchtitan.config import Configurable +from torchtitan.experiments.rl.examples.verifiers.components.dataset import ( + VerifiersTaskSample, +) +from torchtitan.experiments.rl.examples.verifiers.components.env_server import ( + VerifiersEnvServer, +) +from torchtitan.experiments.rl.rollout.advantage import AdvantageEstimator +from torchtitan.experiments.rl.rollout.rollouter import Rollouter, RolloutWorker +from torchtitan.experiments.rl.rollout.types import ( + GenerateFn, + Rollout, + RolloutGroup, + RolloutStatus, + RolloutTurn, +) +from torchtitan.experiments.rl.rubrics import RewardFn, Rubric +from torchtitan.experiments.rl.types import RolloutTurnID + +if TYPE_CHECKING: + from torchtitan.experiments.rl.actors.generator import SamplingConfig + from torchtitan.experiments.rl.examples.verifiers.components.model_adapter import ( + GenerationMetadata, + GeneratorModelAdapter, + ) + from torchtitan.experiments.rl.renderer import RendererConfig + + +VERIFIERS_REWARD_KEY = "verifiers_reward" + + +class VerifiersRewardFn(RewardFn): + """Return the reward produced by Verifiers.""" + + @dataclass(kw_only=True, slots=True) + class Config(RewardFn.Config): + pass + + async def __call__(self, rollout: Rollout, env_input: object) -> float: + del env_input + for turn in reversed(rollout.turns): + if VERIFIERS_REWARD_KEY in turn.env_rewards: + return float(turn.env_rewards[VERIFIERS_REWARD_KEY]) + return 0.0 + + +class VerifiersRollouter(Rollouter): + """Run rollout groups through a locally managed Verifiers EnvServer.""" + + @dataclass(kw_only=True, slots=True) + class Config(Rollouter.Config): + # Unused because this class replaces the base rollout-worker execution path. + worker: RolloutWorker.Config | None = None + # Configuration for the locally managed Verifiers environment server. + env_server: VerifiersEnvServer.Config + # TorchTitan rubric that consumes rewards returned by Verifiers. + rubric: Rubric.Config + # Converts sibling rollout rewards into training advantages. + advantage: Configurable.Config = field( + default_factory=AdvantageEstimator.Config + ) + + # Interface on which the local HTTP model adapter listens. + model_adapter_bind_host: str = "127.0.0.1" + # Adapter port; zero requests an ephemeral port from the operating system. + model_adapter_bind_port: int = 0 + # Base URL given to the Verifiers training client after port substitution. + model_adapter_base_url: str = "http://127.0.0.1:{port}/v1" + # Maximum concurrent rollouts sharing one renderer instance. + renderer_multiplex: int = 256 + # Context limit advertised by the local model adapter. + max_model_len: int + # Maximum time to wait for the Verifiers server to become healthy. + connection_timeout_sec: float = 120.0 + + def __post_init__(self) -> None: + Rollouter.Config.__post_init__(self) + if not 0 <= self.model_adapter_bind_port <= 65535: + raise ValueError("model_adapter_bind_port must be between 0 and 65535") + if ( + self.model_adapter_bind_port == 0 + and "{port}" not in self.model_adapter_base_url + ): + raise ValueError( + "model_adapter_base_url must contain '{port}' when binding " + "an ephemeral port" + ) + if self.renderer_multiplex <= 0: + raise ValueError("renderer_multiplex must be positive") + if self.max_model_len <= 0: + raise ValueError("max_model_len must be positive") + if self.connection_timeout_sec <= 0: + raise ValueError("connection_timeout_sec must be positive") + + def __init__(self, config: Config) -> None: + super().__init__(config) + self._verifiers_config = config + self._rubric: Rubric = config.rubric.build() + self._advantage_estimator: AdvantageEstimator = config.advantage.build() + self._env_server = config.env_server.build() + self._adapter: GeneratorModelAdapter | None = None + self._env_client: Any = None + self._train_client_config: Any = None + + async def setup_async( + self, + *, + renderer_config: RendererConfig, + hf_assets_path: str, + ) -> None: + """Start the EnvServer and connect it to TorchTitan generation.""" + if self._env_client is not None: + return + + from verifiers.v1.configs.client import TrainClientConfig + from verifiers.v1.serve.client import EnvClient + + from torchtitan.experiments.rl.examples.verifiers.components.model_adapter import ( + GeneratorModelAdapter, + ) + + adapter = GeneratorModelAdapter( + host=self._verifiers_config.model_adapter_bind_host, + port=self._verifiers_config.model_adapter_bind_port, + model=hf_assets_path, + max_model_len=self._verifiers_config.max_model_len, + ) + server_address = await self._env_server.start() + env_client = None + try: + await adapter.start() + env_client = EnvClient(server_address) + await env_client.wait_for_server_startup( + timeout=self._verifiers_config.connection_timeout_sec + ) + train_client_config = TrainClientConfig( + base_url=self._verifiers_config.model_adapter_base_url.format( + port=adapter.port + ), + # No API key is needed. This intentionally unset variable makes + # Verifiers use "EMPTY" instead of forwarding PRIME_API_KEY. + api_key_var="TORCHTITAN_VERIFIERS_API_KEY", + renderer=renderer_config.as_renderers_config(), + multiplex=self._verifiers_config.renderer_multiplex, + renderer_model_name=hf_assets_path, + ) + except BaseException: + if env_client is not None: + await env_client.close() + try: + await adapter.close() + finally: + await self._env_server.close() + raise + self._adapter = adapter + self._env_client = env_client + self._train_client_config = train_client_config + + async def close(self) -> None: + """Close the Verifiers client, model adapter, and environment server.""" + try: + if self._env_client is not None: + await self._env_client.close() + finally: + self._env_client = None + self._train_client_config = None + try: + if self._adapter is not None: + await self._adapter.close() + finally: + self._adapter = None + await self._env_server.close() + + async def run_group_rollouts( + self, + *, + generate_fn: GenerateFn, + sample: object, + group_id: int, + group_size: int, + sampling: SamplingConfig, + ) -> RolloutGroup: + """Run sibling episodes through Verifiers, then compute advantages.""" + rollouts = await asyncio.gather( + *( + self._run_single_rollout( + generate_fn=generate_fn, + sample=sample, + sampling=( + sampling + if sampling.seed is None + else replace(sampling, seed=sampling.seed + rollout_id) + ), + group_id=group_id, + rollout_id=rollout_id, + ) + for rollout_id in range(group_size) + ) + ) + + outputs = await self._rubric.score_group(rollouts, sample) + for rollout, output in zip(rollouts, outputs, strict=True): + rollout.reward = output.reward + rollout.reward_breakdown = output.reward_breakdown + + group = RolloutGroup(group_id=group_id, rollouts=rollouts) + advantages = self._advantage_estimator(group) + for rollout, advantage in zip(group.rollouts, advantages, strict=True): + rollout.advantage = advantage + return group + + async def _run_single_rollout( + self, + *, + generate_fn: GenerateFn, + sample: object, + sampling: SamplingConfig, + group_id: int, + rollout_id: int, + ) -> Rollout: + """Send one task to Verifiers and convert its trace to a rollout.""" + if not isinstance(sample, VerifiersTaskSample): + raise TypeError("Verifiers requires a VerifiersTaskSample") + if ( + self._adapter is None + or self._env_client is None + or self._train_client_config is None + ): + raise RuntimeError("Verifiers rollouter is not initialized") + + from verifiers.v1.types import SamplingConfig as VerifiersSamplingConfig + + # Route the adapter's HTTP generation requests through TorchTitan's + # controller-provided generator router. + self._adapter.set_generate_fn(generate_fn) + episode = await self._env_client.run( + task_data=sample.task_data, + client=self._train_client_config, + model=self._adapter.model, + sampling=VerifiersSamplingConfig( + temperature=sampling.temperature, + top_p=sampling.top_p, + max_tokens=sampling.max_tokens, + seed=sampling.seed, + ), + ) + traces = [trace for trace in episode.traces if trace.agent.trainable] + if len(traces) != 1: + raise ValueError( + "Verifiers expects one trainable trace per episode; got " + f"{len(traces)}" + ) + trace = traces[0] + generation_metadata = self._adapter.pop_generation_metadata(trace.id) + turns = self.trace_to_rollout_turns( + trace=trace, + generation_metadata=generation_metadata, + group_id=group_id, + rollout_id=rollout_id, + ) + status = self.rollout_status(episode=episode, trace=trace) + if not turns: + status = RolloutStatus.ERROR + else: + turns[-1].env_rewards[VERIFIERS_REWARD_KEY] = trace.reward + return Rollout( + group_id=group_id, + rollout_id=rollout_id, + status=status, + turns=turns, + ) + + @staticmethod + def rollout_status(*, episode: Any, trace: Any) -> RolloutStatus: + if not episode.ok or not trace.ok: + return RolloutStatus.ERROR + if not trace.is_truncated: + return RolloutStatus.COMPLETED + if trace.stop_condition == "max_turns": + return RolloutStatus.TRUNCATED_MAX_TURNS + return RolloutStatus.TRUNCATED_LENGTH + + @staticmethod + def trace_to_rollout_turns( + *, + trace: Any, + generation_metadata: list[GenerationMetadata], + group_id: int, + rollout_id: int, + ) -> list[RolloutTurn]: + """Flatten a Verifiers trace into TorchTitan's trainable rollout turns. + + Verifiers stores messages in ``trace.nodes`` as an indexed graph. Each + node has per-message ``token_ids``, an aligned trainability ``mask``, a + ``sampled`` flag, and logprobs for the sampled positions. Each entry in + ``trace.branches`` is a root-to-leaf node path whose ``token_ids`` are + the concatenated node tokens and whose ``logprobs`` are aligned to that + full sequence. Each successful entry in ``trace.calls`` identifies its + generated assistant node by index. + + A trace may contain branches that share sampled nodes. This conversion + emits one ``RolloutTurn`` per contiguous trainable token span, emits each + shared sampled node once, and attaches TorchTitan policy metadata from + the matching model-generation call. + """ + from verifiers.v1.dialects.chat import message_to_wire + + successful_calls = [call for call in trace.calls if call.node is not None] + if len(successful_calls) != len(generation_metadata): + raise ValueError( + "Verifiers trace/model-adapter call count mismatch: " + f"trace={len(successful_calls)}, adapter={len(generation_metadata)}" + ) + metadata_by_node = { + call.node: call_metadata + for call, call_metadata in zip( + successful_calls, generation_metadata, strict=True + ) + } + node_index = {id(node): index for index, node in enumerate(trace.nodes)} + trained_nodes: set[int] = set() + turns: list[RolloutTurn] = [] + + for branch in trace.branches: + token_ids = branch.token_ids + logprobs = branch.logprobs + branch_offset = 0 + for node in branch.nodes: + index = node_index[id(node)] + mask = list(node.mask) + if node.sampled and any(mask): + if index in trained_nodes: + mask = [False] * len(mask) + else: + trained_nodes.add(index) + for start, end in _trainable_token_spans(mask): + call_metadata = metadata_by_node.get(index) + if call_metadata is None: + raise ValueError( + f"sampled Verifiers node {index} has no generation metadata" + ) + absolute_start = branch_offset + start + absolute_end = branch_offset + end + turns.append( + RolloutTurn( + rollout_id=RolloutTurnID( + group_id=group_id, + rollout_id=rollout_id, + turn_id=len(turns), + ), + prompt_token_ids=list(token_ids[:absolute_start]), + completion_token_ids=list( + token_ids[absolute_start:absolute_end] + ), + completion_logprobs=list( + logprobs[absolute_start:absolute_end] + ), + min_policy_version=call_metadata.min_policy_version, + max_policy_version=call_metadata.max_policy_version, + completion_message=message_to_wire(node.message), + metrics=list(call_metadata.metrics), + ) + ) + branch_offset += len(node.token_ids) + return turns + + +def _trainable_token_spans(mask: list[bool]) -> list[tuple[int, int]]: + """Return half-open token spans marked trainable by a Verifiers node mask.""" + spans: list[tuple[int, int]] = [] + start: int | None = None + for index, sampled in enumerate([*mask, False]): + if sampled and start is None: + start = index + elif not sampled and start is not None: + spans.append((start, index)) + start = None + return spans diff --git a/torchtitan/experiments/rl/examples/verifiers/config_registry.py b/torchtitan/experiments/rl/examples/verifiers/config_registry.py new file mode 100644 index 0000000000..ca2415843c --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/config_registry.py @@ -0,0 +1,146 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""DAPO math recipes using Verifiers for rollout orchestration.""" + +from __future__ import annotations + +from torchtitan.components.checkpointer import CheckpointManager +from torchtitan.components.loss import ChunkedLossWrapper +from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer +from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.experiments.rl.actors.generator import ( + SamplingConfig, + VLLMCudagraphConfig, + VLLMGenerator, +) +from torchtitan.experiments.rl.actors.trainer import PolicyTrainer +from torchtitan.experiments.rl.controller import ( + AsyncLoopConfig, + Controller, + ValidationConfig, +) +from torchtitan.experiments.rl.examples.verifiers.rollouter import ( + VerifiersMathRollouter, +) +from torchtitan.experiments.rl.losses import DAPOLoss +from torchtitan.experiments.rl.models.cast_linear import LMHeadCastConverter +from torchtitan.experiments.rl.models.vllm_registry import InferenceParallelismConfig +from torchtitan.experiments.rl.observability.metrics import MetricsProcessor +from torchtitan.experiments.rl.renderer import RendererConfig +from torchtitan.experiments.rl.routing.inter_generator_router import ( + InterGeneratorRouter, +) +from torchtitan.experiments.rl.routing.strategies import LeastLoadedRoutingStrategy +from torchtitan.models.qwen3 import model_registry + + +def _qwen3_4b_verifiers_config( + *, + max_response_tokens: int, + max_total_tokens: int, + dump_folder: str, +) -> Controller.Config: + """Build the Qwen3-4B DAPO-Math configuration using Verifiers.""" + num_validation_samples = 30 + return Controller.Config( + model_spec=model_registry( + "4B", + attn_backend="varlen", + converters=[LMHeadCastConverter.Config()], + ), + hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-4B-Base", + dump_folder=dump_folder, + async_loop=AsyncLoopConfig( + num_training_steps=150, + num_prompts_per_train_step=8, + num_samples_per_prompt=16, + target_offpolicy_steps=4, + validation=ValidationConfig(num_samples=num_validation_samples), + ), + compile=CompileConfig(enable=True, backend="aot_eager"), + rollouter=VerifiersMathRollouter.Config(max_model_len=max_total_tokens), + renderer=RendererConfig(name="qwen3", enable_thinking=True), + num_generators=6, + generator_router=InterGeneratorRouter.Config( + strategy=LeastLoadedRoutingStrategy.Config() + ), + metrics=MetricsProcessor.Config( + enable_wandb=True, + console_log_keys_validation=[ + "validation_reward/_mean", + "validation_reward/_max", + "validation/response_length/mean", + "timing/validate", + ], + ), + trainer=PolicyTrainer.Config( + optimizer=default_adamw( + lr=1e-6, + betas=(0.9, 0.98), + weight_decay=0.1, + ), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=0, + min_lr_factor=1.0, + ), + training=TrainingConfig( + num_tokens_per_microbatch_per_dp_rank=max_total_tokens, + max_context_length=max_total_tokens, + ), + parallelism=ParallelismConfig( + data_parallel_replicate_degree=1, + data_parallel_shard_degree=1, + tensor_parallel_degree=2, + ), + checkpoint=CheckpointManager.Config( + enable=True, + initial_load_in_hf=True, + interval=100, + last_save_model_only=False, + keep_latest_k=3, + ), + loss=ChunkedLossWrapper.Config( + num_chunks=8, + loss_fn=DAPOLoss.Config( + ratio_clip_low=0.2, + ratio_clip_high=0.28, + ), + ), + ), + generator=VLLMGenerator.Config( + model_dtype="bfloat16", + parallelism=InferenceParallelismConfig( + data_parallel_degree=1, + tensor_parallel_degree=1, + ), + cudagraph=VLLMCudagraphConfig(enable=True), + checkpoint=CheckpointManager.Config(enable=False), + sampling=SamplingConfig( + temperature=1.0, + top_p=1.0, + max_tokens=max_response_tokens, + ), + ), + ) + + +def rl_dapo_qwen3_4b_verifiers_8k() -> Controller.Config: + """Run the DAPO 8K recipe with Verifiers managing math episodes.""" + return _qwen3_4b_verifiers_config( + max_response_tokens=8192, + max_total_tokens=10240, + dump_folder="outputs/rl/qwen3_4b_verifiers_8k", + ) + + +def rl_dapo_qwen3_4b_verifiers_32k() -> Controller.Config: + """Run the DAPO 32K recipe with Verifiers managing math episodes.""" + return _qwen3_4b_verifiers_config( + max_response_tokens=32768, + max_total_tokens=34816, + dump_folder="outputs/rl/qwen3_4b_verifiers_32k", + ) diff --git a/torchtitan/experiments/rl/examples/verifiers/requirements.txt b/torchtitan/experiments/rl/examples/verifiers/requirements.txt new file mode 100644 index 0000000000..03c0bd48cb --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/requirements.txt @@ -0,0 +1,2 @@ +math-verify==0.9.0 +verifiers==0.3.0 diff --git a/torchtitan/experiments/rl/examples/verifiers/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/rollouter.py new file mode 100644 index 0000000000..b095199348 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/rollouter.py @@ -0,0 +1,56 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path + +from torchtitan.experiments.rl.examples.verifiers.components import ( + VerifiersEnvServer, + VerifiersRewardFn, + VerifiersRollouter, + VerifiersTaskDataset, +) +from torchtitan.experiments.rl.rubrics import Rubric + + +# Dotted module path registered as this example's local Verifiers taskset plugin. +_TASKSET_ID = "torchtitan.experiments.rl.examples.verifiers.taskset" + + +class VerifiersMathRollouter(VerifiersRollouter): + """Run DAPO-Math and AIME rollouts through Verifiers.""" + + @dataclass(kw_only=True, slots=True) + class Config(VerifiersRollouter.Config): + train_dataset: VerifiersTaskDataset.Config = field( + default_factory=lambda: VerifiersTaskDataset.Config( + taskset_id=_TASKSET_ID, + taskset_args={"dataset": "dapo_math"}, + seed=42, + ) + ) + validation_dataset: VerifiersTaskDataset.Config = field( + default_factory=lambda: VerifiersTaskDataset.Config( + taskset_id=_TASKSET_ID, + taskset_args={"dataset": "aime2025"}, + seed=99, + shuffle=False, + ) + ) + env_server: VerifiersEnvServer.Config = field( + default_factory=lambda: VerifiersEnvServer.Config( + config_path=str(Path(__file__).with_name("verifiers_env.toml")), + ) + ) + rubric: Rubric.Config = field( + default_factory=lambda: Rubric.Config( + reward_fns=[VerifiersRewardFn.Config(weight=1.0)], + error_reward=0.0, + ) + ) + max_model_len: int = 10240 diff --git a/torchtitan/experiments/rl/examples/verifiers/taskset.py b/torchtitan/experiments/rl/examples/verifiers/taskset.py new file mode 100644 index 0000000000..11143f44d6 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/taskset.py @@ -0,0 +1,69 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Register TorchTitan's DAPO-Math and AIME datasets with Verifiers. + +Verifiers 0.3.0 does not provide a v1 taskset for these datasets. This module +also serves as an example of exposing a custom dataset as a Verifiers taskset. +""" + +from collections.abc import Iterator +from itertools import islice +from typing import Literal + +import verifiers.v1 as vf + +from torchtitan.experiments.rl.examples.dapo_math import ( + AIME2025Dataset, + DapoMathDataset, + DapoMathSample, + score_math_response, +) + + +class VerifiersMathData(vf.TaskData): + ground_truth: str + + +class VerifiersMathTask(vf.Task[VerifiersMathData]): + data: VerifiersMathData + + @vf.reward(weight=1.0) + async def math_verify(self, trace: vf.Trace) -> float: + return score_math_response(trace.last_reply or "", self.data.ground_truth) + + +class VerifiersMathTasksetConfig(vf.TasksetConfig): + dataset: Literal["dapo_math", "aime2025"] = "dapo_math" + + +class VerifiersMathTaskset(vf.Taskset[VerifiersMathTask, VerifiersMathTasksetConfig]): + config: VerifiersMathTasksetConfig + + def load(self) -> list[VerifiersMathTask]: + dataset, num_tasks = _load_math_dataset(self.config.dataset) + return [ + VerifiersMathTask( + VerifiersMathData( + idx=index, + prompt=sample.prompt, + ground_truth=sample.ground_truth, + ), + self.config.task, + ) + for index, sample in enumerate(islice(dataset, num_tasks)) + ] + + +def _load_math_dataset( + name: Literal["dapo_math", "aime2025"], +) -> tuple[Iterator[DapoMathSample], int]: + if name == "dapo_math": + return DapoMathDataset.Config(shuffle=False).build(), 12643 + return AIME2025Dataset.Config().build(), 30 + + +__all__ = ["VerifiersMathTaskset"] diff --git a/torchtitan/experiments/rl/examples/verifiers/verifiers_env.toml b/torchtitan/experiments/rl/examples/verifiers/verifiers_env.toml new file mode 100644 index 0000000000..e055622a04 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/verifiers_env.toml @@ -0,0 +1,16 @@ +# This math task uses a local subprocess and exposes no tools. + +[env.taskset] +id = "torchtitan.experiments.rl.examples.verifiers.taskset" +dataset = "dapo_math" + +[env.agent] +runtime = { type = "subprocess" } +max_turns = 1 + +[env.agent.harness] +id = "null" + +[serve.pool] +type = "static" +num_workers = 1 diff --git a/torchtitan/experiments/rl/renderer.py b/torchtitan/experiments/rl/renderer.py index 7c5c0ccf72..10191cbb5f 100644 --- a/torchtitan/experiments/rl/renderer.py +++ b/torchtitan/experiments/rl/renderer.py @@ -9,7 +9,12 @@ import logging from dataclasses import dataclass, fields -from renderers import config_from_name, create_renderer, Renderer +from renderers import ( + config_from_name, + create_renderer, + Renderer, + RendererConfig as RenderersConfig, +) from torchtitan.config import Configurable @@ -64,18 +69,13 @@ class RendererConfig(Configurable.Config): preserve_all_thinking: bool | None = None preserve_thinking_between_tool_calls: bool | None = None - def build(self, *, tokenizer_path: str) -> Renderer: - # TODO(renderers#70): use TorchTitan's tokenizer once `renderers` supports - # bring-your-own-tokenizer (PR adds a Tokenizer protocol; drops transformers). - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - + def as_renderers_config(self) -> RenderersConfig | None: + """Convert this TorchTitan config to the underlying renderers config.""" # `name=None` (or "auto") -> let `create_renderer` resolve from the tokenizer. renderer_name = _RENDERER_BY_MODEL.get(self.name, self.name) renderer_config = config_from_name(renderer_name) if renderer_name else None if renderer_config is None: - return create_renderer(tokenizer, None) + return None # Rebuild the typed config and pass parameters # that are not None and are supported @@ -90,4 +90,12 @@ def build(self, *, tokenizer_path: str) -> Renderer: logger.info( f"Using renderer {renderer_name}, of type {config_type}, with args {args}" ) - return create_renderer(tokenizer, config_type(**args)) + return config_type(**args) + + def build(self, *, tokenizer_path: str) -> Renderer: + # TODO(renderers#70): use TorchTitan's tokenizer once `renderers` supports + # bring-your-own-tokenizer (PR adds a Tokenizer protocol; drops transformers). + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + return create_renderer(tokenizer, self.as_renderers_config()) diff --git a/torchtitan/experiments/rl/tests/test_verifiers.py b/torchtitan/experiments/rl/tests/test_verifiers.py new file mode 100644 index 0000000000..7fb34cbca7 --- /dev/null +++ b/torchtitan/experiments/rl/tests/test_verifiers.py @@ -0,0 +1,192 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""CPU tests for the optional Verifiers rollout adapter.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +from aiohttp import ClientSession + +from torchtitan.experiments.rl.examples.verifiers.components.model_adapter import ( + GenerationMetadata, + GeneratorModelAdapter, +) +from torchtitan.experiments.rl.examples.verifiers.components.rollouter import ( + _trainable_token_spans, + VerifiersRollouter, +) +from torchtitan.experiments.rl.types import Completion + + +def test_trainable_token_spans() -> None: + assert _trainable_token_spans([False, True, True, False, True]) == [ + (1, 3), + (4, 5), + ] + + +def test_verifiers_trace_preserves_generation_metadata() -> None: + from verifiers.v1.types import AssistantMessage + + node = SimpleNamespace( + token_ids=[10, 11, 12, 13], + mask=[False, False, True, True], + sampled=True, + message=AssistantMessage(content="Answer: $42$"), + ) + trace = SimpleNamespace( + calls=[SimpleNamespace(node=0)], + nodes=[node], + branches=[ + SimpleNamespace( + nodes=[node], + token_ids=[10, 11, 12, 13], + logprobs=[0.0, 0.0, -0.2, -0.3], + ) + ], + ) + turns = VerifiersRollouter.trace_to_rollout_turns( + trace=trace, + generation_metadata=[ + GenerationMetadata( + min_policy_version=3, + max_policy_version=4, + metrics=[], + ) + ], + group_id=5, + rollout_id=2, + ) + + assert len(turns) == 1 + assert turns[0].prompt_token_ids == [10, 11] + assert turns[0].completion_token_ids == [12, 13] + assert turns[0].completion_logprobs == [-0.2, -0.3] + assert turns[0].completion_message == { + "role": "assistant", + "content": "Answer: $42$", + } + assert turns[0].min_policy_version == 3 + assert turns[0].max_policy_version == 4 + + +def test_model_adapter_forwards_token_request() -> None: + async def run_test() -> None: + received = {} + + async def generate_fn( + prompt_token_ids, + *, + request_id, + routing_session_id=None, + sampling_config=None, + ): + received.update( + prompt_token_ids=prompt_token_ids, + request_id=request_id, + routing_session_id=routing_session_id, + sampling_config=sampling_config, + ) + return Completion( + min_policy_version=7, + max_policy_version=8, + request_id=request_id, + token_ids=[31, 32], + token_logprobs=[-0.1, -0.2], + finish_reason="stop", + ) + + adapter = GeneratorModelAdapter( + host="127.0.0.1", + port=0, + model="test-model", + max_model_len=128, + ) + adapter.set_generate_fn(generate_fn) + await adapter.start() + try: + async with ClientSession() as session: + response = await session.post( + f"http://127.0.0.1:{adapter.port}/inference/v1/generate", + headers={"X-Session-ID": "group=1/rollout=2"}, + json={ + "token_ids": [10, 11], + "sampling_params": { + "temperature": 1.0, + "top_p": 0.9, + "max_tokens": 2, + "seed": 4, + "logprobs": 1, + }, + }, + ) + assert response.status == 200 + payload = await response.json() + generation_metadata = adapter.pop_generation_metadata("group=1/rollout=2") + finally: + await adapter.close() + + assert received["prompt_token_ids"] == [10, 11] + assert received["request_id"] == "group=1/rollout=2/turn=0" + assert received["routing_session_id"] == "group=1/rollout=2" + assert received["sampling_config"].seed == 4 + assert payload["choices"][0]["token_ids"] == [31, 32] + assert [ + (item.min_policy_version, item.max_policy_version) + for item in generation_metadata + ] == [(7, 8)] + + asyncio.run(run_test()) + + +def test_model_adapter_rejects_aborted_generation() -> None: + async def run_test() -> None: + async def generate_fn( + prompt_token_ids, + *, + request_id, + routing_session_id=None, + sampling_config=None, + ): + return Completion( + min_policy_version=7, + max_policy_version=7, + request_id=request_id, + token_ids=[], + token_logprobs=[], + finish_reason="abort", + ) + + adapter = GeneratorModelAdapter( + host="127.0.0.1", + port=0, + model="test-model", + max_model_len=128, + ) + adapter.set_generate_fn(generate_fn) + await adapter.start() + try: + async with ClientSession() as session: + response = await session.post( + f"http://127.0.0.1:{adapter.port}/inference/v1/generate", + headers={"X-Session-ID": "group=1/rollout=2"}, + json={"token_ids": [10, 11], "sampling_params": {}}, + ) + assert response.status == 502 + payload = await response.json() + generation_metadata = adapter.pop_generation_metadata("group=1/rollout=2") + finally: + await adapter.close() + + assert payload == { + "error": "generation finished without a usable completion: abort" + } + assert generation_metadata == [] + + asyncio.run(run_test()) diff --git a/torchtitan/experiments/rl/tests/test_verifiers_example.py b/torchtitan/experiments/rl/tests/test_verifiers_example.py new file mode 100644 index 0000000000..2aba5cd7e0 --- /dev/null +++ b/torchtitan/experiments/rl/tests/test_verifiers_example.py @@ -0,0 +1,90 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""CPU tests for the optional Verifiers math example.""" + +from __future__ import annotations + +import asyncio +import tomllib +from types import SimpleNamespace + +import pytest + +pytest.importorskip("verifiers") + +from torchtitan.experiments.rl.examples.dapo_math import DapoMathSample +from torchtitan.experiments.rl.examples.verifiers import taskset +from torchtitan.experiments.rl.examples.verifiers.components import VerifiersTaskDataset +from torchtitan.experiments.rl.examples.verifiers.config_registry import ( + rl_dapo_qwen3_4b_verifiers_8k, +) +from torchtitan.experiments.rl.examples.verifiers.rollouter import ( + VerifiersMathRollouter, +) + + +def test_verifiers_task_scores_math_response() -> None: + math_task = taskset.VerifiersMathTask( + taskset.VerifiersMathData( + idx=0, + prompt="problem", + ground_truth=r"336^\circ", + ) + ) + assert ( + asyncio.run(math_task.math_verify(SimpleNamespace(last_reply="Answer: $336$"))) + == 1.0 + ) + assert ( + asyncio.run(math_task.math_verify(SimpleNamespace(last_reply="Answer: $335$"))) + == 0.0 + ) + + +def test_verifiers_task_dataset_is_resumable(monkeypatch) -> None: + samples = [ + DapoMathSample(prompt="problem 1", ground_truth="34"), + DapoMathSample(prompt="problem 2", ground_truth="113"), + DapoMathSample(prompt="problem 3", ground_truth="7"), + ] + monkeypatch.setattr(taskset, "_load_math_dataset", lambda name: (iter(samples), 3)) + config = VerifiersTaskDataset.Config( + taskset_id="torchtitan.experiments.rl.examples.verifiers.taskset", + taskset_args={"dataset": "dapo_math"}, + seed=7, + ) + first = config.build() + + second = config.build() + assert [next(first) for _ in range(3)] == [next(second) for _ in range(3)] + + checkpoint = first.state_dict() + expected = [next(first) for _ in range(3)] + resumed = config.build() + resumed.load_state_dict(checkpoint) + assert [next(resumed) for _ in range(3)] == expected + + +def test_verifiers_environment_uses_no_sandbox() -> None: + config_path = VerifiersMathRollouter.Config().env_server.config_path + with open(config_path, "rb") as file: + config = tomllib.load(file) + + assert config["env"]["agent"]["runtime"]["type"] == "subprocess" + assert config["env"]["agent"]["harness"]["id"] == "null" + + +def test_verifiers_config_keeps_dapo_training_recipe() -> None: + config = rl_dapo_qwen3_4b_verifiers_8k() + renderer_config = config.renderer.as_renderers_config() + + assert isinstance(config.rollouter, VerifiersMathRollouter.Config) + assert config.generator.sampling.max_tokens == 8192 + assert config.dump_folder == "outputs/rl/qwen3_4b_verifiers_8k" + assert renderer_config is not None + assert renderer_config.name == "qwen3" + assert renderer_config.enable_thinking