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.
This commit is contained in:
Daniel Han
parent
202d731903
commit
bc134a70c7
4 changed files with 71 additions and 59 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue