vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-10 00:08:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

(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
parent 934f879043
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

77
.github/workflows/version-compat-ci.yml vendored
View file

@ -285,6 +285,83 @@ jobs:
tests/vllm_compat/test_extended_module_imports.py \
-v --tb=short
# Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike
# the static symbol/source greps above, this drives unsloth's actual
# source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only
# runner under the tests/conftest.py spoof harness -- no GPU, no training.
# Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple
# per-token-logps return, restructured PEFT ref-adapter block) by asserting
# the generated Unsloth trainer still satisfies the transform contracts.
grpo-fake-run:
name: GRPO fake-run (latest + main TRL, CPU spoof)
runs-on: ubuntu-latest
timeout-minutes: 18
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
path: unsloth
- name: Clone unsloth-zoo @ main
run: |
for attempt in 1 2 3; do
rm -rf "$RUNNER_TEMP/unsloth-zoo"
if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \
"$RUNNER_TEMP/unsloth-zoo"; then
break
fi
if [ "$attempt" -eq 3 ]; then
echo "::error::git clone unsloth-zoo failed after 3 attempts"
exit 1
fi
delay=$((5 * attempt))
echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..."
sleep "$delay"
done
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- name: Install CPU torch + ecosystem + TRL latest
run: |
python -m pip install --upgrade pip
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10'
# Ecosystem floors unsloth needs; TRL itself is installed last so it
# can pull the transformers/peft it requires.
pip install \
'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \
'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \
'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow
pip install --upgrade trl
pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo"
pip install --no-deps -e ./unsloth
- name: Fake-run vs TRL latest
env:
UNSLOTH_IS_PRESENT: '1'
UNSLOTH_COMPILE_DISABLE: '1'
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
run: |
cd unsloth
python -c "import trl; print('Resolved TRL', trl.__version__)"
PYTHONPATH=. python -m pytest \
tests/version_compat/test_trl_grpo_fake_run.py \
-v --tb=short
# `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge
# TRL break does not red every PR. github.event_name is valid in a step if.
- name: Fake-run vs TRL main (scheduled / dispatch only)
if: ${{ github.event_name != 'pull_request' }}
env:
UNSLOTH_IS_PRESENT: '1'
UNSLOTH_COMPILE_DISABLE: '1'
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
run: |
pip install --upgrade "git+https://github.com/huggingface/trl"
cd unsloth
python -c "import trl; print('Resolved TRL', trl.__version__)"
PYTHONPATH=. python -m pytest \
tests/version_compat/test_trl_grpo_fake_run.py \
-v --tb=short
# Daily-only: same suites but with --strict on importable upstream
# tags. Schedule-only so PR jobs stay fast; cron tolerates a flake.
daily-fresh-fetch:

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"
)

36
unsloth/models/rl.py
View file

@ -1651,7 +1651,36 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
RLTrainer_source = re.sub(pattern, new_options, RLTrainer_source, flags = re.DOTALL)
if trl_version >= Version("0.27.0"):
if trl_version >= Version("1.4.0"):
# The `elif is_peft_model(model) and args.beta != 0.0:` ref-adapter block
# was introduced in TRL 1.4.0 and is used through 1.7.x. Remove only that
# block, anchored on the final ref_param copy so we do NOT also swallow the
# following gradient-checkpointing enable_input_require_grads() block.
peft_pattern = (
r"\s*elif is_peft_model\(model\) and args\.beta != 0\.0:"
r".*?"
r"ref_param\.data\.copy_\(param\.data\)"
)
replacement_comment = (
"\n # PEFT initialization logic removed via script for trl >= 1.4.0\n"
)
RLTrainer_source = re.sub(
peft_pattern, replacement_comment, RLTrainer_source, flags = re.DOTALL
)
if trl_version >= Version("1.7.0"):
# router_aux_loss_coef / aux_loss_enabled were added in TRL 1.7.0. Unsloth's
# optimized GRPO forward cannot compute the MoE router aux loss, so reject
# explicit opt-in (router_aux_loss_coef > 0) at init rather than silently ignoring it.
RLTrainer_source = RLTrainer_source.replace(
"self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0",
"self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0\n"
' if self.aux_loss_enabled: raise NotImplementedError("Unsloth GRPO does not compute the MoE router auxiliary loss; set router_aux_loss_coef = 0 (the Unsloth default).")',
)
elif trl_version >= Version("0.27.0"):
peft_pattern = (
r"\s*if is_peft_available\(\) and is_peft_model\(model\) and args\.beta != 0\.0:"
r".*?"
@ -1689,6 +1718,11 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
'if getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False):',
"if False:",
)
# TRL >= 1.7.0 spells the same QLoRA bf16 cast as `if _is_quantized_model:`.
RLTrainer_source = RLTrainer_source.replace(
"if _is_quantized_model:",
"if False:",
)
if RLTrainer_name == "SFTTrainer":
original_text = (

42
unsloth/models/rl_replacements.py
View file

@ -1203,6 +1203,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
self._autocast_dtype = torch.float16
compute_aux_loss = kwargs.get("compute_aux_loss", None)
pixel_values, image_grid_thw = (
kwargs.get("pixel_values", None),
kwargs.get("image_grid_thw", None),
@ -1846,7 +1848,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
_extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk
with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype):
if pixel_values is None:
logits_chunk = unwrapped_model(
outputs = unwrapped_model(
input_ids = input_ids_chunk,
attention_mask = attention_mask_chunk,
pixel_values = pixel_values_chunk,
@ -1854,7 +1856,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
pixel_attention_mask = pixel_attention_mask_chunk,
image_sizes = image_sizes_chunk,
**_extra_vision_kwargs,
).logits
)
logits_chunk = outputs.logits
del outputs # free hidden_states before chunked log-softmax
completion_input_ids_chunk = input_ids_chunk[
:, -(logits_to_keep + max_left_pad) :
@ -1876,7 +1881,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
else:
# Essentially, for VLMs we do not go via the optimized path in models/,
# so we don't encounter the Flash Attn left-padding issue.
logits_chunk = unwrapped_model(
outputs = unwrapped_model(
input_ids = input_ids_chunk,
attention_mask = attention_mask_chunk,
pixel_values = pixel_values_chunk,
@ -1885,7 +1890,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
image_sizes = image_sizes_chunk,
logits_to_keep = logits_to_keep + 1,
**_extra_vision_kwargs,
).logits
)
logits_chunk = outputs.logits
del outputs # free hidden_states before chunked log-softmax
logits_chunk = logits_chunk[:, :-1, :]
completion_input_ids_chunk = input_ids_chunk[:, -logits_to_keep:]
@ -1914,11 +1922,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
all_logprobs_list.append(logprobs_chunk)
if logprobs is None: # padded fallback when packing was not used
logprobs = torch.cat(all_logprobs_list, dim = 0)
entropies = None
os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0"
return logprobs.detach(), entropies # logps, entropies
# aux loss is unused: it is off by default (router_aux_loss_coef set to 0 in models/rl.py)
# and explicit opt-in is rejected at trainer init, so this is always None (kept in the
# return for TRL >= 1.7.0's 3-tuple contract).
aux_loss = None
return logprobs.detach(), entropies, aux_loss # logps, entropies, aux_loss
# input_ids = input_ids[:, -logits_to_keep:]
# For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
# See https://github.com/huggingface/trl/issues/2770
@ -1937,6 +1949,24 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
# return logps # compute logprobs for the input tokens
function = inspect.getsource(_get_per_token_logps_and_entropies)
if trl_version < Version("1.7.0"):
# TRL < 1.7.0 unpacks (logps, entropies) at every call site; TRL >= 1.7.0
# always unpacks (logps, entropies, aux_loss). Drop the aux_loss element so
# the return arity matches the installed TRL. Regex tolerates comment /
# whitespace drift on the return line; fail loud if the anchor ever stops
# matching rather than silently shipping a 3-tuple to older TRL.
new_function, n = re.subn(
r"return (logprobs\.detach\(\), entropies), aux_loss[^\n]*",
r"return \1 # logps, entropies",
function,
)
if n != 1:
raise RuntimeError(
"Unsloth GRPO: could not downgrade the per-token-logps return to a "
f"2-tuple for TRL {trl_version} (matched {n} times, expected 1). The "
"return line changed; update the arity gate in rl_replacements.py."
)
function = new_function
return function