Add regression tests for the stray-forward compile-cache reset (#6569)

* Add regression tests for the stray-forward compile-cache reset

Follow-up to #6511, which fixed the bug but whose squash merge did not
include the tests. These cover the two issues that fix addressed, under
the GPU-free tests/conftest.py harness:

- _unsloth_reset_stray_compile_cache is an exported module-level symbol in
  unsloth.models._utils (it previously lived only inside the RL trainer
  template string, so every non-RL import silently no-op'd)
- _unsloth_install_pretrain_detector keeps a recorded "seen" forward on an
  idempotent reinstall with a live hook, and only resets it after teardown
- only a grad-enabled pre-train forward marks the cache poisoned
- the reset warns and clears seen when a stray forward was seen, tears the
  hook down even on the clean path, and walks the .model/.base_model/.module
  wrapper chain to reach a nested marker

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

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

* Pin UNSLOTH_COMPILE_DISABLE in the warn-path reset tests

The reset only warns and resets Dynamo when UNSLOTH_COMPILE_DISABLE != "1".
A GPU-free CI env that sets it to "1" would make the warn assertion in
test_reset_clears_seen_and_warns_when_a_stray_forward_was_seen flaky.
monkeypatch it to "0" in both warn-path tests so the warn / no-warn
assertions are deterministic and test the seen flag, not the env.

---------

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

tests/test_pretrain_compile_reset.py

@@ -0,0 +1,153 @@
+# 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.
+
+"""The stray-pre-train-forward detector and its torch.compile cache reset.
+
+A grad-enabled forward/backward run before ``trainer.train()`` poisons the
+AOTAutograd backward-graph cache; the detector records it so train() can drop
+that cache. These cover the idempotent-reinstall evidence guard, the reset's
+chain-walk/teardown behaviour, and that the helper is importable at module
+scope (every non-RL training entry point imports it). Runs under the GPU-free
+``tests/conftest.py`` harness.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import unsloth  # noqa: F401  (installs the unsloth patches the functions live behind)
+
+import torch
+
+from unsloth.models._utils import (
+    _unsloth_install_pretrain_detector,
+    _unsloth_reset_stray_compile_cache,
+)
+
+
+class _Trainer:
+    """Minimal ``self`` stand-in: the reset only reads ``self.model``."""
+
+
+def test_reset_helper_is_importable_and_exported():
+    # Regression: the helper used to live only inside rl.py's RLTrainer_replacement template
+    # string (exec'd into a generated trainer module), so importing it from a real module raised
+    # ImportError and every non-RL consumer (SFT trainer.py, the plain-Trainer loop, the RL
+    # template's own delegation) silently no-op'd. Pin it as an exported module-level symbol.
+    from unsloth.models import _utils
+    assert callable(_utils._unsloth_reset_stray_compile_cache)
+    assert "_unsloth_reset_stray_compile_cache" in _utils.__all__
+
+
+def test_fresh_install_starts_unseen():
+    m = torch.nn.Linear(2, 2)
+    _unsloth_install_pretrain_detector(m)
+    marker = m._unsloth_pretrain_marker
+    assert marker["seen"] is False
+    assert "hook" in marker  # a live hook is registered
+
+
+def test_reinstall_with_live_hook_preserves_seen():
+    # Re-entering get_peft_model/patch_peft_model after a grad-enabled probe must NOT wipe the
+    # recorded poisoning, or train() skips the reset and the NaN/flat-loss bug returns.
+    m = torch.nn.Linear(2, 2)
+    _unsloth_install_pretrain_detector(m)
+    hook = m._unsloth_pretrain_marker["hook"]
+    m._unsloth_pretrain_marker["seen"] = True  # a probe the live hook recorded
+
+    _unsloth_install_pretrain_detector(m)  # idempotent re-install
+    marker = m._unsloth_pretrain_marker
+    assert marker["seen"] is True  # evidence kept
+    assert marker["hook"] is hook  # same hook, not double-registered
+
+
+def test_reinstall_after_teardown_resets_and_reregisters():
+    m = torch.nn.Linear(2, 2)
+    _unsloth_install_pretrain_detector(m)
+    marker = m._unsloth_pretrain_marker
+    marker["seen"] = True
+    marker.pop("hook").remove()  # simulate teardown (what the reset does)
+
+    _unsloth_install_pretrain_detector(m)  # no live hook -> fresh registration
+    assert marker["seen"] is False  # reset for the new session
+    assert "hook" in marker
+
+
+def test_grad_enabled_forward_marks_seen_no_grad_does_not():
+    m = torch.nn.Linear(2, 2)
+    _unsloth_install_pretrain_detector(m)
+    with torch.no_grad():
+        m(torch.zeros(1, 2))
+    assert m._unsloth_pretrain_marker["seen"] is False  # no backward graph -> clean
+    m(torch.zeros(1, 2))  # grad-enabled forward poisons the cache
+    assert m._unsloth_pretrain_marker["seen"] is True
+
+
+def test_reset_clears_seen_and_warns_when_a_stray_forward_was_seen(monkeypatch):
+    # Pin compile on: the reset only warns/resets when UNSLOTH_COMPILE_DISABLE != "1", which a
+    # GPU-free CI env may set, so force it here to make the warn assertion deterministic.
+    monkeypatch.setenv("UNSLOTH_COMPILE_DISABLE", "0")
+    m = torch.nn.Linear(2, 2)
+    _unsloth_install_pretrain_detector(m)
+    m._unsloth_pretrain_marker["seen"] = True  # a stray pre-train forward
+    trainer = _Trainer()
+    trainer.model = m
+
+    with warnings.catch_warnings(record = True) as caught:
+        warnings.simplefilter("always")
+        _unsloth_reset_stray_compile_cache(trainer)
+
+    assert any("manual forward/backward" in str(w.message) for w in caught)
+    assert "hook" not in m._unsloth_pretrain_marker  # hook torn down
+    assert m._unsloth_pretrain_marker["seen"] is False  # evidence consumed
+
+
+def test_reset_tears_down_hook_even_when_not_seen(monkeypatch):
+    # The clean path still removes the one-shot hook so it adds no per-step cost, but must not
+    # warn or reset Dynamo (nothing was poisoned). Pin compile on so the absent warning proves
+    # seen==False is the reason, not a disabled-compile short circuit.
+    monkeypatch.setenv("UNSLOTH_COMPILE_DISABLE", "0")
+    m = torch.nn.Linear(2, 2)
+    _unsloth_install_pretrain_detector(m)  # seen stays False
+    trainer = _Trainer()
+    trainer.model = m
+
+    with warnings.catch_warnings(record = True) as caught:
+        warnings.simplefilter("always")
+        _unsloth_reset_stray_compile_cache(trainer)
+
+    assert not any("manual forward/backward" in str(w.message) for w in caught)
+    assert "hook" not in m._unsloth_pretrain_marker
+    assert m._unsloth_pretrain_marker["seen"] is False
+
+
+def test_reset_walks_wrapper_chain_to_reach_a_nested_marker():
+    # The probe may have run on an inner wrapper (.model/.base_model/.module), not self.model.
+    inner = torch.nn.Linear(2, 2)
+    _unsloth_install_pretrain_detector(inner)
+    inner._unsloth_pretrain_marker["seen"] = True
+
+    class _Wrapper:  # e.g. a PEFT base_model wrapping the real module
+        pass
+
+    outer = _Wrapper()
+    outer.base_model = inner
+    trainer = _Trainer()
+    trainer.model = outer
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        _unsloth_reset_stray_compile_cache(trainer)
+
+    assert "hook" not in inner._unsloth_pretrain_marker  # found and torn down through the chain
+    assert inner._unsloth_pretrain_marker["seen"] is False