diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 18e57dd7af6..27354704e0d 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -10,6 +10,8 @@ import torch from openai import BadRequestError from parameterized import parameterized_class +from ray.util import remove_placement_group +from ray.util.placement_group import placement_group_table from transformers import AutoConfig, AutoTokenizer from tests.tools import ( @@ -21,10 +23,12 @@ get_template_config, get_vision_language_model_path, ) -from trinity.common.config import Config +from trinity.common.config import Config, GenerationConfig from trinity.common.constants import ROLLOUT_WEIGHT_SYNC_GROUP_NAME, SyncMethod from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.thinking_budget_workflow import ThinkingBudgetWorkflow +from trinity.common.workflows.workflow import Task from trinity.manager.synchronizer import Synchronizer DEBUG = False @@ -37,7 +41,29 @@ def print_debug(*args): async def create_test_models(config: Config): allocator = Allocator(config.explorer) - return await allocator.create_all_models() + try: + engines, auxiliary_engines = await allocator.create_all_models() + except BaseException: + if hasattr(allocator, "pg"): + remove_placement_group(allocator.pg) + raise + + for wrapper in engines: + setattr(wrapper, "_test_placement_group", allocator.pg) + setattr( + wrapper, + "_test_actor_names", + tuple(allocator.bundle_result.actor_bundle_map), + ) + for wrappers in auxiliary_engines: + for wrapper in wrappers: + setattr(wrapper, "_test_placement_group", allocator.pg) + setattr( + wrapper, + "_test_actor_names", + tuple(allocator.bundle_result.actor_bundle_map), + ) + return engines, auxiliary_engines def clone_wrapper(wrapper: ModelWrapper, enable_history: bool) -> ModelWrapper: @@ -108,8 +134,43 @@ async def asyncTearDown(self): for model_list in value: wrappers.extend(model_list) - if wrappers: + if not wrappers: + return + + placement_groups = { + getattr(wrapper, "_test_placement_group").id: getattr(wrapper, "_test_placement_group") + for wrapper in wrappers + if hasattr(wrapper, "_test_placement_group") + } + actors = [actor for wrapper in wrappers for actor in wrapper.models] + actor_names = { + name for wrapper in wrappers for name in getattr(wrapper, "_test_actor_names", ()) + } + namespace = self.config.explorer.rollout_model.ray_namespace + try: await asyncio.gather(*[wrapper.shutdown() for wrapper in wrappers]) + finally: + for actor in actors: + ray.kill(actor, no_restart=True) + for pg in placement_groups.values(): + remove_placement_group(pg) + for _ in range(100): + tables = [placement_group_table(pg) for pg in placement_groups.values()] + live_actor_names = set() + for name in actor_names: + try: + ray.get_actor(name, namespace=namespace) + live_actor_names.add(name) + except ValueError: + pass + if ( + all(not table or table.get("state") == "REMOVED" for table in tables) + and not live_actor_names + ): + break + await asyncio.sleep(0.1) + else: + self.fail("Timed out while removing vLLM test actors or placement groups.") @parameterized_class( @@ -733,6 +794,7 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.enable_openai_api = True self.config.explorer.rollout_model.enable_auto_tool_choice = True self.config.explorer.rollout_model.tool_call_parser = "qwen3_coder" + self.config.explorer.rollout_model.reasoning_parser = "qwen3" self.config.explorer.rollout_model.enable_history = True self.config.check_and_update() @@ -797,6 +859,53 @@ async def test_reasoning_content(self): text = self.tokenizer.decode(exps[0].tokens.tolist()) self.assertIn("Use `list_agents` tool to get the list of agents.", text) + async def test_thinking_token_budget_action_mask(self): + thinking_token_budget = 3 + messages = [{"role": "user", "content": "Which is larger, 9.11 or 9.8? Explain."}] + task = Task( + raw_task={"messages": messages}, + workflow_args={"thinking_token_budget": thinking_token_budget}, + rollout_args=GenerationConfig(n=1, max_tokens=32, temperature=0.0), + ) + + exps = ThinkingBudgetWorkflow(task=task, model=self.model_wrapper).run() + + self.assertEqual(len(exps), 1) + exp = exps[0] + response_token_ids = exp.tokens[exp.prompt_length :].tolist() + tokenizer = AutoTokenizer.from_pretrained(self.config.model.model_path) + reasoning_start_token_ids = tokenizer.encode("", add_special_tokens=False) + reasoning_end_token_ids = tokenizer.encode("", add_special_tokens=False) + self.assertGreater(len(reasoning_start_token_ids), 0) + self.assertGreater(len(reasoning_end_token_ids), 0) + + def find_subsequence(token_ids, subsequence): + return next( + ( + index + for index in range(len(token_ids) - len(subsequence) + 1) + if token_ids[index : index + len(subsequence)] == subsequence + ), + -1, + ) + + reasoning_start = find_subsequence(response_token_ids, reasoning_start_token_ids) + reasoning_end_start = find_subsequence(response_token_ids, reasoning_end_token_ids) + self.assertGreaterEqual(reasoning_end_start, 0) + reasoning_content_start = ( + reasoning_start + len(reasoning_start_token_ids) if reasoning_start >= 0 else 0 + ) + self.assertGreaterEqual(reasoning_end_start, reasoning_content_start) + self.assertLessEqual( + reasoning_end_start - reasoning_content_start, + thinking_token_budget, + ) + + reasoning_end_stop = reasoning_end_start + len(reasoning_end_token_ids) + self.assertTrue(torch.all(exp.action_mask[:reasoning_end_start])) + self.assertTrue(torch.all(~exp.action_mask[reasoning_end_start:reasoning_end_stop])) + self.assertTrue(torch.all(exp.action_mask[reasoning_end_stop:])) + class TestQwen35APIServerMultiModal(VLLMTestBase): async def asyncSetUp(self): diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index 7627cdca296..d862218397a 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -50,6 +50,7 @@ # on-policy distillation workflows "on_policy_distill_workflow": "trinity.common.workflows.on_policy_distill_workflow.OnPolicyDistillWorkflow", "on_policy_distill_math_workflow": "trinity.common.workflows.on_policy_distill_workflow.OnPolicyDistillMathWorkflow", + "thinking_budget_workflow": "trinity.common.workflows.thinking_budget_workflow.ThinkingBudgetWorkflow", # custom workflows "sudoku_workflow": "trinity.common.workflows.envs.sudoku.sudoku_workflow.SudokuWorkflow", }, diff --git a/trinity/common/workflows/thinking_budget_workflow.py b/trinity/common/workflows/thinking_budget_workflow.py new file mode 100644 index 00000000000..aa7f0cf0fc9 --- /dev/null +++ b/trinity/common/workflows/thinking_budget_workflow.py @@ -0,0 +1,113 @@ +"""OpenAI-compatible workflow with vLLM thinking-budget support.""" + +from __future__ import annotations + +from dataclasses import asdict +from typing import TYPE_CHECKING, List, Optional + +import torch +from transformers import AutoTokenizer + +from trinity.common.experience import Experience +from trinity.common.workflows.workflow import Task, Workflow + +if TYPE_CHECKING: + from trinity.common.models.model import ModelWrapper + + +class ThinkingBudgetWorkflow(Workflow): + """Run raw OpenAI messages with a per-request reasoning-token budget. + + vLLM forces the configured end-of-reasoning token sequence when the budget + is exhausted. The workflow locates that sequence in the returned completion + token IDs and masks it out so it does not contribute to the policy loss. + """ + + can_repeat = True + DEFAULT_REASONING_END_STR = "" + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, + ): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + raw_task = task.raw_task or {} + self.messages = raw_task.get("messages") + if not isinstance(self.messages, list): + raise ValueError("ThinkingBudgetWorkflow requires raw_task['messages'] to be a list.") + + self.thinking_token_budget = int(task.workflow_args.get("thinking_token_budget")) # type: ignore [arg-type] + if not isinstance(self.thinking_token_budget, int) or self.thinking_token_budget < 0: + raise ValueError( + "workflow_args['thinking_token_budget'] must be a non-negative integer." + ) + self.reasoning_end_str = task.workflow_args.get( + "reasoning_end_str", self.DEFAULT_REASONING_END_STR + ) + if not isinstance(self.reasoning_end_str, str) or not self.reasoning_end_str: + raise ValueError("workflow_args['reasoning_end_str'] must be a non-empty string.") + if not model.enable_history: + raise ValueError( + "ThinkingBudgetWorkflow requires explorer.rollout_model.enable_history=true." + ) + self.repeat_times = task.rollout_args.n + self.run_id_base = 0 + self.client = model.get_openai_client() + self.tokenizer = AutoTokenizer.from_pretrained( + model.model_path, + trust_remote_code=model.config.trust_remote_code, + ) + self.reasoning_end_token_ids = self.tokenizer.encode( + self.reasoning_end_str, add_special_tokens=False + ) + if not self.reasoning_end_token_ids: + raise ValueError("workflow_args['reasoning_end_str'] encodes to no tokens.") + + def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: + self.repeat_times = repeat_times + self.run_id_base = run_id_base + + @staticmethod + def _mask_forced_reasoning_end( + exp: Experience, + reasoning_end_token_ids: List[int], + ) -> None: + """Set the vLLM-forced reasoning-end token action entries to zero.""" + response_len = len(exp.tokens) - exp.prompt_length # type: ignore [arg-type] + if exp.action_mask is None or len(exp.action_mask) != response_len: + exp.action_mask = torch.ones(response_len, dtype=torch.bool) + response_token_ids = exp.tokens[exp.prompt_length :].tolist() # type: ignore [index] + for start in range(response_len - len(reasoning_end_token_ids) + 1): + end = start + len(reasoning_end_token_ids) + if response_token_ids[start:end] == reasoning_end_token_ids: + exp.action_mask[start:end] = False + return + + def run(self) -> List[Experience]: + rollout_args = { + key: value + for key, value in asdict(self.task.rollout_args).items() + if value is not None and key not in {"top_k", "logprobs", "n"} + } + rollout_args["n"] = self.repeat_times + extra_body = {"thinking_token_budget": self.thinking_token_budget} + if self.task.rollout_args.top_k >= 0: + extra_body["top_k"] = self.task.rollout_args.top_k + + self.client.chat.completions.create( + model=self.client.model_path, + messages=self.messages, + extra_body=extra_body, + **rollout_args, + ) + experiences = self.model.extract_experience_from_history() + for index, exp in enumerate(experiences): + self._mask_forced_reasoning_end( + exp, + reasoning_end_token_ids=self.reasoning_end_token_ids, + ) + exp.eid.run = self.run_id_base + index + return experiences