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.
This commit is contained in:
Daniel Han 2026-06-22 10:08:51 +00:00
commit 430af1776b
2 changed files with 24 additions and 0 deletions

View file

@ -3019,6 +3019,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(

View file

@ -596,6 +596,27 @@ 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.rl 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