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