Reset torch.compile cache poisoned by a stray forward before trainer.train() (#6511)

* 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.

* 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.

* 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.

* 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.

* torch.compile stray-forward reset: tighten comments (no code change)

* 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.

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

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

* 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.

* 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.

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

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

---------

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

unsloth/models/_utils.py

@@ -62,6 +62,8 @@ __all__ = [
     "patch_unsloth_smart_gradient_checkpointing",
     "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",
@@ -194,6 +196,109 @@ from unsloth_zoo.temporary_patches import (
 )
 
 
+def _unsloth_install_pretrain_detector(model):
+    """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):
+        # 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:
+            model._unsloth_pretrain_marker = marker
+        except Exception:
+            return model
+
+    def _mark(_module, _inp):
+        # 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
+
+    try:
+        marker["hook"] = model.register_forward_pre_hook(_mark)
+    except Exception:
+        pass
+    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.

unsloth/models/llama.py

@@ -2762,6 +2762,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._utils 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
 
@@ -2944,6 +2958,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)
 
@@ -3022,6 +3039,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(
@@ -3374,6 +3394,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
@@ -3591,6 +3614,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

unsloth/models/rl.py

@@ -390,9 +390,20 @@ try:
     from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
 except:
     def reset_unsloth_gradient_checkpointing_buffers(): pass
+# 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):
+        # 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.

unsloth/models/vision.py

@@ -1386,6 +1386,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)
 
@@ -1583,6 +1586,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

unsloth/trainer.py

@@ -596,6 +596,30 @@ 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._utils 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