| |
| import re |
| from copy import deepcopy |
| from dataclasses import dataclass |
| from types import MethodType |
| from typing import Dict, Optional |
|
|
| import torch |
| from torch import nn |
|
|
| from swift.utils import get_logger |
| from .utils import ActivationMixin, SwiftAdapter, SwiftConfig, SwiftOutput |
|
|
| logger = get_logger() |
|
|
|
|
| @dataclass |
| class PartConfig(SwiftConfig): |
| """ |
| Freeze the model and train a part of it. |
| |
| Args: |
| target_modules(`Optional[str]`): The target modules to be trained in regex format |
| """ |
|
|
| target_modules: Optional[str] = None |
|
|
| def __post_init__(self): |
| from .mapping import SwiftTuners |
| self.swift_type = SwiftTuners.PART |
|
|
|
|
| class Part(SwiftAdapter): |
|
|
| @staticmethod |
| def target_module_matched(module_key: str, config: PartConfig): |
| return re.fullmatch(config.target_modules, module_key) |
|
|
| @staticmethod |
| def prepare_model(model: nn.Module, config: PartConfig, adapter_name: str): |
| name_list = [name for name, _ in model.named_modules(remove_duplicate=False)] |
| for name in name_list: |
| module: nn.Module = model.get_submodule(name) |
| if Part.target_module_matched(name, config) and not getattr(module, 'plugin', False): |
| if hasattr(module, 'base_layer'): |
| module = module.base_layer |
|
|
| def _forward(self, *args, **kwargs): |
| child_list = [ |
| sub_module for name, sub_module in self.named_modules(remove_duplicate=False) |
| if '_part_' in name |
| ] |
| sub_modules = [child for child in child_list if getattr(child, 'activated', False)] |
| assert len(sub_modules) <= 1 |
| if len(sub_modules) == 1: |
| return sub_modules[0].forward(*args, **kwargs) |
| else: |
| return self.forward_origin(*args, **kwargs) |
|
|
| if not hasattr(module, 'forward_origin'): |
| module.forward_origin = module.forward |
| module.forward = MethodType(_forward, module) |
|
|
| new_module = deepcopy(module) |
| for attr in dir(new_module): |
| if '_part_' in attr: |
| delattr(new_module, attr) |
| new_module.part_name = adapter_name |
| ActivationMixin.mark_all_sub_modules_as_plugin(new_module) |
| setattr(module, f'_part_{adapter_name}', new_module) |
| new_module.requires_grad_(True) |
|
|
| def state_dict_callback(state_dict, adapter_name, **kwargs): |
| new_state_dict = {} |
| for key, value in state_dict.items(): |
| if f'_part_{adapter_name}.' in key: |
| if kwargs.get('replace_key', True): |
| new_key = key.replace(f'_part_{adapter_name}.', '').replace('base_layer.', '') |
| else: |
| new_key = key |
| new_state_dict[new_key] = value |
|
|
| return new_state_dict |
|
|
| def mark_trainable_callback(model: nn.Module): |
| pass |
|
|
| def load_state_dict_callback(model: nn.Module, adapter_name: str, state_dict: Dict[str, torch.Tensor]): |
| new_state_dict = {} |
| for name, module in model.named_modules(remove_duplicate=False): |
| module: nn.Module |
| if Part.target_module_matched(name, config): |
| for param_name in state_dict: |
| if param_name.startswith(name): |
| end = param_name[len(name):] |
| if '_part_' not in param_name: |
| if hasattr(module, 'base_layer'): |
| new_state_dict[name + f'.base_layer._part_{adapter_name}' |
| + end] = state_dict[param_name] |
| else: |
| new_state_dict[name + f'._part_{adapter_name}' + end] = state_dict[param_name] |
| else: |
| new_state_dict[param_name] = state_dict[param_name] |
| return new_state_dict |
|
|
| return SwiftOutput( |
| config=config, |
| state_dict_callback=state_dict_callback, |
| mark_trainable_callback=mark_trainable_callback, |
| load_state_dict_callback=load_state_dict_callback) |
|
|
| @staticmethod |
| def activate_adapter(module: torch.nn.Module, adapter_name: str, activate: bool, offload: str = None): |
| name_list = [name for name, _ in module.named_modules(remove_duplicate=False)] |
| for name in name_list: |
| sub_module: nn.Module = module.get_submodule(name) |
| if re.fullmatch(f'.*_part_{adapter_name}$', name): |
| sub_module.activated = activate |
| SwiftAdapter.save_memory(sub_module, adapter_name, name, activate, offload) |
|
|
