From 9026037298c39836c077f2e08901e19d13458f1a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 20 Jun 2026 12:41:15 +0000
Subject: [PATCH 01/11] Reset torch.compile cache poisoned by a stray forward
 before trainer.train()

A manual forward / forward+backward run under model.train() before
trainer.train() (for example a pre-train grad-norm probe like
out = model(**batch); out.loss.backward()) silently poisons training when
torch.compile is enabled. The stray training-mode pass is the first one in the
process, so it compiles and caches the model forward and, via AOTAutograd, its
backward graph in a one-off context that does not match the real training loop.
When trainer.train() reuses that cached graph the gradients come out NaN/Inf,
the loss never moves, and the run looks like it trains but never learns.

Observed on gpt-oss-20b (loss frozen at ~4.25, grad_norm NaN from step 1) with
both use_gradient_checkpointing="unsloth" and =True. It does not reproduce when
the probe runs under torch.no_grad(), nor with UNSLOTH_COMPILE_DISABLE=1, and a
single torch._dynamo.reset() before training fully cures it (loss 4.29 -> 0.0002,
identical to a run with no probe). Resetting the gradient-checkpointing buffers,
zero_grad, empty_cache, or for_training does not help, confirming the corruption
lives in the torch._dynamo / torch.compile cache.

get_peft_model now attaches a one-shot forward pre-hook that records whether a
forward ran before train(). prepare_for_training_mode checks it at the start of
train() and, if a pre-train forward was seen and torch.compile is enabled, calls
torch._dynamo.reset() (plus a pristine gradient-checkpoint reset and zero_grad)
and warns once. On the normal path (no pre-train forward) it is a strict no-op:
no dynamo reset, no recompilation, identical loss curve.
---
 unsloth/models/_utils.py | 28 ++++++++++++++++++++++++
 unsloth/models/llama.py  |  6 ++++++
 unsloth/models/rl.py     | 46 ++++++++++++++++++++++++++++++++++++++++
 unsloth/models/vision.py |  3 +++
 4 files changed, 83 insertions(+)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 742f5f57e0..b3f42eea55 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -62,6 +62,7 @@ __all__ = [
     "patch_unsloth_smart_gradient_checkpointing",
     "unpatch_unsloth_smart_gradient_checkpointing",
     "apply_unsloth_gradient_checkpointing",
+    "_unsloth_install_pretrain_detector",
     "patch_compiled_autograd",
     "process_vision_info",
     "unsloth_compile_transformers",
@@ -194,6 +195,33 @@ from unsloth_zoo.temporary_patches import (
 )
 
 
+def _unsloth_install_pretrain_detector(model):
+    """Attach a one-shot forward pre-hook that records whether a forward ran
+    before trainer.train(). Used by prepare_for_training_mode to drop a
+    torch.compile graph cache poisoned by a stray manual forward/backward.
+    Idempotent and a no-op if the model cannot take hooks."""
+    if model is None or not hasattr(model, "register_forward_pre_hook"):
+        return model
+    marker = getattr(model, "_unsloth_pretrain_marker", None)
+    if isinstance(marker, dict):
+        marker["seen"] = False
+        return model
+    marker = {"seen": False}
+    try:
+        model._unsloth_pretrain_marker = marker
+    except Exception:
+        return model
+
+    def _mark(_module, _inp):
+        marker["seen"] = True
+
+    try:
+        marker["hook"] = model.register_forward_pre_hook(_mark)
+    except Exception:
+        pass
+    return model
+
+
 def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_length, dtype):
     """
     Apply gradient checkpointing with smart heuristics.
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 08802f030e..e818fea68b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -3368,6 +3368,9 @@ class FastLlamaModel:
             m.for_training = functools.partial(FastBaseModel.for_training, m)
             m.for_inference = functools.partial(FastBaseModel.for_inference, m)
             m = m.model
+        # Detect a stray pre-train forward so train() can drop the torch.compile
+        # graph cache it would otherwise poison (see prepare_for_training_mode).
+        _unsloth_install_pretrain_detector(model)
         return model
 
     @staticmethod
@@ -3585,6 +3588,9 @@ class FastLlamaModel:
             m.for_training = functools.partial(FastBaseModel.for_training, m)
             m.for_inference = functools.partial(FastBaseModel.for_inference, m)
             m = m.model
+        # Detect a stray pre-train forward so train() can drop the torch.compile
+        # graph cache it would otherwise poison (see prepare_for_training_mode).
+        _unsloth_install_pretrain_detector(model)
         return model
 
     @staticmethod
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index c3237b6481..a153cf1872 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -390,9 +390,55 @@ try:
     from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
 except:
     def reset_unsloth_gradient_checkpointing_buffers(): pass
+def _unsloth_reset_stray_compile_cache(self):
+    # A manual forward / forward+backward run under torch.compile BEFORE
+    # trainer.train() (e.g. a pre-train grad-norm probe) compiles and caches the
+    # model's forward and (via AOTAutograd) its backward graph in a one-off
+    # context that does not match the training loop. Reusing that cached graph
+    # poisons training with NaN/zero gradients (loss never moves). If a pre-train
+    # forward was seen and torch.compile is enabled, drop the compiled-graph cache
+    # so training recompiles cleanly. No-op on the normal path.
+    import os
+    model = getattr(self, "model", None)
+    if model is None:
+        return
+    marker = getattr(model, "_unsloth_pretrain_marker", None)
+    seen = bool(marker.get("seen")) if isinstance(marker, dict) else False
+    if seen and os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1":
+        try:
+            import torch._dynamo as _dynamo
+            _dynamo.reset()
+        except Exception:
+            pass
+        try:
+            reset_unsloth_gradient_checkpointing_buffers()
+        except Exception:
+            pass
+        try:
+            model.zero_grad(set_to_none = True)
+        except Exception:
+            pass
+        import warnings
+        warnings.warn(
+            "Unsloth: detected a manual forward/backward run before trainer.train(); "
+            "reset the torch.compile graph cache it poisoned so training starts clean. "
+            "To avoid this, run any pre-train probe under `with torch.no_grad():`."
+        )
+    # Tear down the one-shot detector hook so it never adds per-step cost.
+    if isinstance(marker, dict):
+        hook = marker.pop("hook", None)
+        if hook is not None:
+            try: hook.remove()
+            except Exception: pass
+        marker["seen"] = False
 def prepare_for_training_mode(f):
     @functools.wraps(f)
     def wrapper(self, *args, **kwargs):
+        # Drop any torch.compile graph cache poisoned by a stray pre-train forward.
+        try:
+            _unsloth_reset_stray_compile_cache(self)
+        except Exception:
+            pass
         # Finish the previous W&B run if this is a subsequent train() call.
         # We do this at the START of train() (not the end) so that
         # evaluate() / log() still work after train() completes.
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 3f8d9e6ce0..13b66a59d4 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -1577,6 +1577,9 @@ class FastBaseModel:
             m.for_training = functools.partial(FastBaseModel.for_training, m)
             m.for_inference = functools.partial(FastBaseModel.for_inference, m)
             m = m.model
+        # Detect a stray pre-train forward so train() can drop the torch.compile
+        # graph cache it would otherwise poison (see prepare_for_training_mode).
+        _unsloth_install_pretrain_detector(model)
         return model
 
     @staticmethod

From b77b63e498b3713ffa18875a75251fcff19c105d Mon Sep 17 00:00:00 2001
From: danielhanchen <michaelhan2050@gmail.com>
Date: Sun, 21 Jun 2026 13:21:43 +0000
Subject: [PATCH 02/11] Ignore no-grad pre-train probes and detect probes
 across the wrapper chain

A no-grad forward (with torch.no_grad(): model(**batch)) builds no AOTAutograd
backward graph, so it cannot poison the compiled training graph. Gate the marker
on torch.is_grad_enabled() so such probes no longer trigger a needless dynamo
reset, recompile and warning on an otherwise clean run.

Also walk the model wrapper chain (PeftModel / DDP / base model) when resetting
so a probe that ran on a different wrapper than self.model is still detected, and
tear down every detector hook in the chain. Re-installing the detector is now
idempotent and only re-registers when a prior hook was already removed.
---
 unsloth/models/_utils.py | 24 +++++++++++++++++-------
 unsloth/models/rl.py     | 29 +++++++++++++++++++++++------
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index b3f42eea55..6383e14ac9 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -205,15 +205,25 @@ def _unsloth_install_pretrain_detector(model):
     marker = getattr(model, "_unsloth_pretrain_marker", None)
     if isinstance(marker, dict):
         marker["seen"] = False
-        return model
-    marker = {"seen": False}
-    try:
-        model._unsloth_pretrain_marker = marker
-    except Exception:
-        return model
+        # Re-register only if the previous hook was torn down (e.g. by an earlier
+        # train()); if the hook is still live this is a strict no-op so we never
+        # stack duplicate hooks.
+        if "hook" in marker:
+            return model
+    else:
+        marker = {"seen": False}
+        try:
+            model._unsloth_pretrain_marker = marker
+        except Exception:
+            return model
 
     def _mark(_module, _inp):
-        marker["seen"] = True
+        # Only a GRAD-ENABLED forward can poison the AOTAutograd/torch.compile
+        # backward-graph cache. A no-grad probe (`with torch.no_grad(): model(...)`,
+        # the sanity check the warning recommends) builds no backward graph, so
+        # treat it as clean and avoid a needless dynamo reset + recompile + warning.
+        if torch.is_grad_enabled():
+            marker["seen"] = True
 
     try:
         marker["hook"] = model.register_forward_pre_hook(_mark)
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index a153cf1872..cf64ee4e2c 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -402,8 +402,25 @@ def _unsloth_reset_stray_compile_cache(self):
     model = getattr(self, "model", None)
     if model is None:
         return
-    marker = getattr(model, "_unsloth_pretrain_marker", None)
-    seen = bool(marker.get("seen")) if isinstance(marker, dict) else False
+    # The detector hook may sit on any wrapper in the chain (PeftModel / DDP /
+    # the base model), and a pre-train probe could have run on a different
+    # wrapper than self.model. Walk the chain so a "seen" marker anywhere is
+    # detected, and collect every marker so all hooks are torn down below.
+    markers = []
+    seen = False
+    _curr = model
+    _visited = set()
+    while _curr is not None and id(_curr) not in _visited:
+        _visited.add(id(_curr))
+        _m = getattr(_curr, "_unsloth_pretrain_marker", None)
+        if isinstance(_m, dict):
+            markers.append(_m)
+            if _m.get("seen"):
+                seen = True
+        _nxt = getattr(_curr, "model", None)
+        if _nxt is None:
+            _nxt = getattr(_curr, "base_model", None)
+        _curr = _nxt
     if seen and os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1":
         try:
             import torch._dynamo as _dynamo
@@ -424,13 +441,13 @@ def _unsloth_reset_stray_compile_cache(self):
             "reset the torch.compile graph cache it poisoned so training starts clean. "
             "To avoid this, run any pre-train probe under `with torch.no_grad():`."
         )
-    # Tear down the one-shot detector hook so it never adds per-step cost.
-    if isinstance(marker, dict):
-        hook = marker.pop("hook", None)
+    # Tear down every one-shot detector hook in the chain so none adds per-step cost.
+    for _m in markers:
+        hook = _m.pop("hook", None)
         if hook is not None:
             try: hook.remove()
             except Exception: pass
-        marker["seen"] = False
+        _m["seen"] = False
 def prepare_for_training_mode(f):
     @functools.wraps(f)
     def wrapper(self, *args, **kwargs):

From 76d79325034c7489c180e1b75ff8e36bb192bc13 Mon Sep 17 00:00:00 2001
From: danielhanchen <michaelhan2050@gmail.com>
Date: Mon, 22 Jun 2026 02:07:57 +0000
Subject: [PATCH 03/11] Walk DDP/FSDP .module when scanning for the pre-train
 marker

The chain walk followed only .model and .base_model, so a probe that fired on the
model below a DDP/FSDP wrapper (which exposes it via .module) left the marker
undetected and the poisoned compile cache un-reset. Add .module to the walk.
---
 unsloth/models/rl.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index cf64ee4e2c..f399a7b3e8 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -417,9 +417,14 @@ def _unsloth_reset_stray_compile_cache(self):
             markers.append(_m)
             if _m.get("seen"):
                 seen = True
+        # Follow the wrapper chain: Unsloth/HF (.model), PEFT (.base_model) and
+        # DDP / FSDP (.module). A pre-train probe can fire on the model below a
+        # DDP wrapper, so .module must be walked too or the marker is missed.
         _nxt = getattr(_curr, "model", None)
         if _nxt is None:
             _nxt = getattr(_curr, "base_model", None)
+        if _nxt is None:
+            _nxt = getattr(_curr, "module", None)
         _curr = _nxt
     if seen and os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1":
         try:

From 8596138a86a1d08abad8636ef6d8e85a2b3c0ada Mon Sep 17 00:00:00 2001
From: danielhanchen <michaelhan2050@gmail.com>
Date: Mon, 22 Jun 2026 04:59:35 +0000
Subject: [PATCH 04/11] Install pre-train detector on the full-finetuning path
 too

get_peft_model returns early when UNSLOTH_ENABLE_FULL_FINETUNING=1, before the
detector was installed, so full-finetuning runs (which still use torch.compile) did
not drop a graph cache poisoned by a stray pre-train forward. Install the detector
before both full-finetuning early returns (FastLlamaModel and FastBaseModel). The
detector is idempotent, so this never stacks duplicate hooks when get_peft_model is
also called on a LoRA model.
---
 unsloth/models/llama.py  | 3 +++
 unsloth/models/vision.py | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index e818fea68b..c69847c048 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2938,6 +2938,9 @@ class FastLlamaModel:
             )
         if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1":
             print("Unsloth: Full finetuning is enabled, so .get_peft_model has no effect")
+            # Full finetuning still compiles, so a stray pre-train forward can poison the
+            # cache; install the detector here too (it is idempotent).
+            _unsloth_install_pretrain_detector(model)
             return model
         transformers_set_seed(random_state)
 
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 13b66a59d4..fbb47df400 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -1380,6 +1380,9 @@ class FastBaseModel:
     ):
         if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1":
             print("Unsloth: Full finetuning is enabled, so .get_peft_model has no effect")
+            # Full finetuning still compiles, so a stray pre-train forward can poison the
+            # cache; install the detector here too (it is idempotent).
+            _unsloth_install_pretrain_detector(model)
             return model
         transformers_set_seed(random_state)
 

From d2f6eec004252f77b4bd7cd2d402f39b0deb06ed Mon Sep 17 00:00:00 2001
From: danielhanchen <michaelhan2050@gmail.com>
Date: Mon, 22 Jun 2026 07:55:02 +0000
Subject: [PATCH 05/11] torch.compile stray-forward reset: tighten comments (no
 code change)

---
 unsloth/models/_utils.py | 17 ++++++-----------
 unsloth/models/rl.py     | 22 ++++++++--------------
 2 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 6383e14ac9..fbdfb84602 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -196,18 +196,15 @@ from unsloth_zoo.temporary_patches import (
 
 
 def _unsloth_install_pretrain_detector(model):
-    """Attach a one-shot forward pre-hook that records whether a forward ran
-    before trainer.train(). Used by prepare_for_training_mode to drop a
-    torch.compile graph cache poisoned by a stray manual forward/backward.
-    Idempotent and a no-op if the model cannot take hooks."""
+    """Attach a one-shot forward pre-hook recording whether a forward ran before
+    trainer.train(), so prepare_for_training_mode can drop a torch.compile graph cache poisoned
+    by a stray manual forward/backward. Idempotent; no-op if the model cannot take hooks."""
     if model is None or not hasattr(model, "register_forward_pre_hook"):
         return model
     marker = getattr(model, "_unsloth_pretrain_marker", None)
     if isinstance(marker, dict):
         marker["seen"] = False
-        # Re-register only if the previous hook was torn down (e.g. by an earlier
-        # train()); if the hook is still live this is a strict no-op so we never
-        # stack duplicate hooks.
+        # Re-register only if the previous hook was torn down; a live hook stays (no duplicates).
         if "hook" in marker:
             return model
     else:
@@ -218,10 +215,8 @@ def _unsloth_install_pretrain_detector(model):
             return model
 
     def _mark(_module, _inp):
-        # Only a GRAD-ENABLED forward can poison the AOTAutograd/torch.compile
-        # backward-graph cache. A no-grad probe (`with torch.no_grad(): model(...)`,
-        # the sanity check the warning recommends) builds no backward graph, so
-        # treat it as clean and avoid a needless dynamo reset + recompile + warning.
+        # Only a grad-enabled forward poisons the AOTAutograd backward-graph cache; a no-grad
+        # probe builds no backward graph, so treat it as clean (avoids a needless dynamo reset).
         if torch.is_grad_enabled():
             marker["seen"] = True
 
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index f399a7b3e8..1a00f21416 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -391,21 +391,17 @@ try:
 except:
     def reset_unsloth_gradient_checkpointing_buffers(): pass
 def _unsloth_reset_stray_compile_cache(self):
-    # A manual forward / forward+backward run under torch.compile BEFORE
-    # trainer.train() (e.g. a pre-train grad-norm probe) compiles and caches the
-    # model's forward and (via AOTAutograd) its backward graph in a one-off
-    # context that does not match the training loop. Reusing that cached graph
-    # poisons training with NaN/zero gradients (loss never moves). If a pre-train
-    # forward was seen and torch.compile is enabled, drop the compiled-graph cache
-    # so training recompiles cleanly. No-op on the normal path.
+    # A manual forward/backward under torch.compile BEFORE trainer.train() (e.g. a grad-norm
+    # probe) caches a forward + AOTAutograd backward graph in a one-off context; reusing it
+    # poisons training with NaN/zero gradients. If such a forward was seen and compile is on,
+    # drop the compiled-graph cache so training recompiles cleanly. No-op on the normal path.
     import os
     model = getattr(self, "model", None)
     if model is None:
         return
-    # The detector hook may sit on any wrapper in the chain (PeftModel / DDP /
-    # the base model), and a pre-train probe could have run on a different
-    # wrapper than self.model. Walk the chain so a "seen" marker anywhere is
-    # detected, and collect every marker so all hooks are torn down below.
+    # The detector hook can sit on any wrapper in the chain, and the probe may have run on a
+    # different one than self.model, so walk the chain: detect a "seen" marker anywhere and
+    # collect every marker to tear down below.
     markers = []
     seen = False
     _curr = model
@@ -417,9 +413,7 @@ def _unsloth_reset_stray_compile_cache(self):
             markers.append(_m)
             if _m.get("seen"):
                 seen = True
-        # Follow the wrapper chain: Unsloth/HF (.model), PEFT (.base_model) and
-        # DDP / FSDP (.module). A pre-train probe can fire on the model below a
-        # DDP wrapper, so .module must be walked too or the marker is missed.
+        # Follow the wrapper chain: Unsloth/HF (.model), PEFT (.base_model), DDP/FSDP (.module).
         _nxt = getattr(_curr, "model", None)
         if _nxt is None:
             _nxt = getattr(_curr, "base_model", None)

From 430af1776b3e0066cb94fbb0a3cd7208669b02a1 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 22 Jun 2026 10:08:51 +0000
Subject: [PATCH 06/11] Wire stray-forward compile-cache reset into SFT path
 and PEFT pass-through

The pre-train forward detector is installed for plain LoRA/vision models in
get_peft_model, but only RL trainers ran the reset via prepare_for_training_mode.
A grad-enabled probe before SFTTrainer.train() therefore left the poisoned Dynamo
cache in place and the detector hook running on every training forward.

- trainer.py: wrap SFTTrainer.train to run _unsloth_reset_stray_compile_cache,
  which both drops the poisoned cache and tears down the detector hook. For
  UnslothSFTTrainer the later prepare_for_training_mode assignment supersedes it.
- llama.py: arm the detector before the 'Already have LoRA adapters' early return
  so pre-wrapped PEFT models keep the reset capability.
---
 unsloth/models/llama.py |  3 +++
 unsloth/trainer.py      | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c69847c048..baea9674d6 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -3019,6 +3019,9 @@ class FastLlamaModel:
                         model.get_output_embeddings(), DEVICE_TYPE_TORCH
                     )
 
+                # 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)
                 return model
             else:
                 raise TypeError(
diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index 8790099fe3..cd5d498f08 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -596,6 +596,27 @@ def _patch_sft_trainer_auto_packing(trl_module):
             )
             print(message)
 
+        # get_peft_model installs a pre-train forward detector for plain LoRA/vision models,
+        # but only RL trainers run the reset via prepare_for_training_mode. Wire it into the
+        # SFT train() path too, else a grad-enabled probe before train() leaves the poisoned
+        # Dynamo cache in place and the detector hook installed on every training forward.
+        # (For UnslothSFTTrainer the later prepare_for_training_mode assignment supersedes this.)
+        if not getattr(self, "_unsloth_train_reset_wrapped", False):
+            try:
+                from unsloth.models.rl import _unsloth_reset_stray_compile_cache
+                _orig_train = self.train
+                @wraps(_orig_train)
+                def _train_with_reset(*train_args, **train_kwargs):
+                    try:
+                        _unsloth_reset_stray_compile_cache(self)
+                    except Exception:
+                        pass
+                    return _orig_train(*train_args, **train_kwargs)
+                self.train = _train_with_reset
+                self._unsloth_train_reset_wrapped = True
+            except Exception:
+                pass
+
     sft_trainer.__init__ = new_init
     sft_trainer._unsloth_auto_packing_wrapped = True
 

From 7a981a47084aa872bbf401db75513ed3d86b0c9b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 22 Jun 2026 10:09:18 +0000
Subject: [PATCH 07/11] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 unsloth/trainer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index cd5d498f08..62ebefd58a 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -604,7 +604,9 @@ def _patch_sft_trainer_auto_packing(trl_module):
         if not getattr(self, "_unsloth_train_reset_wrapped", False):
             try:
                 from unsloth.models.rl import _unsloth_reset_stray_compile_cache
+
                 _orig_train = self.train
+
                 @wraps(_orig_train)
                 def _train_with_reset(*train_args, **train_kwargs):
                     try:
@@ -612,6 +614,7 @@ def _patch_sft_trainer_auto_packing(trl_module):
                     except Exception:
                         pass
                     return _orig_train(*train_args, **train_kwargs)
+
                 self.train = _train_with_reset
                 self._unsloth_train_reset_wrapped = True
             except Exception:

From 202d731903d7ae38702f5276aabf88ed9f4f82e9 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 22 Jun 2026 12:15:06 +0000
Subject: [PATCH 08/11] Preserve detector evidence on reinstall + wire reset
 into plain Trainer path

P2 (_utils.py): _unsloth_install_pretrain_detector cleared marker['seen'] before the
live-hook early return, so a re-entrant get_peft_model/patch_peft_model after a
grad-enabled probe erased the recorded poisoning while leaving the hook installed,
and train() then skipped the Dynamo reset. Only reset seen when (re)installing a
fresh hook; keep it when a live hook is already recording.

P2 (llama.py): the detector is armed for every LoRA model, but only TRL SFT/RL train
wrappers consumed it. Inject _unsloth_reset_stray_compile_cache(self) at the start of
the generated _fast_inner_training_loop so a bare transformers.Trainer.train() also
drops a poisoned cache and tears down the hook. Idempotent with the TRL-wrapper reset.
---
 unsloth/models/_utils.py |  7 +++++--
 unsloth/models/llama.py  | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index fbdfb84602..bdc10e75bc 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -203,10 +203,13 @@ def _unsloth_install_pretrain_detector(model):
         return model
     marker = getattr(model, "_unsloth_pretrain_marker", None)
     if isinstance(marker, dict):
-        marker["seen"] = False
-        # Re-register only if the previous hook was torn down; a live hook stays (no duplicates).
+        # A live hook is already recording: keep it (no duplicates) and DON'T clear seen -- a
+        # grad-enabled probe may have already flagged the poisoned cache, and a re-entrant
+        # get_peft_model/patch_peft_model call must not erase that before train() resets.
         if "hook" in marker:
             return model
+        # Marker exists but its hook was torn down -> reinstall fresh, so reset seen.
+        marker["seen"] = False
     else:
         marker = {"seen": False}
         try:
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index baea9674d6..1983dc02cd 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2756,6 +2756,20 @@ class FastLlamaModel:
             "is_torch_tpu_available()",
             "False",
         )
+        # Wire the stray-forward compile-cache reset into the plain Trainer path: get_peft_model
+        # arms the pre-train detector for every LoRA model, but only the TRL SFT/RL wrappers run
+        # the reset. A grad-enabled probe before a bare transformers.Trainer.train() would
+        # otherwise keep the poisoned Dynamo cache and leave the detector hook installed. Anchored
+        # on the first body statement; a no-op (and harmless) if upstream drops that line.
+        inner_training_loop = inner_training_loop.replace(
+            "self.accelerator.free_memory()",
+            "self.accelerator.free_memory()\n"
+            "    try:\n"
+            "        from unsloth.models.rl import _unsloth_reset_stray_compile_cache as _unsloth_reset_cc\n"
+            "        _unsloth_reset_cc(self)\n"
+            "    except Exception: pass",
+            1,
+        )
         exec(inner_training_loop, globals())
         Trainer._inner_training_loop = _fast_inner_training_loop
 

From bc134a70c72a8409a5ddf9458ec8d5515aa2aeda Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 22 Jun 2026 12:32:53 +0000
Subject: [PATCH 09/11] Make _unsloth_reset_stray_compile_cache an importable
 module-level helper

The reset was only defined inside the RLTrainer_replacement template string, so
'from unsloth.models.rl import _unsloth_reset_stray_compile_cache' raised ImportError
(swallowed) on the SFT auto-packing wrapper and the injected plain-Trainer loop -
both paths kept the poisoned Dynamo cache and the dangling detector hook.

Move the canonical implementation to unsloth.models._utils (next to the detector,
exported in __all__). The RL trainer template now imports it (no-op fallback if the
import ever fails), and trainer.py / llama.py import it from _utils too, so every
training entry point actually runs the reset.
---
 unsloth/models/_utils.py | 63 ++++++++++++++++++++++++++++++++++++++++
 unsloth/models/llama.py  |  2 +-
 unsloth/models/rl.py     | 63 ++++------------------------------------
 unsloth/trainer.py       |  2 +-
 4 files changed, 71 insertions(+), 59 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index bdc10e75bc..1fcab917ba 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -63,6 +63,7 @@ __all__ = [
     "unpatch_unsloth_smart_gradient_checkpointing",
     "apply_unsloth_gradient_checkpointing",
     "_unsloth_install_pretrain_detector",
+    "_unsloth_reset_stray_compile_cache",
     "patch_compiled_autograd",
     "process_vision_info",
     "unsloth_compile_transformers",
@@ -230,6 +231,68 @@ def _unsloth_install_pretrain_detector(model):
     return model
 
 
+def _unsloth_reset_stray_compile_cache(self):
+    # A manual forward/backward under torch.compile BEFORE trainer.train() (e.g. a grad-norm
+    # probe) caches a forward + AOTAutograd backward graph in a one-off context; reusing it
+    # poisons training with NaN/zero gradients. If such a forward was seen and compile is on,
+    # drop the compiled-graph cache so training recompiles cleanly. No-op on the normal path.
+    # Module-level (not just inside the RL trainer template) so the SFT auto-packing wrapper and
+    # the plain-Trainer loop can import and run it too.
+    import os
+    model = getattr(self, "model", None)
+    if model is None:
+        return
+    # The detector hook can sit on any wrapper in the chain, and the probe may have run on a
+    # different one than self.model, so walk the chain: detect a "seen" marker anywhere and
+    # collect every marker to tear down below.
+    markers = []
+    seen = False
+    _curr = model
+    _visited = set()
+    while _curr is not None and id(_curr) not in _visited:
+        _visited.add(id(_curr))
+        _m = getattr(_curr, "_unsloth_pretrain_marker", None)
+        if isinstance(_m, dict):
+            markers.append(_m)
+            if _m.get("seen"):
+                seen = True
+        # Follow the wrapper chain: Unsloth/HF (.model), PEFT (.base_model), DDP/FSDP (.module).
+        _nxt = getattr(_curr, "model", None)
+        if _nxt is None:
+            _nxt = getattr(_curr, "base_model", None)
+        if _nxt is None:
+            _nxt = getattr(_curr, "module", None)
+        _curr = _nxt
+    if seen and os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1":
+        try:
+            import torch._dynamo as _dynamo
+            _dynamo.reset()
+        except Exception:
+            pass
+        try:
+            from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
+            reset_unsloth_gradient_checkpointing_buffers()
+        except Exception:
+            pass
+        try:
+            model.zero_grad(set_to_none = True)
+        except Exception:
+            pass
+        import warnings
+        warnings.warn(
+            "Unsloth: detected a manual forward/backward run before trainer.train(); "
+            "reset the torch.compile graph cache it poisoned so training starts clean. "
+            "To avoid this, run any pre-train probe under `with torch.no_grad():`."
+        )
+    # Tear down every one-shot detector hook in the chain so none adds per-step cost.
+    for _m in markers:
+        hook = _m.pop("hook", None)
+        if hook is not None:
+            try: hook.remove()
+            except Exception: pass
+        _m["seen"] = False
+
+
 def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_length, dtype):
     """
     Apply gradient checkpointing with smart heuristics.
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 1983dc02cd..3ef8b4c2e4 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2765,7 +2765,7 @@ class FastLlamaModel:
             "self.accelerator.free_memory()",
             "self.accelerator.free_memory()\n"
             "    try:\n"
-            "        from unsloth.models.rl import _unsloth_reset_stray_compile_cache as _unsloth_reset_cc\n"
+            "        from unsloth.models._utils import _unsloth_reset_stray_compile_cache as _unsloth_reset_cc\n"
             "        _unsloth_reset_cc(self)\n"
             "    except Exception: pass",
             1,
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 1a00f21416..a85c1d08a4 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -390,63 +390,12 @@ try:
     from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
 except:
     def reset_unsloth_gradient_checkpointing_buffers(): pass
-def _unsloth_reset_stray_compile_cache(self):
-    # A manual forward/backward under torch.compile BEFORE trainer.train() (e.g. a grad-norm
-    # probe) caches a forward + AOTAutograd backward graph in a one-off context; reusing it
-    # poisons training with NaN/zero gradients. If such a forward was seen and compile is on,
-    # drop the compiled-graph cache so training recompiles cleanly. No-op on the normal path.
-    import os
-    model = getattr(self, "model", None)
-    if model is None:
-        return
-    # The detector hook can sit on any wrapper in the chain, and the probe may have run on a
-    # different one than self.model, so walk the chain: detect a "seen" marker anywhere and
-    # collect every marker to tear down below.
-    markers = []
-    seen = False
-    _curr = model
-    _visited = set()
-    while _curr is not None and id(_curr) not in _visited:
-        _visited.add(id(_curr))
-        _m = getattr(_curr, "_unsloth_pretrain_marker", None)
-        if isinstance(_m, dict):
-            markers.append(_m)
-            if _m.get("seen"):
-                seen = True
-        # Follow the wrapper chain: Unsloth/HF (.model), PEFT (.base_model), DDP/FSDP (.module).
-        _nxt = getattr(_curr, "model", None)
-        if _nxt is None:
-            _nxt = getattr(_curr, "base_model", None)
-        if _nxt is None:
-            _nxt = getattr(_curr, "module", None)
-        _curr = _nxt
-    if seen and os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1":
-        try:
-            import torch._dynamo as _dynamo
-            _dynamo.reset()
-        except Exception:
-            pass
-        try:
-            reset_unsloth_gradient_checkpointing_buffers()
-        except Exception:
-            pass
-        try:
-            model.zero_grad(set_to_none = True)
-        except Exception:
-            pass
-        import warnings
-        warnings.warn(
-            "Unsloth: detected a manual forward/backward run before trainer.train(); "
-            "reset the torch.compile graph cache it poisoned so training starts clean. "
-            "To avoid this, run any pre-train probe under `with torch.no_grad():`."
-        )
-    # Tear down every one-shot detector hook in the chain so none adds per-step cost.
-    for _m in markers:
-        hook = _m.pop("hook", None)
-        if hook is not None:
-            try: hook.remove()
-            except Exception: pass
-        _m["seen"] = False
+# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain
+# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported.
+try:
+    from unsloth.models._utils import _unsloth_reset_stray_compile_cache
+except Exception:
+    def _unsloth_reset_stray_compile_cache(self): pass
 def prepare_for_training_mode(f):
     @functools.wraps(f)
     def wrapper(self, *args, **kwargs):
diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index 62ebefd58a..83cb1758f0 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -603,7 +603,7 @@ def _patch_sft_trainer_auto_packing(trl_module):
         # (For UnslothSFTTrainer the later prepare_for_training_mode assignment supersedes this.)
         if not getattr(self, "_unsloth_train_reset_wrapped", False):
             try:
-                from unsloth.models.rl import _unsloth_reset_stray_compile_cache
+                from unsloth.models._utils import _unsloth_reset_stray_compile_cache
 
                 _orig_train = self.train
 

From de6c777442bb0a35ebe0a27cbd992bd782ab3929 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 22 Jun 2026 12:33:26 +0000
Subject: [PATCH 10/11] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 unsloth/models/_utils.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 1fcab917ba..2365975cdd 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -239,6 +239,7 @@ def _unsloth_reset_stray_compile_cache(self):
     # Module-level (not just inside the RL trainer template) so the SFT auto-packing wrapper and
     # the plain-Trainer loop can import and run it too.
     import os
+
     model = getattr(self, "model", None)
     if model is None:
         return
@@ -270,7 +271,9 @@ def _unsloth_reset_stray_compile_cache(self):
         except Exception:
             pass
         try:
-            from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
+            from unsloth_zoo.gradient_checkpointing import (
+                reset_unsloth_gradient_checkpointing_buffers,
+            )
             reset_unsloth_gradient_checkpointing_buffers()
         except Exception:
             pass
@@ -279,6 +282,7 @@ def _unsloth_reset_stray_compile_cache(self):
         except Exception:
             pass
         import warnings
+
         warnings.warn(
             "Unsloth: detected a manual forward/backward run before trainer.train(); "
             "reset the torch.compile graph cache it poisoned so training starts clean. "
@@ -288,8 +292,10 @@ def _unsloth_reset_stray_compile_cache(self):
     for _m in markers:
         hook = _m.pop("hook", None)
         if hook is not None:
-            try: hook.remove()
-            except Exception: pass
+            try:
+                hook.remove()
+            except Exception:
+                pass
         _m["seen"] = False
 
 

From 83900ffb3a69f18b5a899499fcc198206e8d5ef7 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 22 Jun 2026 12:42:13 +0000
Subject: [PATCH 11/11] Add regression tests for the stray-forward
 compile-cache reset

Covers the two codex findings fixed in this PR under the GPU-free
tests/conftest.py harness:

- the reset helper 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
- _unsloth_reset_stray_compile_cache 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
---
 tests/test_pretrain_compile_reset.py | 149 +++++++++++++++++++++++++++
 1 file changed, 149 insertions(+)
 create mode 100644 tests/test_pretrain_compile_reset.py

diff --git a/tests/test_pretrain_compile_reset.py b/tests/test_pretrain_compile_reset.py
new file mode 100644
index 0000000000..c16ba37a3b
--- /dev/null
+++ b/tests/test_pretrain_compile_reset.py
@@ -0,0 +1,149 @@
+# 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():
+    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():
+    # 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).
+    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