Support for Seq2Seq Models (T5, T5Gemma, etc.)

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,16 +2564,22 @@ 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,
             lora_alpha          = lora_alpha,
             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!"
             )

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,