From 8b2b99be036dc114700867392f1b6a149ee3fdaf Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 27 May 2026 16:50:44 +0530 Subject: [PATCH] tool mask support (#5682) * tool mask support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle tool masks with older zoo builds * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep tool mask implementation in zoo --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_tool_mask_zoo_compat.py | 103 +++++++++++++++++++++++++++++ unsloth/models/rl.py | 19 ++++++ unsloth/models/rl_replacements.py | 45 ++++++++++++- 3 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 tests/test_tool_mask_zoo_compat.py diff --git a/tests/test_tool_mask_zoo_compat.py b/tests/test_tool_mask_zoo_compat.py new file mode 100644 index 0000000000..b6b68a7561 --- /dev/null +++ b/tests/test_tool_mask_zoo_compat.py @@ -0,0 +1,103 @@ +"""Compatibility checks for env/tool mask support with older unsloth_zoo.""" + +from __future__ import annotations + +import ast +import os +import textwrap + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)) +RL_SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl.py") +RL_REPLACEMENTS_SOURCE_PATH = os.path.join( + REPO_ROOT, "unsloth", "models", "rl_replacements.py" +) + + +def _read(path: str) -> str: + with open(path, "r") as fh: + return fh.read() + + +def _load_local_align_completion_tool_mask(): + src = _read(RL_SOURCE_PATH) + tree = ast.parse(src) + for node in tree.body: + if isinstance(node, ast.If): + for item in node.body: + if ( + isinstance(item, ast.FunctionDef) + and item.name == "align_completion_tool_mask" + ): + function_src = ast.get_source_segment(src, item) + break + else: + continue + break + else: + raise AssertionError("local align_completion_tool_mask fallback is missing") + + calls = [] + + def align_logprobs_with_mask(logprob_tensor, completion_mask, pad_value = None): + calls.append((logprob_tensor, completion_mask, pad_value)) + return torch.tensor( + [[1, 0, 1], [0, 1, 1]], + device = completion_mask.device, + dtype = logprob_tensor.dtype, + ) + + namespace = { + "torch": torch, + "align_logprobs_with_mask": align_logprobs_with_mask, + } + exec(textwrap.dedent(function_src), namespace) + return namespace["align_completion_tool_mask"], calls + + +def test_rl_uses_optional_zoo_tool_mask_helper(): + src = _read(RL_SOURCE_PATH) + assert 'RL_REPLACEMENTS.get("align_completion_tool_mask")' in src + assert 'RL_REPLACEMENTS["align_completion_tool_mask"]' not in src + + +def test_local_tool_mask_fallback_is_only_old_zoo_compat_shim(): + align_completion_tool_mask, calls = _load_local_align_completion_tool_mask() + completion_mask = torch.tensor( + [[1, 1, 0], [1, 1, 1]], + dtype = torch.float32, + ) + + assert align_completion_tool_mask(None, completion_mask) is completion_mask + assert calls == [] + + same_shape_tool_mask = torch.tensor([[1, 0, 1], [0, 1, 1]], dtype = torch.bool) + with pytest.raises(RuntimeError, match = "Please upgrade unsloth_zoo"): + align_completion_tool_mask(same_shape_tool_mask, completion_mask) + + +def test_grpo_accumulated_loss_omits_none_tool_mask_for_old_zoo(): + src = _read(RL_REPLACEMENTS_SOURCE_PATH) + assert "_grpo_accumulated_loss_kwargs = {}" in src + assert ( + 'if tool_mask is not None:\n _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask' + in src + ) + assert src.count("**_grpo_accumulated_loss_kwargs") == 2 + + accelerated_loss_start = src.find('if hasattr(self.args, "loss_type"):') + assert accelerated_loss_start != -1 + accelerated_loss_body = src[ + accelerated_loss_start : src.find( + 'if "train" in self._metrics:', accelerated_loss_start + ) + ] + assert "tool_mask = tool_mask" not in accelerated_loss_body + + +def test_rollout_output_patch_requires_real_tool_mask_symbol(): + src = _read(RL_REPLACEMENTS_SOURCE_PATH) + assert 're.search(r"\\btool_mask\\b", function)' in src + assert 'output["tool_mask"]' in src diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index c82c8364b3..9bf0d3e968 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -365,6 +365,22 @@ calculate_pad_tokens_in_prompt = RL_REPLACEMENTS["calculate_pad_tokens_in_prompt create_completion_attention_mask = RL_REPLACEMENTS["create_completion_attention_mask"] left_pack_padding = RL_REPLACEMENTS["left_pack_padding"] align_logprobs_with_mask = RL_REPLACEMENTS["align_logprobs_with_mask"] +align_completion_tool_mask = RL_REPLACEMENTS.get("align_completion_tool_mask") +if align_completion_tool_mask is None: + + def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, + ) -> torch.Tensor: + if tool_mask is None: + return completion_mask + raise RuntimeError( + "env_mask/tool_mask GRPO requires an unsloth_zoo build whose " + "grpo_accumulated_loss handles tool_mask. Please upgrade " + "unsloth_zoo." + ) + + autotune_batch_and_chunks = RL_REPLACEMENTS["grpo_autotune_batch_and_chunks"] sanitize_logprob = RL_REPLACEMENTS["sanitize_logprob"] @@ -452,6 +468,7 @@ torch_compile_options = {{ {create_completion_attention_mask_code} {left_pack_padding_code} {align_logprobs_with_mask_code} +{align_completion_tool_mask_code} {autotune_batch_and_chunks_code} {sanitize_logprob_code} @@ -1577,6 +1594,7 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): ) left_pack_padding_code = inspect.getsource(left_pack_padding) align_logprobs_with_mask_code = inspect.getsource(align_logprobs_with_mask) + align_completion_tool_mask_code = inspect.getsource(align_completion_tool_mask) autotune_batch_and_chunks_code = inspect.getsource(autotune_batch_and_chunks) sanitize_logprob_code = inspect.getsource(sanitize_logprob) # Get final source code @@ -1607,6 +1625,7 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): autotune_batch_and_chunks_code = autotune_batch_and_chunks_code, left_pack_padding_code = left_pack_padding_code, align_logprobs_with_mask_code = align_logprobs_with_mask_code, + align_completion_tool_mask_code = align_completion_tool_mask_code, sanitize_logprob_code = sanitize_logprob_code, ) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 0f9a324d5b..9ddb5e9453 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -946,6 +946,14 @@ def grpo_trainer__generate_and_score_completions(function_name, function): ) function = function.replace(_save_search, _save_replace) + if re.search(r"\btool_mask\b", function) and 'output["tool_mask"]' not in function: + function = function.replace( + " return output", + " if tool_mask is not None:\n" + ' output["tool_mask"] = tool_mask\n' + " return output", + ) + return function @@ -1523,6 +1531,7 @@ def grpo_trainer_compute_loss(function_name, function): mm_token_type_ids = inputs.get("mm_token_type_ids", None) num_items_in_batch = inputs.get("num_items_in_batch", None) sampling_per_token_logps = inputs.get("sampling_per_token_logps", None) + tool_mask = inputs.get("tool_mask", None) current_gradient_accumulation_steps = self.current_gradient_accumulation_steps num_processes = self.accelerator.num_processes @@ -1598,6 +1607,16 @@ def grpo_trainer_compute_loss(function_name, function): max_left_pad = inputs.get("max_left_pad", 0) if per_token_logps is not None: + loss_mask = completion_mask + if tool_mask is not None: + if tool_mask.shape != completion_mask.shape: + raise ValueError( + "tool_mask/env_mask must have the same shape as completion_mask" + ) + loss_mask = completion_mask * tool_mask.to( + device = completion_mask.device, + dtype = completion_mask.dtype, + ) ( loss, completion_length, @@ -1612,7 +1631,7 @@ def grpo_trainer_compute_loss(function_name, function): old_logps, sampling_per_token_logps, input_ids, - completion_mask, + loss_mask, self.beta, advantages, pixel_values = pixel_values, @@ -1662,6 +1681,28 @@ def grpo_trainer_compute_loss(function_name, function): "unsloth_zoo (see https://github.com/unslothai/unsloth-zoo/pull/613)." ) self._unsloth_grpo_zoo_checked = True + if tool_mask is not None and not getattr( + self, "_unsloth_grpo_tool_mask_zoo_checked", False + ): + _supports_tool_mask = ( + "tool_mask" in inspect.signature(grpo_accumulated_loss).parameters + ) + if not _supports_tool_mask: + try: + _zoo_src = inspect.getsource(grpo_accumulated_loss) + except (TypeError, OSError): + _zoo_src = "" + _supports_tool_mask = "tool_mask" in _zoo_src + if not _supports_tool_mask: + raise RuntimeError( + "env_mask/tool_mask GRPO requires an unsloth_zoo build whose " + "grpo_accumulated_loss handles tool_mask. Please upgrade " + "unsloth_zoo." + ) + self._unsloth_grpo_tool_mask_zoo_checked = True + _grpo_accumulated_loss_kwargs = {} + if tool_mask is not None: + _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask if hasattr(self.args, "loss_type"): ( loss, @@ -1703,6 +1744,7 @@ def grpo_trainer_compute_loss(function_name, function): sampling_per_token_logps = sampling_per_token_logps, token_type_ids = token_type_ids, mm_token_type_ids = mm_token_type_ids, + **_grpo_accumulated_loss_kwargs, ) else: # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17 @@ -1728,6 +1770,7 @@ def grpo_trainer_compute_loss(function_name, function): attention_mask = attention_mask, token_type_ids = token_type_ids, mm_token_type_ids = mm_token_type_ids, + **_grpo_accumulated_loss_kwargs, ) ) if "train" in self._metrics: