ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

(GRPO) Fix PEFT replacement for TRL >= 1.7.0, add missing compute_aux_loss for TRL >= 1.7.0 (#6904)

Browse source 
* Fix PEFT replacement for TRL >= 1.7.0, add missing compute_aux_loss for TRL 1.7.0

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix GRPO for TRL >= 1.7.0: PEFT ref-adapter removal and return arity

rl.py: for trl >= 1.7.0, scope the PEFT removal regex to the ref-adapter
block only by anchoring the end on ref_param.data.copy_(param.data), so it
no longer also deletes the following gradient-checkpointing
enable_input_require_grads() block. Neutralize TRL 1.7.0's
`if _is_quantized_model:` bf16 cast the same way the existing
is_loaded_in_4bit cast is handled.

rl_replacements.py: initialize _extra_moe_kwargs before use (it was
referenced before assignment whenever compute_aux_loss was passed) and only
request output_router_logits when the aux loss is actually wanted.

rl_replacements.py: _get_per_token_logps_and_entropies now returns a 3-tuple
(logps, entropies, aux_loss) for trl >= 1.7.0 and a 2-tuple for older TRL,
matching how every TRL call site unpacks the result. Without this, TRL 1.7.x
_generate_and_score_completions unpacks 3 values from a 2-tuple and raises
"not enough values to unpack (expected 3, got 2)".

* Return zero aux_loss placeholder and drop inference-mode aux collection

* GRPO TRL >= 1.7.0: reject router aux-loss opt-in at init; drop zero aux placeholder

Unsloth's optimized GRPO forward cannot compute the MoE router auxiliary loss.
Previously an explicit opt-in (router_aux_loss_coef > 0) returned a fabricated
zero, silently training without the requested load-balancing penalty. Now reject
it at trainer init with a clear NotImplementedError, and return None (not zero)
for the aux slot of TRL's 3-tuple. Default stays off (coef 0), so the common
path is unaffected.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* GRPO hidden-states fallback: free ModelOutput before chunked log-softmax

The old/ref logprob fallback binds the full ModelOutput (which holds every
layer's hidden_states when output_hidden_states=True) and kept it alive across
chunked_hidden_states_selective_log_softmax, an avoidable OOM on large models.
Extract logits then del outputs in both the text and VLM branches.

* Version-compat CI: proactively catch TRL GRPO breakage

The existing TRL canary is a static symbol/source grep: it verifies symbols
exist but is blind to structural changes (TRL 1.7.0's 2->3-tuple per-token-logps
return arity and restructured PEFT ref-adapter block, which the fix in this PR
addresses, both slipped past it because the methods still existed).

Two additions:
- test_trl_grpo_pinned_symbols.py: extend TRL_TAGS to 1.5/1.6/1.7 and pin the
  exact source-string contracts the rl.py / rl_replacements.py transforms depend
  on for TRL >= 1.7.0 (PEFT elif ref-adapter block + enable_input_require_grads
  survival, if _is_quantized_model, aux_loss_enabled anchor, compute_aux_loss
  arity). A future TRL change fails on main a few days before the PyPI release.
- test_trl_grpo_fake_run.py + a version-compat-ci job: fake-CUDA run that drives
  the real GRPO/SFT/DPO source-transform patchers against latest + main TRL on a
  CPU-only runner (no training) and asserts the generated Unsloth trainer still
  satisfies the transform contracts. Catches behavioral regressions the grep
  cannot see.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fake-run test: use a normal Version import for the aux gate

* version-compat CI: fix fake-run job gate + torch-absent collection

- Drop the invalid job-level matrix if (matrix is not available in
  jobs.<id>.if -> 'Unrecognized named-value: matrix' fails the whole
  workflow). Use a single job that runs vs TRL latest always and re-runs
  vs TRL main only on schedule/dispatch via a step-level github.event_name
  guard. Validated with actionlint.
- Module-level skip the fake-run test when torch is absent so
  daily-fresh-fetch (pytest-only, collects tests/version_compat/) does not
  crash on the top-level spoof import.

* fake-run test: do not skip on import failure

unsloth/trl are installed in the grpo-fake-run job, so a failing import is the
import-time drift this canary must catch. Keep only the not-installed find_spec
skips; let a real import error fail the test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* GRPO arity gate: regex downgrade + fail loud + CI coverage

The TRL < 1.7.0 per-token-logps return downgrade was an exact-string replace
anchored on the full return line incl. its comment, so a reformat (e.g.
pre-commit) could silently no-op it and ship a 3-tuple to older TRL. Switch to
a regex tolerant of comment/whitespace drift, and raise if the anchor stops
matching (re.subn count != 1) instead of failing silently. Add a monkeypatched
trl_version unit test asserting both arities, since CI only installs TRL >= 1.7.0
and never exercised the downgrade otherwise.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fake-run: give SFT/DPO a real contract, not just ast-parse

The SFT/DPO fake patch runs only checked the generated trainer parses. Also
assert the shared QLoRA _is_quantized_model bf16 cast is neutralized (TRL 1.7's
spelling, present in both sft_trainer and dpo_trainer), so a structural TRL
change to that block is caught for SFT/DPO too, not just GRPO.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* GRPO PEFT ref-adapter removal: lower gate to the TRL 1.4.0 floor

The elif is_peft_model(model) and args.beta != 0.0: ref-adapter block was
introduced in TRL 1.4.0 and is unchanged through 1.7.x, but the removal was
gated at >= 1.7.0, so for 1.4 <= TRL < 1.7 the transform fell through to the
0.27 branch (which matches the older if is_peft_available()... form) and
silently no-oped: a PEFT + beta != 0 GRPO run then computed the KL reference
from the copied ref adapter instead of the base model. Lower the gate to
1.4.0 and keep the 1.7.0-only router aux-loss fail-fast nested. Widen the
pinned-symbol contract test to run from 1.4.0 so the covered versions are
actually exercised.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
marcandrelarochelle 2026-07-08 07:05:03 -04:00 • committed by GitHub
commit 07c8bbbf5a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 504 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

246
tests/version_compat/test_trl_grpo_fake_run.py Normal file
View file

@ -0,0 +1,246 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""Fake-CUDA GRPO patch run against the *installed* TRL (CPU-only, no training).
The static symbol/source-string canaries (test_trl_grpo_pinned_symbols.py)
grep raw TRL source; they never execute unsloth's transforms. This test drives
the real pipeline: under the aggressive CUDA spoof it imports unsloth and calls
`_patch_trl_rl_trainers_impl`, which reads the installed GRPOTrainer via
inspect.getsource, applies every rl.py/rl_replacements.py rewrite, and compiles
the result into an UnslothGRPOTrainer. A structural TRL change that slips past
the greps (e.g. TRL 1.7.0's 2->3-tuple return arity, or a restructured PEFT
ref-adapter block) surfaces here as a transform error, a broken generated
source, or a violated contract -- with no GPU and no training run.
Meant to run in CI against `trl==latest` and `trl @ main` (see
version-compat-ci.yml). The tests/conftest.py harness pre-loads device_type
with DEVICE_COUNT=0 so unsloth's kernel init takes the CPU-safe path.
"""
from __future__ import annotations
import ast
import importlib
import importlib.machinery
import importlib.util
import inspect
import sys
import types
from pathlib import Path
import pytest
# daily-fresh-fetch collects tests/version_compat/ with only pytest installed;
# the spoof and the rest of this module need the real torch runtime. Skip the
# whole module cleanly when torch is absent rather than crashing collection.
if importlib.util.find_spec("torch") is None:
pytest.skip("torch not installed; fake-run needs the real runtime", allow_module_level = True)
# Apply the spoof BEFORE any unsloth-touching import (mirrors
# tests/vllm_compat/test_extended_module_imports.py).
_SPOOF_DIR = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_SPOOF_DIR))
import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402
_spoof.apply()
def _stub_module(name: str, attrs: dict | None = None) -> None:
if name in sys.modules:
return
m = types.ModuleType(name)
m.__spec__ = importlib.machinery.ModuleSpec(name = name, loader = None, origin = "<test stub>")
for k, v in (attrs or {}).items():
setattr(m, k, v)
sys.modules[name] = m
_stub_module("torchcodec")
def _trl_version():
import trl
from packaging.version import Version
return Version(trl.__version__.split("+")[0])
def _patch_grpo_and_get_source() -> str:
"""Run the GRPO patcher against the installed TRL and return the generated
UnslothGRPOTrainer source. Calls the impl (not the try/except wrapper) so a
transform/compile regression surfaces as a hard error instead of a silent
no-op."""
import trl.trainer.grpo_trainer as _g
from unsloth.models import rl as _rl
_rl._patch_trl_rl_trainers_impl("grpo_trainer")
patched = _g.GRPOTrainer
assert patched.__name__ == "UnslothGRPOTrainer", (
f"GRPO patch silently no-oped: trl.trainer.grpo_trainer.GRPOTrainer is "
f"{patched.__name__!r}, expected 'UnslothGRPOTrainer' (transform failed "
f"or dispatch key drifted on this TRL)"
)
# The transformed body (__init__ rewrites, injected per-token-logps) lives in
# the generated module's `_UnslothGRPOTrainer` base + module-level funcs, not
# the thin UnslothGRPOTrainer subclass -- read the whole generated module.
mod = inspect.getmodule(patched)
return inspect.getsource(mod) if mod is not None else inspect.getsource(patched)
@pytest.fixture(scope = "module")
def generated_grpo_source():
if importlib.util.find_spec("unsloth") is None:
pytest.skip("unsloth not installed")
if importlib.util.find_spec("trl") is None:
pytest.skip("trl not installed")
# Do NOT swallow import errors: unsloth is installed here, so a failing
# `import unsloth` is exactly the import-time TRL/transformers drift this
# canary must surface as a failure, not a skip.
import unsloth # noqa: F401 -- _gpu_init bootstrap under spoof
return _patch_grpo_and_get_source()
def test_grpo_patch_generates_valid_source(generated_grpo_source):
"""The generated UnslothGRPOTrainer must be syntactically valid Python."""
ast.parse(generated_grpo_source)
def test_grpo_patch_aux_fail_fast_injected(generated_grpo_source):
"""TRL >= 1.7.0: rl.py injects a fail-fast for the unsupported MoE router
aux-loss opt-in right after `self.aux_loss_enabled = ...`."""
from packaging.version import Version
if _trl_version() < Version("1.7.0"):
pytest.skip("aux_loss_enabled / router_aux_loss_coef are TRL >= 1.7.0")
assert "does not compute the MoE router auxiliary loss" in generated_grpo_source, (
"aux fail-fast raise missing from generated trainer; rl.py's "
"aux_loss_enabled .replace() anchor did not match this TRL"
)
def test_grpo_patch_three_tuple_return(generated_grpo_source):
"""TRL >= 1.7.0 call sites unpack a 3-tuple from
_get_per_token_logps_and_entropies; the injected replacement must return
(logps, entropies, aux_loss)."""
from packaging.version import Version
if _trl_version() >= Version("1.7.0"):
assert "return logprobs.detach(), entropies, aux_loss" in generated_grpo_source, (
"3-tuple per-token-logps return missing; the arity version-gate in "
"rl_replacements.py did not emit the >=1.7.0 form"
)
else:
assert (
"return logprobs.detach(), entropies, aux_loss" not in generated_grpo_source
), "2-tuple TRL got the 3-tuple return; arity gate mis-fired"
def test_grpo_patch_preserves_grad_checkpointing_block(generated_grpo_source):
"""The tightened PR #6904 PEFT regex must remove only the ref-adapter init,
not the following enable_input_require_grads gradient-checkpointing block."""
from packaging.version import Version
if _trl_version() < Version("1.7.0"):
pytest.skip("ref-adapter elif block is the TRL >= 1.7.0 shape")
assert "enable_input_require_grads" in generated_grpo_source, (
"gradient-checkpointing enable_input_require_grads() block was swallowed "
"by the PEFT-removal regex (over-reach regression)"
)
def test_grpo_patch_neutralizes_ref_adapter_and_qlora_cast(generated_grpo_source):
"""TRL >= 1.7.0: the ref-adapter copy and the hardcoded QLoRA bf16 cast must
both be gone from the generated trainer."""
from packaging.version import Version
if _trl_version() < Version("1.7.0"):
pytest.skip("targets the TRL >= 1.7.0 PEFT / _is_quantized_model shapes")
assert (
"ref_param.data.copy_(param.data)" not in generated_grpo_source
), "TRL's PEFT ref-adapter init survived; rl.py peft_pattern re.sub no-oped"
assert (
"if _is_quantized_model:" not in generated_grpo_source
), "TRL's hardcoded QLoRA bf16 cast survived; rl.py neutralization no-oped"
# SFT / DPO: the same source-transform patcher runs on them (a fake patch run,
# no training), so a structural TRL change can break generation. Assert the patch
# produces a valid, importable Unsloth trainer AND that the shared QLoRA
# `_is_quantized_model` bf16 cast is neutralized (TRL 1.7's spelling), which the
# patcher applies to every trainer. Catches "and or others" beyond GRPO.
def _patch_and_get_source(trainer_file: str, trainer_cls: str) -> str:
if importlib.util.find_spec("unsloth") is None or importlib.util.find_spec("trl") is None:
pytest.skip("unsloth or trl not installed")
# Let a real import failure fail the test (import-time drift is the target).
import unsloth # noqa: F401
import trl.trainer # noqa: F401
from unsloth.models import rl as _rl
_rl._patch_trl_rl_trainers_impl(trainer_file)
mod = importlib.import_module(f"trl.trainer.{trainer_file}")
patched = getattr(mod, trainer_cls)
assert patched.__name__ == f"Unsloth{trainer_cls}", (
f"{trainer_cls} patch silently no-oped on this TRL "
f"(got {patched.__name__!r}); source-transform dispatch drifted"
)
gen = inspect.getmodule(patched)
src = inspect.getsource(gen) if gen is not None else inspect.getsource(patched)
ast.parse(src)
return src
def _assert_quantized_cast_neutralized(src: str, trainer_cls: str) -> None:
from packaging.version import Version
if _trl_version() < Version("1.7.0"):
pytest.skip("pre-1.7.0 spells the QLoRA cast differently (is_loaded_in_4bit)")
assert "if _is_quantized_model:" not in src, (
f"{trainer_cls}: TRL's hardcoded QLoRA bf16 cast survived; the shared "
f"rl.py `if _is_quantized_model:` -> `if False:` neutralization no-oped"
)
def test_sft_patch_generates_valid_source():
src = _patch_and_get_source("sft_trainer", "SFTTrainer")
_assert_quantized_cast_neutralized(src, "SFTTrainer")
def test_dpo_patch_generates_valid_source():
src = _patch_and_get_source("dpo_trainer", "DPOTrainer")
_assert_quantized_cast_neutralized(src, "DPOTrainer")
# The installed TRL in CI is always >= 1.7.0, so the < 1.7.0 return-arity
# downgrade is never exercised by the fake-run above. Lock both arities by
# monkeypatching rl_replacements.trl_version and re-generating the injected
# _get_per_token_logps_and_entropies source directly (no TRL install needed).
def test_per_token_logps_arity_gate_both_directions(monkeypatch):
if importlib.util.find_spec("unsloth") is None:
pytest.skip("unsloth not installed")
import unsloth # noqa: F401
from packaging.version import Version
from unsloth.models import rl_replacements as _rlr
gate = _rlr.grpo_trainer__get_per_token_logps_and_entropies
# >= 1.7.0: 3-tuple return kept.
monkeypatch.setattr(_rlr, "trl_version", Version("1.7.0"), raising = False)
src_new = gate("_get_per_token_logps_and_entropies", None)
assert (
"return logprobs.detach(), entropies, aux_loss" in src_new
), "3-tuple return missing for TRL >= 1.7.0"
# < 1.7.0: aux_loss element dropped -> 2-tuple. A no-op downgrade must raise
# (fail loud), never silently ship a 3-tuple to older TRL.
monkeypatch.setattr(_rlr, "trl_version", Version("1.6.0"), raising = False)
src_old = gate("_get_per_token_logps_and_entropies", None)
assert (
"return logprobs.detach(), entropies # logps, entropies" in src_old
), "2-tuple return missing for TRL < 1.7.0"
assert (
"entropies, aux_loss" not in src_old
), "aux_loss element still present in the TRL < 1.7.0 downgrade"

110
tests/version_compat/test_trl_grpo_pinned_symbols.py
View file

@ -51,10 +51,27 @@ TRL_TAGS = [
"v1.2.0",
"v1.3.0",
"v1.4.0",
"v1.5.0",
"v1.5.1",
"v1.6.0",
"v1.7.0", # anchor: first release unsloth's TRL>=1.7.0 GRPO patch targets
"v1.7.1", # current PyPI latest
"main",
]
def _tag_ge(tag: str, floor: str) -> bool:
"""True if `tag` is `main` or a version >= `floor` (e.g. "1.7.0")."""
if tag == "main":
return True
from packaging.version import Version
try:
return Version(tag.lstrip("v")) >= Version(floor)
except Exception:
return False
# unsloth/trainer.py + unsloth/models/rl.py rebind these top-level names.
@ -537,3 +554,96 @@ def test_trl_truncate_with_protected_tokens_optional(tag: str):
assert src is not None
has_it = "truncate_with_protected_tokens" in src
_ = has_it # informational; pass either way.
# 24-27. TRL >= 1.7.0 GRPO source contracts. Unlike the has_def existence
# checks above, these pin the exact source strings unsloth/models/rl.py and
# rl_replacements.py transform for TRL >= 1.7.0 (the window PR #6904 fixes).
# The 1.7.0 break was invisible to the existence checks because the methods
# still existed -- only their internal structure / return arity changed. If
# TRL restructures one of these, the transform silently no-ops (or the
# generated trainer breaks), so failing here on `main` gives a few-day lead.
@pytest.mark.parametrize("tag", TRL_TAGS)
def test_trl_grpo_peft_ref_adapter_block_contract(tag: str):
"""rl.py (trl>=1.4.0) strips TRL's PEFT ref-adapter init with a re.DOTALL
regex anchored on `elif is_peft_model(model) and args.beta != 0.0:` ...
`ref_param.data.copy_(param.data)`. Both anchors must exist (else the
regex no-ops and the ref adapter is created under Unsloth), and the
following `enable_input_require_grads` gradient-checkpointing block must
remain present -- the tightened regex must NOT swallow it (PR #6904). The
`elif` block shape appeared in TRL 1.4.0, so this contract runs from there."""
if not _tag_ge(tag, "1.4.0"):
pytest.skip(
f"{tag}: pre-1.4.0 uses the `if is_peft_available()...` form (rl.py 0.27 branch)"
)
src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
assert src is not None
assert "elif is_peft_model(model) and args.beta != 0.0:" in src, (
f"{tag}: PEFT ref-adapter `elif` anchor gone; unsloth/models/rl.py "
f"peft_pattern re.sub no-ops and TRL's ref adapter init runs under Unsloth"
)
assert "ref_param.data.copy_(param.data)" in src, (
f"{tag}: `ref_param.data.copy_(param.data)` end-anchor gone; "
f"unsloth/models/rl.py peft_pattern loses its DOTALL end match"
)
assert "enable_input_require_grads" in src, (
f"{tag}: `enable_input_require_grads` block gone from grpo_trainer.py; "
f"the tightened PR #6904 regex assumed it follows the ref-adapter block"
)
@pytest.mark.parametrize("tag", TRL_TAGS)
def test_trl_grpo_quantized_model_cast_contract(tag: str):
"""rl.py (trl>=1.7.0) neutralizes TRL's hardcoded QLoRA bf16 cast
`if _is_quantized_model:` -> `if False:`. A rename leaves the cast active,
which ignores the user's dtype and breaks GradScaler with fp16=True."""
if not _tag_ge(tag, "1.7.0"):
pytest.skip(f"{tag}: pre-1.7.0 spells the cast differently (is_loaded_in_4bit)")
src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
assert src is not None
assert "if _is_quantized_model:" in src, (
f"{tag}: `if _is_quantized_model:` gone; unsloth/models/rl.py cannot "
f"neutralize TRL's hardcoded QLoRA bf16 cast and it runs under Unsloth"
)
@pytest.mark.parametrize("tag", TRL_TAGS)
def test_trl_grpo_aux_loss_enabled_contract(tag: str):
"""rl.py (trl>=1.7.0) appends a fail-fast after
`self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0` so an
explicit MoE router-aux opt-in errors instead of silently training without
the penalty (the optimized forward cannot compute it). A change to this
line drops the guard silently (PR #6904)."""
if not _tag_ge(tag, "1.7.0"):
pytest.skip(f"{tag}: aux_loss_enabled / router_aux_loss_coef added in TRL 1.7.0")
src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
assert src is not None
assert "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0" in src, (
f"{tag}: `aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0` "
f"changed; unsloth/models/rl.py's fail-fast .replace() anchor no-ops"
)
@pytest.mark.parametrize("tag", TRL_TAGS)
def test_trl_grpo_per_token_logps_aux_arity_contract(tag: str):
"""TRL 1.7.0 added `compute_aux_loss` to
_get_per_token_logps_and_entropies and made every call site unpack a
3-tuple. rl_replacements.py version-gates its injected replacement to emit
a 3-tuple for trl>=1.7.0 (2-tuple below). This is the exact change the
has_def existence checks miss: the method still exists, only its arity
changed. If TRL drops/renames the aux return, the gate needs revisiting."""
if not _tag_ge(tag, "0.20.0"):
pytest.skip(f"{tag}: pre-0.20 uses legacy _get_per_token_logps (2-tuple, no aux)")
src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
assert src is not None
assert has_def(src, "_get_per_token_logps_and_entropies", "func"), (
f"{tag}: _get_per_token_logps_and_entropies missing on TRL >=0.20; "
f"unsloth's per-token-logps injection dispatch key no longer matches"
)
if _tag_ge(tag, "1.7.0"):
assert "compute_aux_loss" in src, (
f"{tag}: TRL >=1.7.0 dropped `compute_aux_loss`; the 3-tuple "
f"injection gate in unsloth/models/rl_replacements.py must be revisited"
)