chore: import upstream snapshot with attribution

swift/tuner_plugin/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .base import PeftTuner, Tuner
+from .mapping import tuners_map

swift/tuner_plugin/base.py

@@ -0,0 +1,80 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from peft import PeftModel
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from swift.arguments import SftArguments
+
+
+class Tuner:
+    """Base class for model tuners that adapt pre-trained models for specific tasks."""
+
+    @staticmethod
+    def prepare_model(args: 'SftArguments', model: torch.nn.Module) -> torch.nn.Module:
+        """Prepare a new model with a tuner.
+
+        Args:
+            args: The training arguments containing tuner configuration.
+            model: The model instance to be wrapped.
+
+        Returns:
+            The wrapped model with tuner applied.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def save_pretrained(
+        model: torch.nn.Module,
+        save_directory: str,
+        state_dict: Optional[dict] = None,
+        safe_serialization: bool = True,
+        **kwargs,
+    ) -> None:
+        """Save the model checkpoint.
+
+        Args:
+            model: The wrapped model by `prepare_model`.
+            save_directory: The directory path where the model will be saved.
+            state_dict: The model's state_dict, used during DeepSpeed training.
+                Only contains trainable parameters
+            safe_serialization: Whether to use safetensors format for serialization. Defaults to True.
+            **kwargs: Additional keyword arguments for saving.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def from_pretrained(model: torch.nn.Module, model_id: str, **kwargs) -> torch.nn.Module:
+        """Load a model from a checkpoint directory.
+
+        Args:
+            model: The original model instance.
+            model_id: The model identifier or checkpoint directory path to load from.
+            **kwargs: Additional keyword arguments for loading.
+
+        Returns:
+            The wrapped model instance with loaded weights.
+        """
+        raise NotImplementedError
+
+
+class PeftTuner(Tuner):
+    """Tuner implementation using the PEFT library."""
+
+    @staticmethod
+    def save_pretrained(
+        model: torch.nn.Module,
+        save_directory: str,
+        state_dict: Optional[dict] = None,
+        safe_serialization: bool = True,
+        **kwargs,
+    ) -> None:
+        """Save the PEFT model checkpoint."""
+        if isinstance(model, PeftModel):
+            if 'selected_adapters' not in kwargs:
+                kwargs['selected_adapters'] = ['default']
+        model.save_pretrained(save_directory, safe_serialization=safe_serialization, **kwargs)
+
+    @staticmethod
+    def from_pretrained(model: torch.nn.Module, model_id: str, **kwargs) -> torch.nn.Module:
+        return PeftModel.from_pretrained(model, model_id, **kwargs)

swift/tuner_plugin/dummy.py

@@ -0,0 +1,15 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from typing import TYPE_CHECKING
+
+from .base import PeftTuner
+
+if TYPE_CHECKING:
+    from swift.arguments import SftArguments
+
+
+class DummyTuner(PeftTuner):
+
+    @staticmethod
+    def prepare_model(args: 'SftArguments', model: torch.nn.Module) -> torch.nn.Module:
+        return model

swift/tuner_plugin/ia3.py

@@ -0,0 +1,22 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from peft import IA3Config, get_peft_model
+from typing import TYPE_CHECKING
+
+from swift.model import ModelKeys
+from swift.utils import find_all_linears
+from .base import PeftTuner
+
+if TYPE_CHECKING:
+    from swift.arguments import SftArguments
+
+
+# Here gives a simple example of IA3
+class IA3Tuner(PeftTuner):
+
+    @staticmethod
+    def prepare_model(args: 'SftArguments', model: torch.nn.Module) -> torch.nn.Module:
+        model_arch: ModelKeys = model.model_meta.model_arch
+        ia3_config = IA3Config(
+            target_modules=find_all_linears(model), feedforward_modules='.*' + model_arch.mlp.split('{}.')[1] + '.*')
+        return get_peft_model(model, ia3_config)

swift/tuner_plugin/lora_llm.py

@@ -0,0 +1,78 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import safetensors.torch
+import torch
+from peft import LoraConfig, PeftModel, get_peft_model
+from transformers.integrations import is_deepspeed_zero3_enabled
+from typing import TYPE_CHECKING, Optional
+
+from swift.utils import deep_getattr, get_logger, get_multimodal_target_regex
+from .base import Tuner
+
+logger = get_logger()
+
+if TYPE_CHECKING:
+    from swift.arguments import SftArguments
+
+
+def is_vit_aligner_param(model_arch, parameter_name: str) -> bool:
+    for module_prefix in model_arch.vision_tower + model_arch.aligner:
+        if f'.{module_prefix}.' in parameter_name:
+            return True
+    return False
+
+
+class LoRALLMTuner(Tuner):
+    """Full-parameter training of ViT/Aligner while LoRA training LLM"""
+
+    @staticmethod
+    def from_pretrained(model: torch.nn.Module, model_id: str, **kwargs) -> torch.nn.Module:
+        model = PeftModel.from_pretrained(model, model_id, **kwargs)
+        state_dict = safetensors.torch.load_file(os.path.join(model_id, 'vit.safetensors'))
+        if is_deepspeed_zero3_enabled():
+            import deepspeed
+            params_dict = dict(model.named_parameters())
+            params_to_load = {name: params_dict[name] for name in state_dict if name in params_dict}
+            if params_to_load:
+                with deepspeed.zero.GatheredParameters(list(params_to_load.values()), modifier_rank=0):
+                    if deepspeed.comm.get_rank() == 0:
+                        for name, param in params_to_load.items():
+                            param.data.copy_(state_dict[name])
+        else:
+            model.load_state_dict(state_dict, strict=False)
+        model_arch = model.model_meta.model_arch
+        for module_prefix in model_arch.vision_tower + model_arch.aligner:
+            deep_getattr(model, module_prefix).requires_grad_(True)
+        return model
+
+    @staticmethod
+    def save_pretrained(
+        model: torch.nn.Module,
+        save_directory: str,
+        state_dict: Optional[dict] = None,
+        safe_serialization: bool = True,
+        **kwargs,
+    ) -> None:
+        if state_dict is None:
+            state_dict = {}
+            for n, p in model.named_parameters():
+                if p.requires_grad:
+                    state_dict[n] = p.detach().cpu()
+        model.save_pretrained(save_directory, state_dict=state_dict, safe_serialization=safe_serialization, **kwargs)
+        # vit/aligner
+        model_arch = model.model_meta.model_arch
+        state_dict = {k: v for k, v in state_dict.items() if is_vit_aligner_param(model_arch, k)}
+        safetensors.torch.save_file(
+            state_dict, os.path.join(save_directory, 'vit.safetensors'), metadata={'format': 'pt'})
+
+    @staticmethod
+    def prepare_model(args: 'SftArguments', model: torch.nn.Module) -> torch.nn.Module:
+        model_arch = model.model_meta.model_arch
+        target_regex = get_multimodal_target_regex(model)
+        logger.info(f'target_regex: {target_regex}')
+        lora_config = LoraConfig(
+            task_type=args.task_type.upper(), r=args.lora_rank, lora_alpha=args.lora_alpha, target_modules=target_regex)
+        model = get_peft_model(model, lora_config)
+        for module_prefix in model_arch.vision_tower + model_arch.aligner:
+            deep_getattr(model, module_prefix).requires_grad_(True)
+        return model

swift/tuner_plugin/mapping.py

@@ -0,0 +1,10 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .dummy import DummyTuner
+from .ia3 import IA3Tuner
+from .lora_llm import LoRALLMTuner
+
+tuners_map = {
+    'ia3': IA3Tuner,
+    'lora_llm': LoRALLMTuner,
+    'dummy': DummyTuner,
+}