diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index e492d21e9..b15d5bfa2 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -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: diff --git a/tests/version_compat/test_trl_grpo_fake_run.py b/tests/version_compat/test_trl_grpo_fake_run.py new file mode 100644 index 000000000..c79c9955d --- /dev/null +++ b/tests/version_compat/test_trl_grpo_fake_run.py @@ -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 = "") + 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" diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py index 834692e2d..2f935dde4 100644 --- a/tests/version_compat/test_trl_grpo_pinned_symbols.py +++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py @@ -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" + ) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index b5cadf2de..eeab8fbac 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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 = ( diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 098950de0..ffb845b04 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -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