Fix TrainingArguments silently disabling unsloth gradient checkpointing (#6829)

* Fix TrainingArguments silently disabling unsloth gradient checkpointing

* Cover loaded adapters and preserve explicit None in GC restore

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

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

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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:
+#   <obj>._unsloth_gradient_checkpointing if hasattr(<obj>, '...') else getattr(<args>, 'gradient_checkpointing', True)
+_TERNARY = re.compile(
+    r"(?P<model>[\w.]+)\._unsloth_gradient_checkpointing "
+    r"if hasattr\((?P=model), '_unsloth_gradient_checkpointing'\) "
+    r"else getattr\((?P<args>[\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"

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,

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"

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)
 

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