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:
parent
d2f6eec004
commit
430af1776b
2 changed files with 24 additions and 0 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue