wehub-resource-sync
119 lines
4.9 KiB
Python
119 lines
4.9 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import re
|
|
import torch
|
|
from copy import deepcopy
|
|
from dataclasses import dataclass
|
|
from torch import nn
|
|
from types import MethodType
|
|
from typing import Dict, Optional
|
|
|
|
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)
|