From 0c295342bbc77c850653952baf4c6b9b1906a55f Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 16 Jul 2026 19:36:11 +0800 Subject: [PATCH 1/6] add thinking budget --- tests/common/vllm_test.py | 56 +++++++- trinity/common/workflows/__init__.py | 1 + .../workflows/thinking_budget_workflow.py | 124 ++++++++++++++++++ 3 files changed, 180 insertions(+), 1 deletion(-) create mode 100644 trinity/common/workflows/thinking_budget_workflow.py diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 18e57dd7af6..7f1fb038342 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -21,10 +21,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 @@ -733,6 +735,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 +800,57 @@ 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, + "reasoning_end_str": "", + }, + 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.logprobs[reasoning_end_start:reasoning_end_stop] == 0)) + 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..e9888dcc0b7 --- /dev/null +++ b/trinity/common/workflows/thinking_budget_workflow.py @@ -0,0 +1,124 @@ +"""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 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 enforces the end-of-reasoning sequence by forcing its logits. Those + tokens therefore have an exact log probability of zero. The workflow masks + the first such contiguous forced run at/after the budget boundary out of the + action mask so it does not contribute to the policy loss. + """ + + can_repeat = True + + 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") + if self.reasoning_end_str is not None and not isinstance(self.reasoning_end_str, str): + raise ValueError("workflow_args['reasoning_end_str'] must be a 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() + + 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, + budget: int, + reasoning_end_token_ids: Optional[List[int]] = None, + ) -> 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] + if reasoning_end_token_ids: + 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 + if exp.logprobs is None or len(exp.logprobs) != response_len: + return + + # Forced tokens receive probability 1 after vLLM's logits override, so + # their returned log probability is exactly 0. Search from the earliest + # possible budget boundary and mask only the first contiguous run. + forced = torch.eq(exp.logprobs, 0) + start = min(budget, response_len) + while start < response_len and not bool(forced[start]): + start += 1 + end = start + while end < response_len and bool(forced[end]): + end += 1 + if start < end: + exp.action_mask[start:end] = False + + 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() + reasoning_end_token_ids = None + if self.reasoning_end_str: + reasoning_end_token_ids = self.model.tokenizer.encode( + self.reasoning_end_str, add_special_tokens=False + ) + for index, exp in enumerate(experiences): + self._mask_forced_reasoning_end( + exp, + self.thinking_token_budget, + reasoning_end_token_ids=reasoning_end_token_ids, + ) + exp.eid.run = self.run_id_base + index + return experiences From a5401237260fd0931af57e5030a54b0dcbc2d5c3 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 16 Jul 2026 19:54:21 +0800 Subject: [PATCH 2/6] fix tests --- tests/common/vllm_test.py | 1 - .../workflows/thinking_budget_workflow.py | 58 +++++++------------ 2 files changed, 22 insertions(+), 37 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 7f1fb038342..9436510fd54 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -848,7 +848,6 @@ def find_subsequence(token_ids, subsequence): 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.logprobs[reasoning_end_start:reasoning_end_stop] == 0)) self.assertTrue(torch.all(exp.action_mask[reasoning_end_stop:])) diff --git a/trinity/common/workflows/thinking_budget_workflow.py b/trinity/common/workflows/thinking_budget_workflow.py index e9888dcc0b7..7606fdf1d8a 100644 --- a/trinity/common/workflows/thinking_budget_workflow.py +++ b/trinity/common/workflows/thinking_budget_workflow.py @@ -6,6 +6,7 @@ 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 @@ -17,10 +18,9 @@ class ThinkingBudgetWorkflow(Workflow): """Run raw OpenAI messages with a per-request reasoning-token budget. - vLLM enforces the end-of-reasoning sequence by forcing its logits. Those - tokens therefore have an exact log probability of zero. The workflow masks - the first such contiguous forced run at/after the budget boundary out of the - action mask so it does not contribute to the policy loss. + 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 @@ -44,8 +44,8 @@ def __init__( "workflow_args['thinking_token_budget'] must be a non-negative integer." ) self.reasoning_end_str = task.workflow_args.get("reasoning_end_str") - if self.reasoning_end_str is not None and not isinstance(self.reasoning_end_str, str): - raise ValueError("workflow_args['reasoning_end_str'] must be a string.") + 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." @@ -53,6 +53,15 @@ def __init__( 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 @@ -61,35 +70,18 @@ def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: @staticmethod def _mask_forced_reasoning_end( exp: Experience, - budget: int, - reasoning_end_token_ids: Optional[List[int]] = None, + 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] - if reasoning_end_token_ids: - 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 - if exp.logprobs is None or len(exp.logprobs) != response_len: - return - - # Forced tokens receive probability 1 after vLLM's logits override, so - # their returned log probability is exactly 0. Search from the earliest - # possible budget boundary and mask only the first contiguous run. - forced = torch.eq(exp.logprobs, 0) - start = min(budget, response_len) - while start < response_len and not bool(forced[start]): - start += 1 - end = start - while end < response_len and bool(forced[end]): - end += 1 - if start < end: - exp.action_mask[start:end] = False + 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 = { @@ -109,16 +101,10 @@ def run(self) -> List[Experience]: **rollout_args, ) experiences = self.model.extract_experience_from_history() - reasoning_end_token_ids = None - if self.reasoning_end_str: - reasoning_end_token_ids = self.model.tokenizer.encode( - self.reasoning_end_str, add_special_tokens=False - ) for index, exp in enumerate(experiences): self._mask_forced_reasoning_end( exp, - self.thinking_token_budget, - reasoning_end_token_ids=reasoning_end_token_ids, + reasoning_end_token_ids=self.reasoning_end_token_ids, ) exp.eid.run = self.run_id_base + index return experiences From d6016e3e0dfeada95bcb0820ae08b50ef91c4ba3 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 16 Jul 2026 20:09:17 +0800 Subject: [PATCH 3/6] update tests --- tests/common/vllm_test.py | 5 +---- trinity/common/workflows/thinking_budget_workflow.py | 5 ++++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 9436510fd54..1691774858c 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -805,10 +805,7 @@ async def test_thinking_token_budget_action_mask(self): 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, - "reasoning_end_str": "", - }, + workflow_args={"thinking_token_budget": thinking_token_budget}, rollout_args=GenerationConfig(n=1, max_tokens=32, temperature=0.0), ) diff --git a/trinity/common/workflows/thinking_budget_workflow.py b/trinity/common/workflows/thinking_budget_workflow.py index 7606fdf1d8a..aa7f0cf0fc9 100644 --- a/trinity/common/workflows/thinking_budget_workflow.py +++ b/trinity/common/workflows/thinking_budget_workflow.py @@ -24,6 +24,7 @@ class ThinkingBudgetWorkflow(Workflow): """ can_repeat = True + DEFAULT_REASONING_END_STR = "" def __init__( self, @@ -43,7 +44,9 @@ def __init__( 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.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: From ac83ecb9d861e507f031681d8ffb4d974d43405b Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 16 Jul 2026 21:17:31 +0800 Subject: [PATCH 4/6] fix tests --- tests/common/vllm_test.py | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 1691774858c..6c94a62605e 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 ( @@ -39,7 +41,19 @@ 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) + for wrappers in auxiliary_engines: + for wrapper in wrappers: + setattr(wrapper, "_test_placement_group", allocator.pg) + return engines, auxiliary_engines def clone_wrapper(wrapper: ModelWrapper, enable_history: bool) -> ModelWrapper: @@ -110,8 +124,26 @@ 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") + } + try: await asyncio.gather(*[wrapper.shutdown() for wrapper in wrappers]) + finally: + 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()] + if all(not table or table.get("state") == "REMOVED" for table in tables): + break + await asyncio.sleep(0.1) + else: + self.fail("Timed out while removing vLLM test placement groups.") @parameterized_class( From ff000446283478444eeb30b6821a509c9d426b37 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 17 Jul 2026 10:04:45 +0800 Subject: [PATCH 5/6] clean env --- tests/common/vllm_test.py | 33 +++++++++++++++++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 6c94a62605e..da279d1d97f 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -50,9 +50,19 @@ async def create_test_models(config: Config): 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 @@ -132,18 +142,37 @@ async def asyncTearDown(self): 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()] - if all(not table or table.get("state") == "REMOVED" for table in tables): + 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 placement groups.") + self.fail("Timed out while removing vLLM test actors or placement groups.") @parameterized_class( From 4a25003fa006482240d30c62ec8dd2364f766605 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 17 Jul 2026 10:05:06 +0800 Subject: [PATCH 6/6] fix pre-commit --- tests/common/vllm_test.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index da279d1d97f..27354704e0d 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -144,9 +144,7 @@ async def asyncTearDown(self): } 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", ()) + name for wrapper in wrappers for name in getattr(wrapper, "_test_actor_names", ()) } namespace = self.config.explorer.rollout_model.ray_namespace try: