diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c69847c048..baea9674d6 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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( diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 8790099fe3..cd5d498f08 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -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