AsyncGRPOTrainer: optional SAO-style importance sampling mask (#13820)

* AsyncGRPOTrainer: optional SAO-style importance sampling mask

importance_sampling_mask = True on AsyncGRPOTrainer zeroes the gradient of every
token whose pi_theta / pi_rollout ratio falls outside (1 - epsilon, 1 + epsilon_high)
instead of PPO clipping it, as in SAO's direct double-sided importance sampling
(arXiv 2607.07508). Off by default; the class keeps TRL's own compute_loss.

* Give the installed-TRL test skip an explanatory reason

* Record the masked compute_loss rewrite in the exec-literal baseline

* Run CI on the converged head

scripts/exec_literals_baseline.json

@@ -289,6 +289,13 @@
       "count": 1,
       "reason": "Re-executes TRL's own prepare_peft_model with one branch condition corrected, matching the clause TRL itself added in 0.24.0. The source is read with inspect.getsource from the installed trl package, never from user input or a network fetch, and it is only rewritten when the exact expected branch text is present exactly once; otherwise the patch bails and leaves TRL alone. See tests/test_kbit_prep_guard.py."
     },
+    {
+      "file": "unsloth/models/rl.py",
+      "sink": "compile",
+      "digest": "4be3a8fe5e437a0c",
+      "count": 1,
+      "reason": "Compiles TRL's own AsyncGRPOTrainer.compute_loss, read with inspect.getsource from the installed trl package, with the PPO loss line swapped for the masked one only when it is present exactly once (two metric lines likewise); built only when a caller passes importance_sampling_mask=True. Never user input or a network fetch. See tests/test_async_grpo_patch.py."
+    },
     {
       "file": "unsloth/models/rl.py",
       "sink": "eval",
@@ -387,6 +394,13 @@
       "count": 1,
       "reason": "Re-imports TRL trainer and config classes by name, and installs the RL trainer source this file generates. The names come from this file's own tables."
     },
+    {
+      "file": "unsloth/models/rl.py",
+      "sink": "exec",
+      "digest": "dc0878fa52de4e4f",
+      "count": 1,
+      "reason": "Executes the code object compiled on the line above against TRL's own module globals into a fresh namespace, defining only the rewritten compute_loss. The source is the installed trl package's code, not user or network input."
+    },
     {
       "file": "unsloth/models/rl.py",
       "sink": "exec",

tests/test_async_grpo_patch.py

@@ -47,6 +47,7 @@ NAMES = (
     "_hold_peft_requires_grad",
     "_allow_moe_parameter_lora",
     "_check_async_grpo_adapter_servable",
+    "_async_grpo_masked_compute_loss",
     "patch_trl_async_grpo",
 )
 
@@ -60,6 +61,8 @@ def _load(fast_backbone_types = ()):
     src = open(RL_SOURCE_PATH, encoding = "utf-8").read()
     ns = {
         "copy": __import__("copy"),
+        "inspect": __import__("inspect"),
+        "textwrap": __import__("textwrap"),
         "types": types,
         "torch": torch,
         "functools": functools,
@@ -71,7 +74,10 @@ def _load(fast_backbone_types = ()):
         "importlib_version": __import__("importlib.metadata", fromlist = ["version"]).version,
     }
     for node in ast.parse(src).body:
-        if isinstance(node, ast.FunctionDef) and node.name in NAMES:
+        is_constant = isinstance(node, ast.Assign) and any(
+            isinstance(t, ast.Name) and t.id.startswith("_ASYNC_GRPO_") for t in node.targets
+        )
+        if is_constant or (isinstance(node, ast.FunctionDef) and node.name in NAMES):
             exec(compile(ast.Module([node], []), RL_SOURCE_PATH, "exec"), ns)
     assert all(n in ns for n in NAMES)
     ns["_is_unsloth_fast_backbone"] = lambda b: isinstance(b, fast_backbone_types)
@@ -659,3 +665,96 @@ def test_image_check_does_not_consume_an_iterable_dataset():
     assert next(data) == {"prompt": "a"}
     assert ns["_dataset_has_images"]([{"image": object()}]) is True
     assert ns["_dataset_has_images"]([{"prompt": "hi"}]) is False
+
+
+def _ppo_compute_loss(
+    self,
+    model,
+    inputs,
+    return_outputs = False,
+    num_items_in_batch = None,
+):
+    # The policy-loss lines of TRL 1.13 to 1.15 AsyncGRPOTrainer.compute_loss, verbatim.
+    log_probs = model(inputs)
+    advantages = inputs["advantages"]
+    log_ratio = log_probs - inputs["old_log_probs"]
+    coef_1 = torch.exp(log_ratio)
+    coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high)
+    per_token_loss1 = coef_1 * advantages
+    per_token_loss2 = coef_2 * advantages
+    per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
+    with torch.no_grad():
+        is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0)
+        is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0)
+        self.clipped = (is_low_clipped | is_high_clipped).tolist()
+    return per_token_loss.sum()
+
+
+def _loss_grad(trainer, ratio, advantages):
+    log_probs = torch.zeros(len(ratio), dtype = torch.float64, requires_grad = True)
+    inputs = {
+        "old_log_probs": -torch.tensor(ratio, dtype = torch.float64).log(),
+        "advantages": torch.tensor(advantages, dtype = torch.float64),
+    }
+    trainer.compute_loss(lambda _: log_probs, inputs).backward()
+    return log_probs.grad.tolist()
+
+
+def test_importance_sampling_mask_gives_the_sao_gradient(fake_trl):
+    mod, _ = fake_trl
+    mod.AsyncGRPOTrainer.compute_loss = _ppo_compute_loss
+    _load(fast_backbone_types = (FastBackbone,))["patch_trl_async_grpo"]()
+    ratio = [0.5, 0.7, 0.75, 1.0, 1.3, 5.9, 6.0, 9.0, 9.0, 0.5]
+    advantages = [1.0, 1.0, -2.0, 0.5, -1.0, 1.0, 1.0, 1.0, -1.0, -1.0]
+    masked = mod.AsyncGRPOTrainer(
+        Peft(CausalLM(FastBackbone())), None, importance_sampling_mask = True
+    )
+    masked.epsilon_low, masked.epsilon_high = 0.3, 5.0
+    # SAO Eq. 1-3: d/dlog pi of -f(r) A log pi, f(r) = r inside the open window (0.7, 6.0), else 0.
+    expected = [-(r * a if 0.7 < r < 6.0 else 0.0) for r, a in zip(ratio, advantages)]
+    assert _loss_grad(masked, ratio, advantages) == pytest.approx(expected)
+    assert masked.clipped == [True, True, False, False, False, False, True, True, True, True]
+    # Default stays TRL's PPO loss, on TRL's own method: a negative advantage at r = 9 still gets a gradient.
+    plain = mod.AsyncGRPOTrainer(Peft(CausalLM(FastBackbone())), None)
+    plain.epsilon_low, plain.epsilon_high = 0.3, 5.0
+    assert "compute_loss" not in vars(plain)
+    assert _loss_grad(plain, ratio, advantages)[8] == pytest.approx(9.0)
+    assert plain.clipped[8] is False
+
+
+def test_importance_sampling_mask_refuses_an_unknown_loss(fast_inference_trl):
+    mod, _ = fast_inference_trl
+
+    def compute_loss(
+        self,
+        model,
+        inputs,
+        return_outputs = False,
+        num_items_in_batch = None,
+    ):
+        return model(inputs).sum()
+
+    mod.AsyncGRPOTrainer.compute_loss = compute_loss
+    peft = _fast_inference_peft(CausalLM(FastBackbone()))
+    with pytest.raises(ValueError, match = "cannot be masked"):
+        mod.AsyncGRPOTrainer(
+            peft, None, args = mod.AsyncGRPOConfig("o"), importance_sampling_mask = True
+        )
+    # Refused before the in-process server starts or TRL's init runs.
+    assert _FakeServer.started == [] and mod._calls["init"] == []
+
+
+def test_importance_sampling_mask_rewrites_the_installed_trl_loss():
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            from trl.experimental.async_grpo.async_grpo_trainer import AsyncGRPOTrainer
+    except Exception as e:
+        pytest.skip(reason = f"needs TRL >= 1.13, which ships trl.experimental.async_grpo: {e}")
+    ns = _load()
+    masked = ns["_async_grpo_masked_compute_loss"](AsyncGRPOTrainer.compute_loss)
+    code = masked.__code__
+    assert masked is ns["_async_grpo_masked_compute_loss"](AsyncGRPOTrainer.compute_loss)
+    assert code is not AsyncGRPOTrainer.compute_loss.__code__
+    assert "where" in code.co_names and "min" not in code.co_names
+    assert masked.__globals__ is AsyncGRPOTrainer.compute_loss.__globals__

unsloth/models/rl.py

@@ -27,6 +27,7 @@ import inspect
 import os
 import re
 import sys
+import textwrap
 import warnings
 import weakref
 from contextlib import contextmanager
@@ -4688,6 +4689,46 @@ def _check_async_grpo_adapter_servable(model):
         model.load_lora(os.path.join(directory, "adapter"), load_tensors = True)
 
 
+_ASYNC_GRPO_PPO_LOSS = "per_token_loss = -torch.min(per_token_loss1, per_token_loss2)"
+# Direct double-sided IS (SAO, arXiv 2607.07508 Eq. 3): a token whose pi_theta / pi_rollout ratio leaves
+# (1 - epsilon, 1 + epsilon_high) gets no gradient at all, whatever the sign of its advantage.
+_ASYNC_GRPO_MASKED_LOSS = (
+    "per_token_loss = torch.where("
+    "(coef_1 > 1 - self.epsilon_low) & (coef_1 < 1 + self.epsilon_high), -per_token_loss1, 0.0)"
+)
+_ASYNC_GRPO_MASK_METRICS = (
+    (
+        "is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0)",
+        "is_low_clipped = coef_1 <= 1 - self.epsilon_low",
+    ),
+    (
+        "is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0)",
+        "is_high_clipped = coef_1 >= 1 + self.epsilon_high",
+    ),
+)
+
+
+@functools.lru_cache(maxsize = None)
+def _async_grpo_masked_compute_loss(compute_loss):
+    source = textwrap.dedent(inspect.getsource(compute_loss))
+    if source.count(_ASYNC_GRPO_PPO_LOSS) != 1:
+        raise ValueError(
+            "Unsloth: importance_sampling_mask = True could not find the PPO loss in "
+            f"{compute_loss.__qualname__}, so it cannot be masked. Drop importance_sampling_mask for this trainer."
+        )
+    source = source.replace(_ASYNC_GRPO_PPO_LOSS, _ASYNC_GRPO_MASKED_LOSS)
+    # The clip_ratio metrics then count masked tokens (either advantage sign) instead of PPO-clipped ones.
+    for old, new in _ASYNC_GRPO_MASK_METRICS:
+        source = source.replace(old, new)
+    namespace = {}
+    exec(
+        compile(source, f"<unsloth masked {compute_loss.__qualname__}>", "exec"),
+        compute_loss.__globals__,
+        namespace,
+    )
+    return functools.wraps(compute_loss)(namespace[compute_loss.__name__])
+
+
 def patch_trl_async_grpo():
     # Let TRL's AsyncGRPOTrainer take a loaded Unsloth model; keep packed rows from attending across sequences.
     if importlib.util.find_spec("trl.experimental") is None:
@@ -4730,7 +4771,18 @@ def patch_trl_async_grpo():
     original_train = getattr(trainer_class, "train", None)
 
     @functools.wraps(original_init)
-    def __init__(self, model, *args, **kwargs):
+    def __init__(
+        self,
+        model,
+        *args,
+        importance_sampling_mask = False,
+        **kwargs,
+    ):
+        masked_loss = (
+            _async_grpo_masked_compute_loss(type(self).compute_loss)
+            if importance_sampling_mask
+            else None
+        )
         config = args[1] if len(args) >= 2 else kwargs.get("args")
         if not isinstance(model, str) and config is None:
             name = getattr(getattr(model, "config", None), "_name_or_path", "") or "model"
@@ -4793,6 +4845,8 @@ def patch_trl_async_grpo():
         if server is not None:
             self._unsloth_vllm_server = server
             weakref.finalize(self, server.stop)
+        if masked_loss is not None:
+            self.compute_loss = types.MethodType(masked_loss, self)
 
     trainer_class.__init__ = __init__
     if original_train is not None: