From dc7811b4e9f82d49e1a0fe4e8165eaeca289347f Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Thu, 27 Aug 2026 14:21:56 -0700 Subject: [PATCH 01/11] [rl] Add optional Verifiers rollout integration --- pyproject.toml | 3 + torchtitan/experiments/__init__.py | 1 + torchtitan/experiments/rl/README.md | 8 +- .../rl/examples/verifiers/README.md | 30 ++ .../rl/examples/verifiers/__init__.py | 11 + .../rl/examples/verifiers/config_registry.py | 51 +++ .../rl/examples/verifiers/requirements.txt | 2 + .../rl/examples/verifiers/rollouter.py | 58 +++ .../rl/examples/verifiers/taskset.py | 67 ++++ .../rl/examples/verifiers/verifiers_env.toml | 16 + .../rl/rollout/verifiers/__init__.py | 23 ++ .../rl/rollout/verifiers/dataset.py | 104 +++++ .../rl/rollout/verifiers/env_server.py | 174 +++++++++ .../rl/rollout/verifiers/model_adapter.py | 248 ++++++++++++ .../rl/rollout/verifiers/rollouter.py | 363 ++++++++++++++++++ .../experiments/rl/tests/test_verifiers.py | 137 +++++++ .../rl/tests/test_verifiers_example.py | 85 ++++ 17 files changed, 1380 insertions(+), 1 deletion(-) create mode 100644 torchtitan/experiments/rl/examples/verifiers/README.md create mode 100644 torchtitan/experiments/rl/examples/verifiers/__init__.py create mode 100644 torchtitan/experiments/rl/examples/verifiers/config_registry.py create mode 100644 torchtitan/experiments/rl/examples/verifiers/requirements.txt create mode 100644 torchtitan/experiments/rl/examples/verifiers/rollouter.py create mode 100644 torchtitan/experiments/rl/examples/verifiers/taskset.py create mode 100644 torchtitan/experiments/rl/examples/verifiers/verifiers_env.toml create mode 100644 torchtitan/experiments/rl/rollout/verifiers/__init__.py create mode 100644 torchtitan/experiments/rl/rollout/verifiers/dataset.py create mode 100644 torchtitan/experiments/rl/rollout/verifiers/env_server.py create mode 100644 torchtitan/experiments/rl/rollout/verifiers/model_adapter.py create mode 100644 torchtitan/experiments/rl/rollout/verifiers/rollouter.py create mode 100644 torchtitan/experiments/rl/tests/test_verifiers.py create mode 100644 torchtitan/experiments/rl/tests/test_verifiers_example.py 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..c01e10d1cf --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/README.md @@ -0,0 +1,30 @@ +# 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 imported only by this example and the Verifiers rollout adapter. 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/config_registry.py b/torchtitan/experiments/rl/examples/verifiers/config_registry.py new file mode 100644 index 0000000000..269c0ee34d --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/config_registry.py @@ -0,0 +1,51 @@ +# 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 dataclasses import replace + +from torchtitan.experiments.rl.controller import Controller +from torchtitan.experiments.rl.examples.dapo_math.config_registry import ( + rl_dapo_qwen3_4b_math_32k, + rl_dapo_qwen3_4b_math_8k, +) +from torchtitan.experiments.rl.examples.verifiers.rollouter import ( + VerifiersMathRollouter, +) + + +def _with_verifiers( + config: Controller.Config, + *, + max_model_len: int, + dump_folder: str, +) -> Controller.Config: + return replace( + config, + dump_folder=dump_folder, + rollouter=VerifiersMathRollouter.Config(max_model_len=max_model_len), + ) + + +def rl_dapo_qwen3_4b_verifiers_8k() -> Controller.Config: + """Run the DAPO 8K recipe with Verifiers managing math episodes.""" + return _with_verifiers( + rl_dapo_qwen3_4b_math_8k(), + max_model_len=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 _with_verifiers( + rl_dapo_qwen3_4b_math_32k(), + max_model_len=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..6b3a76cfb7 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/rollouter.py @@ -0,0 +1,58 @@ +# 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.rollout.verifiers import ( + VerifiersEnvServer, + VerifiersRewardFn, + VerifiersRollouter, + VerifiersTaskDataset, +) +from torchtitan.experiments.rl.rubrics import Rubric + + +_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"}, + num_tasks=12643, + seed=42, + ) + ) + validation_dataset: VerifiersTaskDataset.Config = field( + default_factory=lambda: VerifiersTaskDataset.Config( + taskset_id=_TASKSET_ID, + taskset_args={"dataset": "aime2025"}, + num_tasks=30, + 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, + ) + ) + model_name: str = "Qwen/Qwen3-4B-Base" + 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..e60cf96227 --- /dev/null +++ b/torchtitan/experiments/rl/examples/verifiers/taskset.py @@ -0,0 +1,67 @@ +# 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. + +"""Verifiers taskset for the DAPO math datasets.""" + +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 = _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 _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/rollout/verifiers/__init__.py b/torchtitan/experiments/rl/rollout/verifiers/__init__.py new file mode 100644 index 0000000000..cf91aa7f59 --- /dev/null +++ b/torchtitan/experiments/rl/rollout/verifiers/__init__.py @@ -0,0 +1,23 @@ +# 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.rollout.verifiers.dataset import ( + VerifiersTaskDataset, + VerifiersTaskSample, +) +from torchtitan.experiments.rl.rollout.verifiers.env_server import VerifiersEnvServer +from torchtitan.experiments.rl.rollout.verifiers.rollouter import ( + VerifiersRewardFn, + VerifiersRollouter, +) + +__all__ = [ + "VerifiersEnvServer", + "VerifiersRewardFn", + "VerifiersRollouter", + "VerifiersTaskDataset", + "VerifiersTaskSample", +] diff --git a/torchtitan/experiments/rl/rollout/verifiers/dataset.py b/torchtitan/experiments/rl/rollout/verifiers/dataset.py new file mode 100644 index 0000000000..000b2e069a --- /dev/null +++ b/torchtitan/experiments/rl/rollout/verifiers/dataset.py @@ -0,0 +1,104 @@ +# 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 _load_taskset_id(taskset_id: str) -> str: + """Load a dotted local taskset module under a Verifiers-compatible alias.""" + 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): + taskset_id: str + num_tasks: int + taskset_args: dict[str, Any] = field(default_factory=dict) + seed: int = 42 + shuffle: bool = True + + def __post_init__(self) -> None: + if 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 = _load_taskset_id(config.taskset_id) + taskset_config = taskset_config_type(taskset_id).model_validate( + {"id": taskset_id, **config.taskset_args} + ) + tasks = list(load_taskset(taskset_config).head(config.num_tasks)) + if 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/rollout/verifiers/env_server.py b/torchtitan/experiments/rl/rollout/verifiers/env_server.py new file mode 100644 index 0000000000..1c8bf41e6a --- /dev/null +++ b/torchtitan/experiments/rl/rollout/verifiers/env_server.py @@ -0,0 +1,174 @@ +# 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.rollout.verifiers.dataset import _load_taskset_id + + +def _serve_env_from_config( + 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"] = _load_taskset_id(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): + config_path: str + bind_address: str = "tcp://127.0.0.1:0" + 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=_serve_env_from_config, + 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_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_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_resources( + process=process, + address_queue=address_queue, + parent_conn=parent_conn, + ) + + @staticmethod + async def _close_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/rollout/verifiers/model_adapter.py b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py new file mode 100644 index 0000000000..aba89e8564 --- /dev/null +++ b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py @@ -0,0 +1,248 @@ +# 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 GenerationEvidence: + """TorchTitan metadata that is not represented in Verifiers traces.""" + + 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._port: int | None = None + self._turn_counts: dict[str, int] = {} + self._evidence: dict[str, list[GenerationEvidence]] = {} + + @property + def port(self) -> int: + if self._port is None: + raise RuntimeError("GeneratorModelAdapter has not started") + return self._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._health) + app.router.add_get("/v1/models", self._models) + app.router.add_post("/inference/v1/generate", self._generate) + 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._port = int(sockets[0].getsockname()[1]) + logger.info( + "Verifiers model adapter listening on http://%s:%d", + self._host, + self._port, + ) + + async def close(self) -> None: + runner = self._runner + if runner is not None: + await runner.cleanup() + self._runner = None + self._port = None + self._turn_counts.clear() + self._evidence.clear() + + def take_evidence(self, session_id: str) -> list[GenerationEvidence]: + self._turn_counts.pop(session_id, None) + return self._evidence.pop(session_id, []) + + async def _health(self, request: web.Request) -> web.Response: + del request + return web.json_response({"status": "ok"}) + + async def _models(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 _generate(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 = _token_id_list( + body.get("token_ids"), field_name="token_ids" + ) + sampling = _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, + ) + + self._evidence.setdefault(session_id, []).append( + GenerationEvidence( + min_policy_version=completion.min_policy_version, + max_policy_version=completion.max_policy_version, + metrics=list(completion.metrics), + ) + ) + finish_reason = completion.finish_reason + if finish_reason not in ("stop", "length"): + finish_reason = "stop" + 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": finish_reason, + } + ], + "prompt_logprobs": None, + "kv_transfer_params": None, + } + ) + + +def _token_id_list(value: object, *, field_name: str) -> list[int]: + 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 _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 = _token_id_list( + 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/rollout/verifiers/rollouter.py b/torchtitan/experiments/rl/rollout/verifiers/rollouter.py new file mode 100644 index 0000000000..c769f3f890 --- /dev/null +++ b/torchtitan/experiments/rl/rollout/verifiers/rollouter.py @@ -0,0 +1,363 @@ +# 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 renderers import config_from_name + +from torchtitan.config import Configurable +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.rollout.verifiers.dataset import VerifiersTaskSample +from torchtitan.experiments.rl.rollout.verifiers.env_server import VerifiersEnvServer +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.renderer import RendererConfig + from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( + GenerationEvidence, + ) + + +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): + worker: RolloutWorker.Config | None = None + env_server: VerifiersEnvServer.Config + rubric: Rubric.Config + advantage: Configurable.Config = field( + default_factory=AdvantageEstimator.Config + ) + + model_adapter_bind_host: str = "127.0.0.1" + model_adapter_bind_port: int = 0 + model_adapter_base_url: str = "http://127.0.0.1:{port}/v1" + model_name: str + renderer_model_name: str | None = None + renderer_name: str = "qwen3" + renderer_kwargs: dict[str, Any] = field( + default_factory=lambda: { + "enable_thinking": True, + "thinking_retention": "all", + } + ) + renderer_multiplex: int = 256 + max_model_len: int + 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: + from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( + GeneratorModelAdapter, + ) + + 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( + host=config.model_adapter_bind_host, + port=config.model_adapter_bind_port, + model=config.model_name, + max_model_len=config.max_model_len, + ) + 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.""" + del renderer_config + if self._env_client is not None: + return + + from verifiers.v1.configs.client import TrainClientConfig + from verifiers.v1.serve.client import EnvClient + + server_address = await self._env_server.start() + env_client = None + try: + await self._adapter.start() + env_client = EnvClient(server_address) + await env_client.wait_for_server_startup( + timeout=self._verifiers_config.connection_timeout_sec + ) + renderer = config_from_name(self._verifiers_config.renderer_name) + if renderer is None: + raise ValueError("renderer_name must resolve to a concrete renderer") + renderer = type(renderer)(**self._verifiers_config.renderer_kwargs) + train_client_config = TrainClientConfig( + base_url=self._verifiers_config.model_adapter_base_url.format( + port=self._adapter.port + ), + api_key_var="TORCHTITAN_VERIFIERS_API_KEY", + renderer=renderer, + multiplex=self._verifiers_config.renderer_multiplex, + renderer_model_name=( + self._verifiers_config.renderer_model_name or hf_assets_path + ), + ) + except BaseException: + if env_client is not None: + await env_client.close() + try: + await self._adapter.close() + finally: + await self._env_server.close() + raise + 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: + await self._adapter.close() + finally: + 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._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 + + 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._verifiers_config.model_name, + 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] + evidence = self._adapter.take_evidence(trace.id) + turns = self.trace_to_rollout_turns( + trace=trace, + evidence=evidence, + 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, + evidence: list[GenerationEvidence], + group_id: int, + rollout_id: int, + ) -> list[RolloutTurn]: + successful_calls = [call for call in trace.calls if call.node is not None] + if len(successful_calls) != len(evidence): + raise ValueError( + "Verifiers trace/model-adapter call count mismatch: " + f"trace={len(successful_calls)}, adapter={len(evidence)}" + ) + evidence_by_node = { + call.node: call_evidence + for call, call_evidence in zip(successful_calls, evidence, 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_evidence = evidence_by_node.get(index) + if call_evidence is None: + raise ValueError( + f"sampled Verifiers node {index} has no generation evidence" + ) + 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_evidence.min_policy_version, + max_policy_version=call_evidence.max_policy_version, + metrics=list(call_evidence.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/tests/test_verifiers.py b/torchtitan/experiments/rl/tests/test_verifiers.py new file mode 100644 index 0000000000..a3b82199a7 --- /dev/null +++ b/torchtitan/experiments/rl/tests/test_verifiers.py @@ -0,0 +1,137 @@ +# 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.rollout.verifiers.model_adapter import ( + GenerationEvidence, + GeneratorModelAdapter, +) +from torchtitan.experiments.rl.rollout.verifiers.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_evidence() -> None: + node = SimpleNamespace( + token_ids=[10, 11, 12, 13], + mask=[False, False, True, True], + sampled=True, + ) + 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, + evidence=[ + GenerationEvidence( + 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].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() + evidence = adapter.take_evidence("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 evidence + ] == [(7, 8)] + + 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..05b338e20b --- /dev/null +++ b/torchtitan/experiments/rl/tests/test_verifiers_example.py @@ -0,0 +1,85 @@ +# 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.config_registry import ( + rl_dapo_qwen3_4b_verifiers_8k, +) +from torchtitan.experiments.rl.examples.verifiers.rollouter import ( + VerifiersMathRollouter, +) +from torchtitan.experiments.rl.rollout.verifiers import VerifiersTaskDataset + + +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, "_dataset", lambda name: (iter(samples), 3)) + config = VerifiersTaskDataset.Config( + taskset_id="torchtitan.experiments.rl.examples.verifiers.taskset", + taskset_args={"dataset": "dapo_math"}, + num_tasks=3, + 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() + + assert isinstance(config.rollouter, VerifiersMathRollouter.Config) + assert config.generator.sampling.max_tokens == 8192 + assert config.dump_folder == "outputs/rl/qwen3_4b_verifiers_8k" From a13674c0442393298bdb45b6d4a53e848733c838 Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Mon, 31 Aug 2026 10:47:09 -0700 Subject: [PATCH 02/11] [rl] Clean up Verifiers rollout integration --- .../rl/examples/verifiers/rollouter.py | 2 - .../rl/examples/verifiers/taskset.py | 4 +- .../rl/rollout/verifiers/dataset.py | 20 ++++++-- .../rl/rollout/verifiers/model_adapter.py | 13 +++-- .../experiments/rl/tests/test_verifiers.py | 47 +++++++++++++++++++ .../rl/tests/test_verifiers_example.py | 15 +++--- 6 files changed, 81 insertions(+), 20 deletions(-) diff --git a/torchtitan/experiments/rl/examples/verifiers/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/rollouter.py index 6b3a76cfb7..92729af3cd 100644 --- a/torchtitan/experiments/rl/examples/verifiers/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/rollouter.py @@ -30,7 +30,6 @@ class Config(VerifiersRollouter.Config): default_factory=lambda: VerifiersTaskDataset.Config( taskset_id=_TASKSET_ID, taskset_args={"dataset": "dapo_math"}, - num_tasks=12643, seed=42, ) ) @@ -38,7 +37,6 @@ class Config(VerifiersRollouter.Config): default_factory=lambda: VerifiersTaskDataset.Config( taskset_id=_TASKSET_ID, taskset_args={"dataset": "aime2025"}, - num_tasks=30, seed=99, shuffle=False, ) diff --git a/torchtitan/experiments/rl/examples/verifiers/taskset.py b/torchtitan/experiments/rl/examples/verifiers/taskset.py index e60cf96227..6155497991 100644 --- a/torchtitan/experiments/rl/examples/verifiers/taskset.py +++ b/torchtitan/experiments/rl/examples/verifiers/taskset.py @@ -36,9 +36,7 @@ class VerifiersMathTasksetConfig(vf.TasksetConfig): dataset: Literal["dapo_math", "aime2025"] = "dapo_math" -class VerifiersMathTaskset( - vf.Taskset[VerifiersMathTask, VerifiersMathTasksetConfig] -): +class VerifiersMathTaskset(vf.Taskset[VerifiersMathTask, VerifiersMathTasksetConfig]): config: VerifiersMathTasksetConfig def load(self) -> list[VerifiersMathTask]: diff --git a/torchtitan/experiments/rl/rollout/verifiers/dataset.py b/torchtitan/experiments/rl/rollout/verifiers/dataset.py index 000b2e069a..914f159f92 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/dataset.py +++ b/torchtitan/experiments/rl/rollout/verifiers/dataset.py @@ -43,13 +43,13 @@ class VerifiersTaskDataset(Configurable): @dataclass(kw_only=True, slots=True) class Config(Configurable.Config): taskset_id: str - num_tasks: int taskset_args: dict[str, Any] = field(default_factory=dict) + num_tasks: int | None = None seed: int = 42 shuffle: bool = True def __post_init__(self) -> None: - if self.num_tasks <= 0: + 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: @@ -61,8 +61,20 @@ def __init__(self, config: Config) -> None: taskset_config = taskset_config_type(taskset_id).model_validate( {"id": taskset_id, **config.taskset_args} ) - tasks = list(load_taskset(taskset_config).head(config.num_tasks)) - if len(tasks) != config.num_tasks: + 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}" diff --git a/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py index aba89e8564..49a64e0493 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py +++ b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py @@ -162,6 +162,14 @@ async def _generate(self, request: web.Request) -> web.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._evidence.setdefault(session_id, []).append( GenerationEvidence( @@ -170,9 +178,6 @@ async def _generate(self, request: web.Request) -> web.Response: metrics=list(completion.metrics), ) ) - finish_reason = completion.finish_reason - if finish_reason not in ("stop", "length"): - finish_reason = "stop" return web.json_response( { "request_id": completion.request_id, @@ -193,7 +198,7 @@ async def _generate(self, request: web.Request) -> web.Response: ) ] }, - "finish_reason": finish_reason, + "finish_reason": completion.finish_reason, } ], "prompt_logprobs": None, diff --git a/torchtitan/experiments/rl/tests/test_verifiers.py b/torchtitan/experiments/rl/tests/test_verifiers.py index a3b82199a7..22c62a2488 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers.py +++ b/torchtitan/experiments/rl/tests/test_verifiers.py @@ -135,3 +135,50 @@ async def generate_fn( ] == [(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() + evidence = adapter.take_evidence("group=1/rollout=2") + finally: + await adapter.close() + + assert payload == { + "error": "generation finished without a usable completion: abort" + } + assert evidence == [] + + asyncio.run(run_test()) diff --git a/torchtitan/experiments/rl/tests/test_verifiers_example.py b/torchtitan/experiments/rl/tests/test_verifiers_example.py index 05b338e20b..b60343dde6 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers_example.py +++ b/torchtitan/experiments/rl/tests/test_verifiers_example.py @@ -35,12 +35,14 @@ def test_verifiers_task_scores_math_response() -> None: 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 + 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: @@ -53,7 +55,6 @@ def test_verifiers_task_dataset_is_resumable(monkeypatch) -> None: config = VerifiersTaskDataset.Config( taskset_id="torchtitan.experiments.rl.examples.verifiers.taskset", taskset_args={"dataset": "dapo_math"}, - num_tasks=3, seed=7, ) first = config.build() From da007559f9a473c6c0fd299057386525812f2e3a Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Mon, 31 Aug 2026 11:34:01 -0700 Subject: [PATCH 03/11] [rl] Make Verifiers recipes explicit --- .../rl/examples/verifiers/config_registry.py | 131 +++++++++++++++--- 1 file changed, 113 insertions(+), 18 deletions(-) diff --git a/torchtitan/experiments/rl/examples/verifiers/config_registry.py b/torchtitan/experiments/rl/examples/verifiers/config_registry.py index 269c0ee34d..ca2415843c 100644 --- a/torchtitan/experiments/rl/examples/verifiers/config_registry.py +++ b/torchtitan/experiments/rl/examples/verifiers/config_registry.py @@ -8,44 +8,139 @@ from __future__ import annotations -from dataclasses import replace - -from torchtitan.experiments.rl.controller import Controller -from torchtitan.experiments.rl.examples.dapo_math.config_registry import ( - rl_dapo_qwen3_4b_math_32k, - rl_dapo_qwen3_4b_math_8k, +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 _with_verifiers( - config: Controller.Config, +def _qwen3_4b_verifiers_config( *, - max_model_len: int, + max_response_tokens: int, + max_total_tokens: int, dump_folder: str, ) -> Controller.Config: - return replace( - 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, - rollouter=VerifiersMathRollouter.Config(max_model_len=max_model_len), + 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 _with_verifiers( - rl_dapo_qwen3_4b_math_8k(), - max_model_len=10240, + 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 _with_verifiers( - rl_dapo_qwen3_4b_math_32k(), - max_model_len=34816, + return _qwen3_4b_verifiers_config( + max_response_tokens=32768, + max_total_tokens=34816, dump_folder="outputs/rl/qwen3_4b_verifiers_32k", ) From e4fb47fd9677591e6e58695106858c1bee28cece Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Mon, 31 Aug 2026 11:54:00 -0700 Subject: [PATCH 04/11] [rl] Align Verifiers adapter configuration --- .../rl/examples/verifiers/rollouter.py | 1 - torchtitan/experiments/rl/renderer.py | 28 ++++--- .../rl/rollout/verifiers/dataset.py | 5 ++ .../rl/rollout/verifiers/env_server.py | 47 ++++++------ .../rl/rollout/verifiers/model_adapter.py | 71 +++++++++--------- .../rl/rollout/verifiers/rollouter.py | 75 ++++++++++--------- .../rl/tests/test_verifiers_example.py | 4 + 7 files changed, 127 insertions(+), 104 deletions(-) diff --git a/torchtitan/experiments/rl/examples/verifiers/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/rollouter.py index 92729af3cd..9bc30d49a3 100644 --- a/torchtitan/experiments/rl/examples/verifiers/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/rollouter.py @@ -52,5 +52,4 @@ class Config(VerifiersRollouter.Config): error_reward=0.0, ) ) - model_name: str = "Qwen/Qwen3-4B-Base" max_model_len: int = 10240 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/rollout/verifiers/dataset.py b/torchtitan/experiments/rl/rollout/verifiers/dataset.py index 914f159f92..185711f994 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/dataset.py +++ b/torchtitan/experiments/rl/rollout/verifiers/dataset.py @@ -42,10 +42,15 @@ class VerifiersTaskDataset(Configurable): @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: diff --git a/torchtitan/experiments/rl/rollout/verifiers/env_server.py b/torchtitan/experiments/rl/rollout/verifiers/env_server.py index 1c8bf41e6a..91f993686a 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/env_server.py +++ b/torchtitan/experiments/rl/rollout/verifiers/env_server.py @@ -63,8 +63,11 @@ class VerifiersEnvServer(Configurable): @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: @@ -76,16 +79,16 @@ def __post_init__(self) -> None: 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 + 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 + if self.address is not None: + return self.address context = multiprocessing.get_context("spawn") address_queue = context.Queue() @@ -93,8 +96,8 @@ async def start(self) -> str: process = context.Process( target=_serve_env_from_config, args=( - self._config.config_path, - self._config.bind_address, + self.config.config_path, + self.config.bind_address, address_queue, child_conn, ), @@ -103,7 +106,7 @@ async def start(self) -> str: process.start() child_conn.close() - deadline = asyncio.get_running_loop().time() + self._config.startup_timeout_sec + deadline = asyncio.get_running_loop().time() + self.config.startup_timeout_sec while True: try: address = address_queue.get_nowait() @@ -127,25 +130,25 @@ async def start(self) -> str: ) raise TimeoutError( "Verifiers EnvServer did not publish its address within " - f"{self._config.startup_timeout_sec} seconds" + 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 + 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 + 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_resources( process=process, address_queue=address_queue, diff --git a/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py index 49a64e0493..ff8a4b2cbc 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py +++ b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py @@ -39,27 +39,27 @@ def __init__( 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._port: int | None = None - self._turn_counts: dict[str, int] = {} - self._evidence: dict[str, list[GenerationEvidence]] = {} + 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.evidence: dict[str, list[GenerationEvidence]] = {} @property def port(self) -> int: - if self._port is None: + if self.bound_port is None: raise RuntimeError("GeneratorModelAdapter has not started") - return self._port + return self.bound_port def set_generate_fn(self, generate_fn: GenerateFn) -> None: - self._generate_fn = generate_fn + self.generate_fn = generate_fn async def start(self) -> None: - if self._runner is not None: + if self.runner is not None: return app = web.Application() app.router.add_get("/healthz", self._health) @@ -67,32 +67,32 @@ async def start(self) -> None: app.router.add_post("/inference/v1/generate", self._generate) runner = web.AppRunner(app) await runner.setup() - site = web.TCPSite(runner, self._host, self._requested_port) + 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._port = int(sockets[0].getsockname()[1]) + self.runner = runner + self.bound_port = int(sockets[0].getsockname()[1]) logger.info( "Verifiers model adapter listening on http://%s:%d", - self._host, - self._port, + self.host, + self.bound_port, ) async def close(self) -> None: - runner = self._runner + runner = self.runner if runner is not None: await runner.cleanup() - self._runner = None - self._port = None - self._turn_counts.clear() - self._evidence.clear() + self.runner = None + self.bound_port = None + self.turn_counts.clear() + self.evidence.clear() def take_evidence(self, session_id: str) -> list[GenerationEvidence]: - self._turn_counts.pop(session_id, None) - return self._evidence.pop(session_id, []) + self.turn_counts.pop(session_id, None) + return self.evidence.pop(session_id, []) async def _health(self, request: web.Request) -> web.Response: del request @@ -105,17 +105,17 @@ async def _models(self, request: web.Request) -> web.Response: "object": "list", "data": [ { - "id": self._model, + "id": self.model, "object": "model", "owned_by": "torchtitan", - "max_model_len": self._max_model_len, + "max_model_len": self.max_model_len, } ], } ) async def _generate(self, request: web.Request) -> web.Response: - if self._generate_fn is None: + if self.generate_fn is None: return web.json_response( {"error": "TorchTitan GenerateFn is not ready"}, status=503 ) @@ -127,7 +127,7 @@ async def _generate(self, request: web.Request) -> web.Response: try: body = await request.json() - prompt_token_ids = _token_id_list( + prompt_token_ids = _validate_token_ids( body.get("token_ids"), field_name="token_ids" ) sampling = _sampling_config(body.get("sampling_params")) @@ -136,11 +136,11 @@ async def _generate(self, request: web.Request) -> web.Response: 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 + 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( + completion = await self.generate_fn( prompt_token_ids, request_id=request_id, routing_session_id=session_id, @@ -171,7 +171,7 @@ async def _generate(self, request: web.Request) -> web.Response: status=502, ) - self._evidence.setdefault(session_id, []).append( + self.evidence.setdefault(session_id, []).append( GenerationEvidence( min_policy_version=completion.min_policy_version, max_policy_version=completion.max_policy_version, @@ -207,7 +207,8 @@ async def _generate(self, request: web.Request) -> web.Response: ) -def _token_id_list(value: object, *, field_name: str) -> list[int]: +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 @@ -240,7 +241,7 @@ def _sampling_config(value: object): stop_token_ids = value.get("stop_token_ids") if stop_token_ids is not None: - stop_token_ids = _token_id_list( + stop_token_ids = _validate_token_ids( stop_token_ids, field_name="stop_token_ids", ) diff --git a/torchtitan/experiments/rl/rollout/verifiers/rollouter.py b/torchtitan/experiments/rl/rollout/verifiers/rollouter.py index c769f3f890..18fe9df4a6 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/rollouter.py +++ b/torchtitan/experiments/rl/rollout/verifiers/rollouter.py @@ -12,8 +12,6 @@ from dataclasses import dataclass, field, replace from typing import Any, TYPE_CHECKING -from renderers import config_from_name - from torchtitan.config import Configurable from torchtitan.experiments.rl.rollout.advantage import AdvantageEstimator from torchtitan.experiments.rl.rollout.rollouter import Rollouter, RolloutWorker @@ -34,6 +32,7 @@ from torchtitan.experiments.rl.renderer import RendererConfig from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( GenerationEvidence, + GeneratorModelAdapter, ) @@ -60,27 +59,28 @@ class VerifiersRollouter(Rollouter): @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" - model_name: str - renderer_model_name: str | None = None - renderer_name: str = "qwen3" - renderer_kwargs: dict[str, Any] = field( - default_factory=lambda: { - "enable_thinking": True, - "thinking_retention": "all", - } - ) + # 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: @@ -103,21 +103,12 @@ def __post_init__(self) -> None: raise ValueError("connection_timeout_sec must be positive") def __init__(self, config: Config) -> None: - from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( - GeneratorModelAdapter, - ) - 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( - host=config.model_adapter_bind_host, - port=config.model_adapter_bind_port, - model=config.model_name, - max_model_len=config.max_model_len, - ) + self._adapter: GeneratorModelAdapter | None = None self._env_client: Any = None self._train_client_config: Any = None @@ -128,44 +119,50 @@ async def setup_async( hf_assets_path: str, ) -> None: """Start the EnvServer and connect it to TorchTitan generation.""" - del renderer_config 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.rollout.verifiers.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 self._adapter.start() + await adapter.start() env_client = EnvClient(server_address) await env_client.wait_for_server_startup( timeout=self._verifiers_config.connection_timeout_sec ) - renderer = config_from_name(self._verifiers_config.renderer_name) - if renderer is None: - raise ValueError("renderer_name must resolve to a concrete renderer") - renderer = type(renderer)(**self._verifiers_config.renderer_kwargs) train_client_config = TrainClientConfig( base_url=self._verifiers_config.model_adapter_base_url.format( - port=self._adapter.port + port=adapter.port ), + # Keep real provider credentials away from the unauthenticated local + # adapter; Verifiers uses "EMPTY" when this variable is unset. api_key_var="TORCHTITAN_VERIFIERS_API_KEY", - renderer=renderer, + renderer=renderer_config.as_renderers_config(), multiplex=self._verifiers_config.renderer_multiplex, - renderer_model_name=( - self._verifiers_config.renderer_model_name or hf_assets_path - ), + renderer_model_name=hf_assets_path, ) except BaseException: if env_client is not None: await env_client.close() try: - await self._adapter.close() + 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 @@ -178,8 +175,10 @@ async def close(self) -> None: self._env_client = None self._train_client_config = None try: - await self._adapter.close() + if self._adapter is not None: + await self._adapter.close() finally: + self._adapter = None await self._env_server.close() async def run_group_rollouts( @@ -232,7 +231,11 @@ async def _run_single_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._env_client is None or self._train_client_config is None: + 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 @@ -241,7 +244,7 @@ async def _run_single_rollout( episode = await self._env_client.run( task_data=sample.task_data, client=self._train_client_config, - model=self._verifiers_config.model_name, + model=self._adapter.model, sampling=VerifiersSamplingConfig( temperature=sampling.temperature, top_p=sampling.top_p, diff --git a/torchtitan/experiments/rl/tests/test_verifiers_example.py b/torchtitan/experiments/rl/tests/test_verifiers_example.py index b60343dde6..7ffe07a70f 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers_example.py +++ b/torchtitan/experiments/rl/tests/test_verifiers_example.py @@ -80,7 +80,11 @@ def test_verifiers_environment_uses_no_sandbox() -> None: 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 From 0a1707991d6d2971eb67b20e90ad9af4bbd66932 Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Mon, 31 Aug 2026 11:59:39 -0700 Subject: [PATCH 05/11] [rl] Clarify Verifiers helper naming --- .../rl/examples/verifiers/taskset.py | 4 +-- .../rl/rollout/verifiers/dataset.py | 6 ++-- .../rl/rollout/verifiers/env_server.py | 18 +++++----- .../rl/rollout/verifiers/model_adapter.py | 33 ++++++++++--------- .../rl/rollout/verifiers/rollouter.py | 32 +++++++++--------- .../experiments/rl/tests/test_verifiers.py | 17 +++++----- .../rl/tests/test_verifiers_example.py | 2 +- 7 files changed, 59 insertions(+), 53 deletions(-) diff --git a/torchtitan/experiments/rl/examples/verifiers/taskset.py b/torchtitan/experiments/rl/examples/verifiers/taskset.py index 6155497991..54fb1605fe 100644 --- a/torchtitan/experiments/rl/examples/verifiers/taskset.py +++ b/torchtitan/experiments/rl/examples/verifiers/taskset.py @@ -40,7 +40,7 @@ class VerifiersMathTaskset(vf.Taskset[VerifiersMathTask, VerifiersMathTasksetCon config: VerifiersMathTasksetConfig def load(self) -> list[VerifiersMathTask]: - dataset, num_tasks = _dataset(self.config.dataset) + dataset, num_tasks = _load_math_dataset(self.config.dataset) return [ VerifiersMathTask( VerifiersMathData( @@ -54,7 +54,7 @@ def load(self) -> list[VerifiersMathTask]: ] -def _dataset( +def _load_math_dataset( name: Literal["dapo_math", "aime2025"], ) -> tuple[Iterator[DapoMathSample], int]: if name == "dapo_math": diff --git a/torchtitan/experiments/rl/rollout/verifiers/dataset.py b/torchtitan/experiments/rl/rollout/verifiers/dataset.py index 185711f994..2d03d023b3 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/dataset.py +++ b/torchtitan/experiments/rl/rollout/verifiers/dataset.py @@ -16,8 +16,8 @@ from torchtitan.config import Configurable -def _load_taskset_id(taskset_id: str) -> str: - """Load a dotted local taskset module under a Verifiers-compatible alias.""" +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 @@ -62,7 +62,7 @@ def __init__(self, config: Config) -> None: # rest of TorchTitan RL does not require it. from verifiers.v1.utils.loaders import load_taskset, taskset_config_type - taskset_id = _load_taskset_id(config.taskset_id) + taskset_id = register_local_taskset_alias(config.taskset_id) taskset_config = taskset_config_type(taskset_id).model_validate( {"id": taskset_id, **config.taskset_args} ) diff --git a/torchtitan/experiments/rl/rollout/verifiers/env_server.py b/torchtitan/experiments/rl/rollout/verifiers/env_server.py index 91f993686a..c9f84c5595 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/env_server.py +++ b/torchtitan/experiments/rl/rollout/verifiers/env_server.py @@ -16,10 +16,12 @@ from typing import Any from torchtitan.config import Configurable -from torchtitan.experiments.rl.rollout.verifiers.dataset import _load_taskset_id +from torchtitan.experiments.rl.rollout.verifiers.dataset import ( + register_local_taskset_alias, +) -def _serve_env_from_config( +def _run_env_server_process( config_path: str, address: str, address_queue: Any, @@ -43,7 +45,7 @@ def _serve_env_from_config( data = tomllib.load(file) taskset = data.get("env", {}).get("taskset", {}) if taskset_id := taskset.get("id"): - taskset["id"] = _load_taskset_id(taskset_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( @@ -94,7 +96,7 @@ async def start(self) -> str: address_queue = context.Queue() parent_conn, child_conn = context.Pipe() process = context.Process( - target=_serve_env_from_config, + target=_run_env_server_process, args=( self.config.config_path, self.config.bind_address, @@ -114,7 +116,7 @@ async def start(self) -> str: except Empty: if not process.is_alive(): exit_code = process.exitcode - await self._close_resources( + await self._close_server_process_resources( process=process, address_queue=address_queue, parent_conn=parent_conn, @@ -123,7 +125,7 @@ async def start(self) -> str: f"Verifiers EnvServer exited with code {exit_code}" ) from None if asyncio.get_running_loop().time() >= deadline: - await self._close_resources( + await self._close_server_process_resources( process=process, address_queue=address_queue, parent_conn=parent_conn, @@ -149,14 +151,14 @@ async def close(self) -> None: self.address_queue = None self.parent_conn = None self.address = None - await self._close_resources( + await self._close_server_process_resources( process=process, address_queue=address_queue, parent_conn=parent_conn, ) @staticmethod - async def _close_resources( + async def _close_server_process_resources( *, process: Any, address_queue: Any, diff --git a/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py index ff8a4b2cbc..552bc53f80 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py +++ b/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py @@ -20,8 +20,8 @@ @dataclass(frozen=True, slots=True) -class GenerationEvidence: - """TorchTitan metadata that is not represented in Verifiers traces.""" +class GenerationMetadata: + """Metadata for one successful model call not stored in its Verifiers trace.""" min_policy_version: int max_policy_version: int @@ -47,7 +47,7 @@ def __init__( self.runner: web.AppRunner | None = None self.bound_port: int | None = None self.turn_counts: dict[str, int] = {} - self.evidence: dict[str, list[GenerationEvidence]] = {} + self.generation_metadata: dict[str, list[GenerationMetadata]] = {} @property def port(self) -> int: @@ -62,9 +62,9 @@ async def start(self) -> None: if self.runner is not None: return app = web.Application() - app.router.add_get("/healthz", self._health) - app.router.add_get("/v1/models", self._models) - app.router.add_post("/inference/v1/generate", self._generate) + 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) @@ -88,17 +88,18 @@ async def close(self) -> None: self.runner = None self.bound_port = None self.turn_counts.clear() - self.evidence.clear() + self.generation_metadata.clear() - def take_evidence(self, session_id: str) -> list[GenerationEvidence]: + 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.evidence.pop(session_id, []) + return self.generation_metadata.pop(session_id, []) - async def _health(self, request: web.Request) -> web.Response: + async def _handle_health_request(self, request: web.Request) -> web.Response: del request return web.json_response({"status": "ok"}) - async def _models(self, request: web.Request) -> web.Response: + async def _handle_models_request(self, request: web.Request) -> web.Response: del request return web.json_response( { @@ -114,7 +115,7 @@ async def _models(self, request: web.Request) -> web.Response: } ) - async def _generate(self, request: web.Request) -> web.Response: + 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 @@ -130,7 +131,7 @@ async def _generate(self, request: web.Request) -> web.Response: prompt_token_ids = _validate_token_ids( body.get("token_ids"), field_name="token_ids" ) - sampling = _sampling_config(body.get("sampling_params")) + 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: @@ -171,8 +172,8 @@ async def _generate(self, request: web.Request) -> web.Response: status=502, ) - self.evidence.setdefault(session_id, []).append( - GenerationEvidence( + 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), @@ -217,7 +218,7 @@ def _validate_token_ids(value: object, *, field_name: str) -> list[int]: return list(value) -def _sampling_config(value: object): +def _parse_sampling_config(value: object): """Convert Verifiers' vLLM sampling payload to TorchTitan config.""" from torchtitan.experiments.rl.actors.generator import SamplingConfig diff --git a/torchtitan/experiments/rl/rollout/verifiers/rollouter.py b/torchtitan/experiments/rl/rollout/verifiers/rollouter.py index 18fe9df4a6..51a8088dfb 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/rollouter.py +++ b/torchtitan/experiments/rl/rollout/verifiers/rollouter.py @@ -31,7 +31,7 @@ from torchtitan.experiments.rl.actors.generator import SamplingConfig from torchtitan.experiments.rl.renderer import RendererConfig from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( - GenerationEvidence, + GenerationMetadata, GeneratorModelAdapter, ) @@ -259,10 +259,10 @@ async def _run_single_rollout( f"{len(traces)}" ) trace = traces[0] - evidence = self._adapter.take_evidence(trace.id) + generation_metadata = self._adapter.pop_generation_metadata(trace.id) turns = self.trace_to_rollout_turns( trace=trace, - evidence=evidence, + generation_metadata=generation_metadata, group_id=group_id, rollout_id=rollout_id, ) @@ -292,19 +292,21 @@ def rollout_status(*, episode: Any, trace: Any) -> RolloutStatus: def trace_to_rollout_turns( *, trace: Any, - evidence: list[GenerationEvidence], + generation_metadata: list[GenerationMetadata], group_id: int, rollout_id: int, ) -> list[RolloutTurn]: successful_calls = [call for call in trace.calls if call.node is not None] - if len(successful_calls) != len(evidence): + if len(successful_calls) != len(generation_metadata): raise ValueError( "Verifiers trace/model-adapter call count mismatch: " - f"trace={len(successful_calls)}, adapter={len(evidence)}" + 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 ) - evidence_by_node = { - call.node: call_evidence - for call, call_evidence in zip(successful_calls, evidence, strict=True) } node_index = {id(node): index for index, node in enumerate(trace.nodes)} trained_nodes: set[int] = set() @@ -323,10 +325,10 @@ def trace_to_rollout_turns( else: trained_nodes.add(index) for start, end in _trainable_token_spans(mask): - call_evidence = evidence_by_node.get(index) - if call_evidence is None: + call_metadata = metadata_by_node.get(index) + if call_metadata is None: raise ValueError( - f"sampled Verifiers node {index} has no generation evidence" + f"sampled Verifiers node {index} has no generation metadata" ) absolute_start = branch_offset + start absolute_end = branch_offset + end @@ -344,9 +346,9 @@ def trace_to_rollout_turns( completion_logprobs=list( logprobs[absolute_start:absolute_end] ), - min_policy_version=call_evidence.min_policy_version, - max_policy_version=call_evidence.max_policy_version, - metrics=list(call_evidence.metrics), + min_policy_version=call_metadata.min_policy_version, + max_policy_version=call_metadata.max_policy_version, + metrics=list(call_metadata.metrics), ) ) branch_offset += len(node.token_ids) diff --git a/torchtitan/experiments/rl/tests/test_verifiers.py b/torchtitan/experiments/rl/tests/test_verifiers.py index 22c62a2488..efba4f4e73 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers.py +++ b/torchtitan/experiments/rl/tests/test_verifiers.py @@ -14,7 +14,7 @@ from aiohttp import ClientSession from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( - GenerationEvidence, + GenerationMetadata, GeneratorModelAdapter, ) from torchtitan.experiments.rl.rollout.verifiers.rollouter import ( @@ -31,7 +31,7 @@ def test_trainable_token_spans() -> None: ] -def test_verifiers_trace_preserves_generation_evidence() -> None: +def test_verifiers_trace_preserves_generation_metadata() -> None: node = SimpleNamespace( token_ids=[10, 11, 12, 13], mask=[False, False, True, True], @@ -50,8 +50,8 @@ def test_verifiers_trace_preserves_generation_evidence() -> None: ) turns = VerifiersRollouter.trace_to_rollout_turns( trace=trace, - evidence=[ - GenerationEvidence( + generation_metadata=[ + GenerationMetadata( min_policy_version=3, max_policy_version=4, metrics=[], @@ -121,7 +121,7 @@ async def generate_fn( ) assert response.status == 200 payload = await response.json() - evidence = adapter.take_evidence("group=1/rollout=2") + generation_metadata = adapter.pop_generation_metadata("group=1/rollout=2") finally: await adapter.close() @@ -131,7 +131,8 @@ async def generate_fn( 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 evidence + (item.min_policy_version, item.max_policy_version) + for item in generation_metadata ] == [(7, 8)] asyncio.run(run_test()) @@ -172,13 +173,13 @@ async def generate_fn( ) assert response.status == 502 payload = await response.json() - evidence = adapter.take_evidence("group=1/rollout=2") + 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 evidence == [] + 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 index 7ffe07a70f..18c2f33df3 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers_example.py +++ b/torchtitan/experiments/rl/tests/test_verifiers_example.py @@ -51,7 +51,7 @@ def test_verifiers_task_dataset_is_resumable(monkeypatch) -> None: DapoMathSample(prompt="problem 2", ground_truth="113"), DapoMathSample(prompt="problem 3", ground_truth="7"), ] - monkeypatch.setattr(taskset, "_dataset", lambda name: (iter(samples), 3)) + 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"}, From d8afda6f7f1aa0e69d110d704fdb5fccf4e366fa Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Mon, 31 Aug 2026 12:02:16 -0700 Subject: [PATCH 06/11] [rl] Document custom Verifiers taskset --- torchtitan/experiments/rl/examples/verifiers/taskset.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/torchtitan/experiments/rl/examples/verifiers/taskset.py b/torchtitan/experiments/rl/examples/verifiers/taskset.py index 54fb1605fe..11143f44d6 100644 --- a/torchtitan/experiments/rl/examples/verifiers/taskset.py +++ b/torchtitan/experiments/rl/examples/verifiers/taskset.py @@ -4,7 +4,11 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Verifiers taskset for the DAPO math datasets.""" +"""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 f72f690a971e1c248460cc48ab1efa8ce5f876e7 Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Mon, 31 Aug 2026 12:04:32 -0700 Subject: [PATCH 07/11] [rl] Document Verifiers taskset identifier --- torchtitan/experiments/rl/examples/verifiers/rollouter.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchtitan/experiments/rl/examples/verifiers/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/rollouter.py index 9bc30d49a3..92a831643d 100644 --- a/torchtitan/experiments/rl/examples/verifiers/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/rollouter.py @@ -18,6 +18,7 @@ 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" From 36c417d8f36739012c7847f634c9632fb818f9ee Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Mon, 31 Aug 2026 12:11:10 -0700 Subject: [PATCH 08/11] [rl] Move Verifiers integration into example --- .../rl/examples/verifiers/README.md | 3 ++- .../verifiers/components}/__init__.py | 8 +++++--- .../verifiers/components}/dataset.py | 0 .../verifiers/components}/env_server.py | 2 +- .../verifiers/components}/model_adapter.py | 0 .../verifiers/components}/rollouter.py | 18 +++++++++++------- .../rl/examples/verifiers/rollouter.py | 2 +- .../experiments/rl/tests/test_verifiers.py | 4 ++-- .../rl/tests/test_verifiers_example.py | 2 +- 9 files changed, 23 insertions(+), 16 deletions(-) rename torchtitan/experiments/rl/{rollout/verifiers => examples/verifiers/components}/__init__.py (63%) rename torchtitan/experiments/rl/{rollout/verifiers => examples/verifiers/components}/dataset.py (100%) rename torchtitan/experiments/rl/{rollout/verifiers => examples/verifiers/components}/env_server.py (98%) rename torchtitan/experiments/rl/{rollout/verifiers => examples/verifiers/components}/model_adapter.py (100%) rename torchtitan/experiments/rl/{rollout/verifiers => examples/verifiers/components}/rollouter.py (96%) diff --git a/torchtitan/experiments/rl/examples/verifiers/README.md b/torchtitan/experiments/rl/examples/verifiers/README.md index c01e10d1cf..f32026a8bd 100644 --- a/torchtitan/experiments/rl/examples/verifiers/README.md +++ b/torchtitan/experiments/rl/examples/verifiers/README.md @@ -4,7 +4,8 @@ This example keeps the existing [DAPO Math](../dapo_math) recipe unchanged and r 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 imported only by this example and the Verifiers rollout adapter. Other TitanRL recipes do not require it. +Verifiers is optional, and all Verifiers integration code lives in this example. +Other TitanRL recipes do not require it. ## Setup diff --git a/torchtitan/experiments/rl/rollout/verifiers/__init__.py b/torchtitan/experiments/rl/examples/verifiers/components/__init__.py similarity index 63% rename from torchtitan/experiments/rl/rollout/verifiers/__init__.py rename to torchtitan/experiments/rl/examples/verifiers/components/__init__.py index cf91aa7f59..e45b4ea5d3 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/__init__.py +++ b/torchtitan/experiments/rl/examples/verifiers/components/__init__.py @@ -4,12 +4,14 @@ # 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.rollout.verifiers.dataset import ( +from torchtitan.experiments.rl.examples.verifiers.components.dataset import ( VerifiersTaskDataset, VerifiersTaskSample, ) -from torchtitan.experiments.rl.rollout.verifiers.env_server import VerifiersEnvServer -from torchtitan.experiments.rl.rollout.verifiers.rollouter import ( +from torchtitan.experiments.rl.examples.verifiers.components.env_server import ( + VerifiersEnvServer, +) +from torchtitan.experiments.rl.examples.verifiers.components.rollouter import ( VerifiersRewardFn, VerifiersRollouter, ) diff --git a/torchtitan/experiments/rl/rollout/verifiers/dataset.py b/torchtitan/experiments/rl/examples/verifiers/components/dataset.py similarity index 100% rename from torchtitan/experiments/rl/rollout/verifiers/dataset.py rename to torchtitan/experiments/rl/examples/verifiers/components/dataset.py diff --git a/torchtitan/experiments/rl/rollout/verifiers/env_server.py b/torchtitan/experiments/rl/examples/verifiers/components/env_server.py similarity index 98% rename from torchtitan/experiments/rl/rollout/verifiers/env_server.py rename to torchtitan/experiments/rl/examples/verifiers/components/env_server.py index c9f84c5595..0bfdf35014 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/env_server.py +++ b/torchtitan/experiments/rl/examples/verifiers/components/env_server.py @@ -16,7 +16,7 @@ from typing import Any from torchtitan.config import Configurable -from torchtitan.experiments.rl.rollout.verifiers.dataset import ( +from torchtitan.experiments.rl.examples.verifiers.components.dataset import ( register_local_taskset_alias, ) diff --git a/torchtitan/experiments/rl/rollout/verifiers/model_adapter.py b/torchtitan/experiments/rl/examples/verifiers/components/model_adapter.py similarity index 100% rename from torchtitan/experiments/rl/rollout/verifiers/model_adapter.py rename to torchtitan/experiments/rl/examples/verifiers/components/model_adapter.py diff --git a/torchtitan/experiments/rl/rollout/verifiers/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py similarity index 96% rename from torchtitan/experiments/rl/rollout/verifiers/rollouter.py rename to torchtitan/experiments/rl/examples/verifiers/components/rollouter.py index 51a8088dfb..ec3d2d2baa 100644 --- a/torchtitan/experiments/rl/rollout/verifiers/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py @@ -13,6 +13,12 @@ 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 ( @@ -22,18 +28,16 @@ RolloutStatus, RolloutTurn, ) -from torchtitan.experiments.rl.rollout.verifiers.dataset import VerifiersTaskSample -from torchtitan.experiments.rl.rollout.verifiers.env_server import VerifiersEnvServer 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.renderer import RendererConfig - from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( + from torchtitan.experiments.rl.examples.verifiers.components.model_adapter import ( GenerationMetadata, GeneratorModelAdapter, ) + from torchtitan.experiments.rl.renderer import RendererConfig VERIFIERS_REWARD_KEY = "verifiers_reward" @@ -125,7 +129,7 @@ async def setup_async( from verifiers.v1.configs.client import TrainClientConfig from verifiers.v1.serve.client import EnvClient - from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( + from torchtitan.experiments.rl.examples.verifiers.components.model_adapter import ( GeneratorModelAdapter, ) @@ -147,8 +151,8 @@ async def setup_async( base_url=self._verifiers_config.model_adapter_base_url.format( port=adapter.port ), - # Keep real provider credentials away from the unauthenticated local - # adapter; Verifiers uses "EMPTY" when this variable is unset. + # 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, diff --git a/torchtitan/experiments/rl/examples/verifiers/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/rollouter.py index 92a831643d..b095199348 100644 --- a/torchtitan/experiments/rl/examples/verifiers/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/rollouter.py @@ -9,7 +9,7 @@ from dataclasses import dataclass, field from pathlib import Path -from torchtitan.experiments.rl.rollout.verifiers import ( +from torchtitan.experiments.rl.examples.verifiers.components import ( VerifiersEnvServer, VerifiersRewardFn, VerifiersRollouter, diff --git a/torchtitan/experiments/rl/tests/test_verifiers.py b/torchtitan/experiments/rl/tests/test_verifiers.py index efba4f4e73..316b798f9e 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers.py +++ b/torchtitan/experiments/rl/tests/test_verifiers.py @@ -13,11 +13,11 @@ from aiohttp import ClientSession -from torchtitan.experiments.rl.rollout.verifiers.model_adapter import ( +from torchtitan.experiments.rl.examples.verifiers.components.model_adapter import ( GenerationMetadata, GeneratorModelAdapter, ) -from torchtitan.experiments.rl.rollout.verifiers.rollouter import ( +from torchtitan.experiments.rl.examples.verifiers.components.rollouter import ( _trainable_token_spans, VerifiersRollouter, ) diff --git a/torchtitan/experiments/rl/tests/test_verifiers_example.py b/torchtitan/experiments/rl/tests/test_verifiers_example.py index 18c2f33df3..2aba5cd7e0 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers_example.py +++ b/torchtitan/experiments/rl/tests/test_verifiers_example.py @@ -18,13 +18,13 @@ 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, ) -from torchtitan.experiments.rl.rollout.verifiers import VerifiersTaskDataset def test_verifiers_task_scores_math_response() -> None: From 0b4e0fca68802f3221421bbd74500e3b52cef3c0 Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Tue, 1 Sep 2026 09:12:53 -0700 Subject: [PATCH 09/11] [rl] Document Verifiers rollout conversion --- .../rl/examples/verifiers/components/rollouter.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py index ec3d2d2baa..530771613d 100644 --- a/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py @@ -244,6 +244,8 @@ async def _run_single_rollout( 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, @@ -300,6 +302,11 @@ def trace_to_rollout_turns( group_id: int, rollout_id: int, ) -> list[RolloutTurn]: + """Flatten a Verifiers trace into TorchTitan's trainable rollout turns. + + A trace may contain branches that share sampled nodes. Emit each sampled + node once and attach metadata from its matching model-generation call. + """ successful_calls = [call for call in trace.calls if call.node is not None] if len(successful_calls) != len(generation_metadata): raise ValueError( From 74251b66eb436eab4be1b3ca1025c092a3ab078e Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Tue, 1 Sep 2026 09:15:45 -0700 Subject: [PATCH 10/11] [rl] Document Verifiers trace format --- .../rl/examples/verifiers/components/rollouter.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py index 530771613d..15910701e7 100644 --- a/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py @@ -304,8 +304,18 @@ def trace_to_rollout_turns( ) -> list[RolloutTurn]: """Flatten a Verifiers trace into TorchTitan's trainable rollout turns. - A trace may contain branches that share sampled nodes. Emit each sampled - node once and attach metadata from its matching model-generation call. + 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. """ successful_calls = [call for call in trace.calls if call.node is not None] if len(successful_calls) != len(generation_metadata): From d77ef9c83c2a111ded7eeec801c137bda4f864c8 Mon Sep 17 00:00:00 2001 From: Jiani Wang Date: Tue, 1 Sep 2026 09:24:52 -0700 Subject: [PATCH 11/11] [rl] Record Verifiers completion messages --- .../rl/examples/verifiers/components/rollouter.py | 3 +++ torchtitan/experiments/rl/tests/test_verifiers.py | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py index 15910701e7..5665b00eba 100644 --- a/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py +++ b/torchtitan/experiments/rl/examples/verifiers/components/rollouter.py @@ -317,6 +317,8 @@ def trace_to_rollout_turns( 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( @@ -369,6 +371,7 @@ def trace_to_rollout_turns( ), 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), ) ) diff --git a/torchtitan/experiments/rl/tests/test_verifiers.py b/torchtitan/experiments/rl/tests/test_verifiers.py index 316b798f9e..7fb34cbca7 100644 --- a/torchtitan/experiments/rl/tests/test_verifiers.py +++ b/torchtitan/experiments/rl/tests/test_verifiers.py @@ -32,10 +32,13 @@ def test_trainable_token_spans() -> None: 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)], @@ -65,6 +68,10 @@ def test_verifiers_trace_preserves_generation_metadata() -> None: 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