add multimodal LLM BLIP-2 and InstructBLIP

src/llmtuner/model/adapter.py

@@ -1,24 +1,25 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 import torch
 from peft import LoraConfig, LoraModel, PeftModel, TaskType, get_peft_model
+from transformers import AutoModelForVision2Seq
 from transformers.integrations import is_deepspeed_zero3_enabled
 
 from ..extras.logging import get_logger
 from .utils import QuantizationMethod, find_all_linear_modules, find_expanded_modules
 
-
 if TYPE_CHECKING:
-    from transformers.modeling_utils import PreTrainedModel
+    from transformers.modeling_utils import PreTrainedModel, AutoModelForVision2Seq
 
     from ..hparams import FinetuningArguments, ModelArguments
 
-
 logger = get_logger(__name__)
 
 
 def init_adapter(
-    model: "PreTrainedModel", model_args: "ModelArguments", finetuning_args: "FinetuningArguments", is_trainable: bool
+        model: "PreTrainedModel", model_args: "ModelArguments",
+        finetuning_args: "FinetuningArguments",
+        is_trainable: bool
 ) -> "PreTrainedModel":
     r"""
     Initializes the adapters.
@@ -43,9 +44,9 @@ def init_adapter(
     if finetuning_args.finetuning_type == "freeze" and is_trainable:
         logger.info("Fine-tuning method: Freeze")
         num_layers = (
-            getattr(model.config, "num_hidden_layers", None)
-            or getattr(model.config, "num_layers", None)
-            or getattr(model.config, "n_layer", None)
+                getattr(model.config, "num_hidden_layers", None)
+                or getattr(model.config, "num_layers", None)
+                or getattr(model.config, "n_layer", None)
         )
         if not num_layers:
             raise ValueError("Current model does not support freeze tuning.")
@@ -135,9 +136,9 @@ def init_adapter(
                 target_modules = find_expanded_modules(model, target_modules, finetuning_args.num_layer_trainable)
 
             if (
-                finetuning_args.use_dora
-                and getattr(model, "quantization_method", None) is not None
-                and getattr(model, "quantization_method", None) != QuantizationMethod.BITS_AND_BYTES
+                    finetuning_args.use_dora
+                    and getattr(model, "quantization_method", None) is not None
+                    and getattr(model, "quantization_method", None) != QuantizationMethod.BITS_AND_BYTES
             ):
                 raise ValueError("DoRA is not compatible with PTQ-quantized models.")
 
@@ -176,3 +177,94 @@ def init_adapter(
             logger.info("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))
 
     return model
+
+
+def init_mm_adapter(
+        model: "AutoModelForVision2Seq", model_args: "ModelArguments",
+        finetuning_args: "FinetuningArguments",
+        is_trainable: bool
+) -> "AutoModelForVision2Seq":
+    if finetuning_args.finetuning_type == "lora":
+        logger.info("Fine-tuning method: {}".format("DoRA" if finetuning_args.use_dora else "LoRA"))
+        adapter_to_resume = None
+
+        if model_args.adapter_name_or_path is not None:
+            is_mergeable = True
+            if getattr(model, "quantization_method", None):  # merge lora in quantized model is unstable
+                assert len(model_args.adapter_name_or_path) == 1, "Quantized model only accepts a single adapter."
+                is_mergeable = False
+
+            if is_deepspeed_zero3_enabled():
+                assert len(model_args.adapter_name_or_path) == 1, "Cannot use multiple adapters in DeepSpeed ZeRO-3."
+                is_mergeable = False
+
+            if (is_trainable and not finetuning_args.create_new_adapter) or (not is_mergeable):
+                adapter_to_merge = model_args.adapter_name_or_path[:-1]
+                adapter_to_resume = model_args.adapter_name_or_path[-1]
+            else:
+                adapter_to_merge = model_args.adapter_name_or_path
+
+            for adapter in adapter_to_merge:
+                model: "LoraModel" = PeftModel.from_pretrained(
+                    model, adapter, offload_folder=model_args.offload_folder
+                )
+                model = model.merge_and_unload()
+
+            if len(adapter_to_merge) > 0:
+                logger.info("Merged {} adapter(s).".format(len(adapter_to_merge)))
+
+            if adapter_to_resume is not None:  # resume lora training
+                model = PeftModel.from_pretrained(
+                    model, adapter_to_resume, is_trainable=is_trainable, offload_folder=model_args.offload_folder
+                )
+
+        if is_trainable and adapter_to_resume is None:  # create new lora weights while training
+            if len(finetuning_args.lora_target) == 1 and finetuning_args.lora_target[0] == "all":
+                target_modules = find_all_linear_modules(model)
+            else:
+                target_modules = finetuning_args.lora_target
+
+            if finetuning_args.use_llama_pro:
+                target_modules = find_expanded_modules(model, target_modules, finetuning_args.num_layer_trainable)
+
+            if (
+                    finetuning_args.use_dora
+                    and getattr(model, "quantization_method", None) is not None
+                    and getattr(model, "quantization_method", None) != QuantizationMethod.BITS_AND_BYTES
+            ):
+                raise ValueError("DoRA is not compatible with PTQ-quantized models.")
+
+            peft_kwargs = {
+                "r": finetuning_args.lora_rank,
+                "target_modules": target_modules,
+                "lora_alpha": finetuning_args.lora_alpha,
+                "lora_dropout": finetuning_args.lora_dropout,
+                "use_rslora": finetuning_args.use_rslora,
+                "modules_to_save": finetuning_args.additional_target,
+            }
+
+            if model_args.use_unsloth:
+                from unsloth import FastLanguageModel  # type: ignore
+
+                unsloth_peft_kwargs = {
+                    "model": model,
+                    "max_seq_length": model_args.model_max_length,
+                    "use_gradient_checkpointing": "unsloth",
+                }
+                model = FastLanguageModel.get_peft_model(**peft_kwargs, **unsloth_peft_kwargs)
+            else:
+                lora_config = LoraConfig(
+                    # task_type=TaskType.CAUSAL_LM,
+                    inference_mode=False,
+                    use_dora=finetuning_args.use_dora,
+                    **peft_kwargs,
+                )
+                model = get_peft_model(model, lora_config)
+
+        if (not finetuning_args.pure_bf16) and (not finetuning_args.use_badam):
+            for param in filter(lambda p: p.requires_grad, model.parameters()):
+                param.data = param.data.to(torch.float32)
+
+        if model_args.adapter_name_or_path is not None:
+            logger.info("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))
+    return model