merge model part to the text stream

src/llmtuner/model/adapter.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 import torch
 from peft import LoraConfig, LoraModel, PeftModel, TaskType, get_peft_model
@@ -21,11 +21,11 @@ logger = get_logger(__name__)
 
 def init_adapter(
     config: "PretrainedConfig",
-    model: "PreTrainedModel",
+    model: Union["PreTrainedModel","AutoModelForVision2Seq"],
     model_args: "ModelArguments",
     finetuning_args: "FinetuningArguments",
     is_trainable: bool,
-) -> "PreTrainedModel":
+) -> Union["PreTrainedModel","AutoModelForVision2Seq"]:
     r"""
     Initializes the adapters.
 
@@ -195,103 +195,4 @@ def init_adapter(
         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
-
-
-def init_mm_adapter(
-        model: "AutoModelForVision2Seq", model_args: "ModelArguments",
-        finetuning_args: "FinetuningArguments",
-        is_trainable: bool,
-        use_clm=True,
-) -> "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:
-                if use_clm:
-                    lora_config = LoraConfig(
-                        task_type=TaskType.CAUSAL_LM,
-                        inference_mode=False,
-                        use_dora=finetuning_args.use_dora,
-                        **peft_kwargs,
-                    )
-                else:
-                    lora_config = LoraConfig(
-                        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
+    return model