From 9dbd40e5b31fdb65fc78d4a17eaf3ceeceea0840 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 22 Jun 2026 05:39:48 -0700 Subject: [PATCH] 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 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/_utils.py | 105 +++++++++++++++++++++++++++++++++++++++ unsloth/models/llama.py | 26 ++++++++++ unsloth/models/rl.py | 11 ++++ unsloth/models/vision.py | 6 +++ unsloth/trainer.py | 24 +++++++++ 5 files changed, 172 insertions(+) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 742f5f57e0..2365975cdd 100644 --- a/unsloth/models/_utils.py +++ b/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. diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 2f449453f9..7b9cf3df0e 100644 --- a/unsloth/models/llama.py +++ b/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 diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index c3237b6481..a85c1d08a4 100644 --- a/unsloth/models/rl.py +++ b/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. diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index b68f0361e3..182030c68b 100644 --- a/unsloth/models/vision.py +++ b/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 diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 8790099fe3..83cb1758f0 100644 --- a/unsloth/trainer.py +++ b/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