File size: 5,010 Bytes
96fe658 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
# Copyright (c) Alibaba, Inc. and its affiliates.
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)
|