diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9537dd63c3..caf259c0bd 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2046,6 +2046,11 @@ def unsloth_fast_generate( class FastLlamaModel: + @staticmethod + def _prepare_for_qat(model, qat_scheme): + model = _prepare_model_for_qat(model, qat_scheme) + return model + @staticmethod def pre_patch(): init_name, function = patch_llama_rope_scaling( @@ -3007,7 +3012,8 @@ class FastLlamaModel: # Apply QAT + LoRA if specified if qat_scheme is not None: print("Unsloth: Applying QAT to mitigate quantization degradation") - model = _prepare_model_for_qat(model, qat_scheme) + model = FastLlamaModel._prepare_for_qat(model, qat_scheme) + model._saved_temp_tokenizer = _saved_temp_tokenizer diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 15f01cf15d..9d8a2ce9fa 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -595,6 +595,11 @@ except: class FastModel(FastBaseModel): + @staticmethod + def _prepare_for_qat(model, qat_scheme): + model = _prepare_model_for_qat(model, qat_scheme) + return model + @staticmethod def from_pretrained( model_name = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit", @@ -1144,7 +1149,7 @@ class FastModel(FastBaseModel): # Apply QAT if specified if qat_scheme is not None: print("Unsloth: Applying QAT to mitigate quantization degradation") - model = _prepare_model_for_qat(model, qat_scheme) + model = FastModel._prepare_for_qat(model, qat_scheme) # Patch Tiled MLP # to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}""