From bc134a70c72a8409a5ddf9458ec8d5515aa2aeda Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 22 Jun 2026 12:32:53 +0000 Subject: [PATCH] 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