diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index eafbd5a43..b72d8defa 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -68,11 +68,18 @@ LlamaFlashAttention2 = LlamaAttention pass -from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, BitsAndBytesConfig, AutoConfig +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + AutoModelForSequenceClassification, + AutoModelForSeq2SeqLM, + BitsAndBytesConfig, + AutoConfig, +) from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING from transformers import set_seed as transformers_set_seed from peft import LoraConfig, TaskType, get_peft_model as _get_peft_model -from peft import PeftModelForCausalLM, PeftModelForSequenceClassification +from peft import PeftModelForCausalLM, PeftModelForSequenceClassification, PeftModelForSeq2SeqLM from ..save import patch_saving_functions import re, os, inspect, math, sys import types @@ -2292,7 +2299,7 @@ def get_peft_model( if r <= 0: raise TypeError(f"Unsloth: Rank of {str(r)} must be larger than 0.") - if isinstance(model, PeftModelForCausalLM) or isinstance(model, PeftModelForSequenceClassification): + if any(functools.partial(isinstance, model), (PeftModelForCausalLM, PeftModelForSequenceClassification, PeftModelForSeq2SeqLM)): # Check if exactly the same and then pass through! assert(hasattr(model, "peft_config")) @@ -2557,8 +2564,14 @@ def get_peft_model( raise NotImplementedError("Unsloth: Currently fast inference does not work with using biases for LoRA.") pass - #d oes not get lora yet, so get name from model, not base model - is_classification = "Classification" in str(type(model)) + # does not get lora yet, so get name from model, not base model + model_type = type(model) + if model_type in AutoModelForSeq2SeqLM._model_mapping.values(): + task_type = TaskType.SEQ_2_SEQ_LM + elif model_type in AutoModelForSequenceClassification._model_mapping.values(): + task_type = TaskType.SEQ_CLS + else: + task_type = TaskType.CAUSAL_LM arguments = dict( r = r, @@ -2566,7 +2579,7 @@ def get_peft_model( target_modules = final_modules, lora_dropout = lora_dropout, bias = bias, - task_type = TaskType.CAUSAL_LM if not is_classification else TaskType.SEQ_CLS, + task_type = task_type, layers_to_transform = layers_to_transform, init_lora_weights = init_lora_weights, loftq_config = loftq_config, @@ -2701,7 +2714,7 @@ def patch_peft_model( use_gradient_checkpointing = use_gradient_checkpointing, ) pass - if not isinstance(model, PeftModelForCausalLM) and not isinstance(model, PeftModelForSequenceClassification): + if not any(functools.partial(isinstance, model), (PeftModelForCausalLM, PeftModelForSequenceClassification, PeftModelForSeq2SeqLM)): raise TypeError( "Unsloth: Your model needs to call `.get_peft_model` first!" ) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 7ac27158a..59fc1e0cf 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -479,6 +479,7 @@ def from_pretrained( from .vision import FastBaseModel from transformers import ( AutoModelForCausalLM, + AutoModelForSeq2SeqLM, ) try: from transformers import AutoModelForImageTextToText @@ -824,7 +825,12 @@ def from_pretrained( is_vlm = any(x.endswith("ForConditionalGeneration") for x in model_config.architectures) is_vlm = is_vlm or hasattr(model_config, "vision_config") if auto_model is None: - auto_model = AutoModelForVision2Seq if is_vlm else AutoModelForCausalLM + if AutoModelForSeq2SeqLM._model_mapping.get(type(model_config), None) is not None: + auto_model = AutoModelForSeq2SeqLM + elif is_vlm: + auto_model = AutoModelForVision2Seq + else: + auto_model = AutoModelForCausalLM model, tokenizer = FastBaseModel.from_pretrained( model_name = model_name,