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 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 08802f030e..3ef8b4c2e4 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._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 @@ -2938,6 +2952,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) @@ -3016,6 +3033,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( @@ -3368,6 +3388,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 +3608,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 3f8d9e6ce0..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) @@ -1577,6 +1580,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