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>
This commit is contained in:
Datta Nimmaturi
GitHub
parent
556f396b3c
commit
8b2b99be03
3 changed files with 166 additions and 1 deletions
103
tests/test_tool_mask_zoo_compat.py
Normal file
103
tests/test_tool_mask_zoo_compat.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue