diff --git a/tests/test_gradient_checkpointing_restore.py b/tests/test_gradient_checkpointing_restore.py new file mode 100644 index 0000000000..4f9f3faccc --- /dev/null +++ b/tests/test_gradient_checkpointing_restore.py @@ -0,0 +1,185 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Regression for #4735: a plain ``TrainingArguments`` silently disabling the +gradient-checkpointing (GC) mode the model was configured with at setup. + +Setup records the effective GC mode as ``_unsloth_gradient_checkpointing``; the +trainer restores *that* value, falling back to ``args.gradient_checkpointing`` +only when nothing was recorded. The restore lines live inside exec'd template +strings, which ``py_compile`` never sees, so these tests pull the real snippets +out of the source and execute them against fakes. GPU-free. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + +_ROOT = Path(__file__).resolve().parent.parent / "unsloth" / "models" +_RL = (_ROOT / "rl.py").read_text() +_RL_REPLACEMENTS = (_ROOT / "rl_replacements.py").read_text() + +# The single-line ternary form used at the trainer call sites: +# ._unsloth_gradient_checkpointing if hasattr(, '...') else getattr(, 'gradient_checkpointing', True) +_TERNARY = re.compile( + r"(?P[\w.]+)\._unsloth_gradient_checkpointing " + r"if hasattr\((?P=model), '_unsloth_gradient_checkpointing'\) " + r"else getattr\((?P[\w.]+), 'gradient_checkpointing', True\)" +) + +_MISSING = object() + + +class _Obj: + """Bare attribute bag; ``_unsloth_gradient_checkpointing`` present only when recorded.""" + + def __init__( + self, + recorded = _MISSING, + gradient_checkpointing = _MISSING, + ): + if recorded is not _MISSING: + self._unsloth_gradient_checkpointing = recorded + if gradient_checkpointing is not _MISSING: + self.gradient_checkpointing = gradient_checkpointing + + +class _Self: + def __init__( + self, + model = None, + args = None, + ): + if model is not None: + self.model = model + self.args = args + + +# (recorded on model, args.gradient_checkpointing, expected restored value) +# The point of the fix: a recorded mode wins over args, and a recorded ``None`` +# (a valid setup value) is restored verbatim rather than collapsing to the +# args fallback the way a ``None`` sentinel would. +_MATRIX = [ + ("unsloth", False, "unsloth"), # the #4735 case: args=False must NOT win + (True, False, True), + (False, True, False), # user turned GC off; args=True must NOT re-enable it + (None, True, None), # explicit None is restored, not treated as "unrecorded" + (_MISSING, True, True), # nothing recorded -> fall back to args + (_MISSING, False, False), +] + + +def _eval_ternary(expr, recorded, args_gc): + """Eval a restore expression that references either ``model``/``args`` or ``self.model``/``self.args``.""" + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + self = _Self(model = model, args = args) + return eval( + expr, {"hasattr": hasattr, "getattr": getattr}, {"model": model, "args": args, "self": self} + ) + + +def test_ternary_restore_semantics(): + exprs = [m.group(0) for m in _TERNARY.finditer(_RL)] + exprs += [m.group(0) for m in _TERNARY.finditer(_RL_REPLACEMENTS)] + # Also guards against the lines being deleted/renamed (which reinstates the bug). + assert len(exprs) >= 3, f"expected the 3 trainer-call restore sites, found {len(exprs)}" + for expr in exprs: + for recorded, args_gc, expected in _MATRIX: + got = _eval_ternary(expr, recorded, args_gc) + assert got == expected and type(got) is type( + expected + ), f"{expr!r}: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def _extract_prepare_restore_block(): + """Pull the multi-line restore block out of ``prepare_for_training_mode``'s wrapper. + + It lives inside an exec'd template string, so grab it textually: from the + ``_model = getattr(self, 'model', None)`` line through the closing + ``else:``/``use_gc = ...`` pair. + """ + lines = _RL.splitlines() + start = next( + i for i, l in enumerate(lines) if l.strip() == "_model = getattr(self, 'model', None)" + ) + # End at the fallback assignment rather than a fixed line count, so inserting + # lines into the block can't silently truncate what gets exec'd. + end = next( + i + for i, l in enumerate(lines) + if i > start and "use_gc = getattr(self.args, 'gradient_checkpointing', True)" in l + ) + block = lines[start : end + 1] + # dedent to column 0 so it execs as a top-level block + indent = len(block[0]) - len(block[0].lstrip()) + return "\n".join(l[indent:] for l in block) + + +def test_prepare_for_training_mode_block_semantics(): + block = _extract_prepare_restore_block() + # Must be valid Python (it's never seen by py_compile in the outer file). + ast.parse(block) + + for recorded, args_gc, expected in _MATRIX: + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + ns = {"self": _Self(model = model, args = args), "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + got = ns["use_gc"] + assert ( + got == expected and type(got) is type(expected) + ), f"prepare block: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def test_prepare_block_tolerates_missing_model(): + # gemini flagged the unguarded self.model access: the block reads self.model via + # getattr(self, 'model', None), so a trainer without a .model attribute must fall + # back to args rather than raising AttributeError. + block = _extract_prepare_restore_block() + args = _Obj(gradient_checkpointing = True) + self_no_model = _Self(model = None, args = args) # _Self leaves .model unset when model is None + assert not hasattr(self_no_model, "model") + ns = {"self": self_no_model, "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + assert ns["use_gc"] is True + + +def test_recording_sites_are_real_module_code(): + # The recording side (unlike the restore side) is real module code, not a template + # string. Assert it's present at the choke point (patch_peft_model, so loaded adapters + # are covered) and at the pre-wrapped pass-through, both of which bypass the old + # get_peft_model-only recording. + llama = (_ROOT / "llama.py").read_text() + tree = ast.parse(llama) + + def assigns_marker(node): + return any( + isinstance(n, ast.Assign) + and any( + isinstance(t, ast.Attribute) and t.attr == "_unsloth_gradient_checkpointing" + for t in n.targets + ) + for n in ast.walk(node) + ) + + fns = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + assert "patch_peft_model" in fns and assigns_marker( + fns["patch_peft_model"] + ), "patch_peft_model must record _unsloth_gradient_checkpointing so loaded adapters are covered" + # The pass-through branch lives in get_peft_model. + assert assigns_marker( + fns["get_peft_model"] + ), "get_peft_model pass-through must record _unsloth_gradient_checkpointing" diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 14ee5ee24e..bb7289dfa8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3047,6 +3047,9 @@ class FastLlamaModel: # Pre-wrapped PEFT model passes through here; still arm the detector so an RL # trainer can reset a compile cache poisoned by a pre-train forward. _unsloth_install_pretrain_detector(model) + # This branch returns before patch_peft_model, so record here too; + # apply_unsloth_gradient_checkpointing above already re-patched global state to match (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing model = _exclude_rope_inv_freq_from_ddp(model) return model else: @@ -3406,6 +3409,11 @@ class FastLlamaModel: @staticmethod def patch_peft_model(model, use_gradient_checkpointing = "unsloth"): + # Persist the effective GC mode so the trainer restores it verbatim: for_inference() + # clears the module flags every GRPO step, and a plain TrainingArguments defaults it to + # False, which would otherwise silently disable it at train time (#4735). Recorded here, + # not in get_peft_model, so adapters loaded via loader.py's from_pretrained path are covered. + model._unsloth_gradient_checkpointing = use_gradient_checkpointing if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1": return FastBaseModel.patch_peft_model( model = model, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 602de69d3f..62ef9e916a 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -423,8 +423,14 @@ def prepare_for_training_mode(f): pass # Enable training mode _was_training = None - # Get gradient checkpointing setting from training arguments - use_gc = getattr(self.args, 'gradient_checkpointing', True) + # Restore the GC mode the model was configured with at setup; fall back to + # the training args only when it wasn't recorded (issue #4735). Use hasattr, + # not a None sentinel, so a deliberately-recorded None is restored verbatim. + _model = getattr(self, 'model', None) + if hasattr(_model, '_unsloth_gradient_checkpointing'): + use_gc = _model._unsloth_gradient_checkpointing + else: + use_gc = getattr(self.args, 'gradient_checkpointing', True) if hasattr(self, 'model') and hasattr(self.model, "training"): _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): @@ -532,7 +538,8 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): if getattr(args, "_n_gpu", 1) != 1: args._n_gpu = 1 if "model" in locals() and hasattr(model, "for_training"): - model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) + model.for_training(use_gradient_checkpointing=_use_gc) super().__init__({RLTrainer_call_args}{RLTrainer_kwargs}) if "model" in locals() and hasattr(model, "for_inference"): model.for_inference() @@ -1165,7 +1172,8 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): if "model" in call_args: training_check = ( "if model is not None and hasattr(model, 'for_training'):\n" - " model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))\n" + " _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True)\n" + " model.for_training(use_gradient_checkpointing=_use_gc)\n" "if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n" "if 'processing_class' in locals():\n" " if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n" diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index d3ada23cf9..3be614cf4a 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -761,7 +761,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() - self.model.for_training(use_gradient_checkpointing=getattr(self.args, 'gradient_checkpointing', True))""" + _use_gc = self.model._unsloth_gradient_checkpointing if hasattr(self.model, '_unsloth_gradient_checkpointing') else getattr(self.args, 'gradient_checkpointing', True) + self.model.for_training(use_gradient_checkpointing=_use_gc)""" function = function.replace(line_to_replace, replacement_lines) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 5ab55152db..689e362f95 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1873,6 +1873,11 @@ class FastBaseModel: float32_mixed_precision = float32_mixed_precision, patch_modules_to_save = True, ) + # Persist the configured GC mode so the trainer restores it verbatim. + # for_inference() clears the module flags (GRPO does this every generation + # step), and a plain TrainingArguments defaults gradient_checkpointing=False, + # which would otherwise silently disable this setting at train time (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing # Gemma3N audio conformer processes variable-length audio tensors # that cause stride mismatches in AOT autograd compiled backward