|
|
|
|
|
import copy |
|
|
import re |
|
|
import types |
|
|
from dataclasses import dataclass, field |
|
|
from typing import Dict, List, Optional, Union |
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from swift.utils import get_logger |
|
|
from swift.utils.torch_utils import find_sub_module |
|
|
from .restuning_components import ResTuner, detach_tensors, probe_input_pre_hook, probe_output_hook |
|
|
from .utils import ActivationMixin, SwiftAdapter, SwiftConfig, SwiftOutput |
|
|
|
|
|
logger = get_logger() |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class ResTuningConfig(SwiftConfig): |
|
|
""" |
|
|
The configuration class for the ResTuning module. |
|
|
|
|
|
ResTuning is a flexible parameter-efficient and memory-efficient tuning paradigm framework. |
|
|
'Res-Tuning: A Flexible and Efficient Tuning Paradigm via Unbinding Tuner from Backbone' |
|
|
by Jiang et al.(2023) |
|
|
See |
|
|
|
|
|
Args: |
|
|
dims(`Union[List[int], int]`): The dimensions of the hidden states |
|
|
root_modules(`str`): The root module to be replaced, can a regex string |
|
|
root_modules_hook(`str`): The hook type of root modules, can be "input" or "output" |
|
|
stem_modules(`Union[List[str], str]`): The stem modules to be replaced, |
|
|
can a regex string or name list of full match format |
|
|
stem_modules_hook(`Union[List[str], str]`): The hook type of stem modules, can be "input" or "output" |
|
|
target_modules(`str`): The target module to be replaced, can a regex string |
|
|
target_modules_hook(`str`): The hook type of target modules, can be "input" or "output" |
|
|
tuner_cfg(`Union[List[Dict], Dict, str]`): The configuration of the tuning module, |
|
|
can a string or customized config |
|
|
use_upsample(bool): Whether to use auxiliary upsample module |
|
|
upsample_out_channels(List[int]): The channels if `use_upsample` |
|
|
zero_init_last(bool): Use zero to initialize the last Linear in every sub tuner. |
|
|
|
|
|
""" |
|
|
|
|
|
dims: Optional[Union[List[int], int]] = field( |
|
|
default=None, metadata={'help': 'The dimensions of the hidden states'}) |
|
|
|
|
|
root_modules: str = field( |
|
|
default=None, |
|
|
metadata={ |
|
|
'help': |
|
|
'The root module to be replaced, can a regex string (use the first matching module) or full match format' |
|
|
}) |
|
|
|
|
|
root_modules_hook: str = field( |
|
|
default='input', metadata={'help': 'The hook type of root modules, can be "input" or "output"'}) |
|
|
|
|
|
stem_modules: Optional[Union[List[str], str]] = field( |
|
|
default=None, |
|
|
metadata={'help': 'The stem modules to be replaced, can a regex string or name list of full match format'}) |
|
|
|
|
|
stem_modules_hook: str = field( |
|
|
default='output', metadata={'help': 'The hook type of stem modules, can be "input" or "output"'}) |
|
|
|
|
|
target_modules: str = field( |
|
|
default=None, |
|
|
metadata={ |
|
|
'help': |
|
|
'The target module to be replaced, can a regex string (use the first matching module) or full match format' |
|
|
}) |
|
|
|
|
|
target_modules_hook: str = field( |
|
|
default='input', metadata={'help': 'The hook type of target modules, can be "input" or "output"'}) |
|
|
|
|
|
target_hidden_pos: Union[int, str] = field( |
|
|
default=None, metadata={'help': 'The position of the hidden state for target modules output'}) |
|
|
|
|
|
tuner_cfg: Optional[Union[List[Dict], Dict, str]] = field( |
|
|
default=None, metadata={'help': 'The configuration of the tuning module, can a string or customized config'}) |
|
|
|
|
|
use_upsample: bool = field(default=False, metadata={'help': 'Whether to use auxiliary upsample module'}) |
|
|
|
|
|
upsample_out_channels: List[int] = field( |
|
|
default=None, metadata={'help': 'The number of output channels when "use_upsample" is set to "True"'}) |
|
|
|
|
|
zero_init_last: bool = field(default=False, metadata={'help': 'Zero init last weight'}) |
|
|
|
|
|
use_bypass: bool = field(default=True, metadata={'help': 'Whether to use bypass'}) |
|
|
|
|
|
def __post_init__(self): |
|
|
from .mapping import SwiftTuners |
|
|
self.swift_type = SwiftTuners.RESTUNING |
|
|
self.target_hidden_pos = 0 if self.target_hidden_pos is None else self.target_hidden_pos |
|
|
|
|
|
|
|
|
class ResTuning(SwiftAdapter): |
|
|
|
|
|
@staticmethod |
|
|
def prepare_model(model: nn.Module, config: ResTuningConfig, adapter_name: str) -> SwiftOutput: |
|
|
"""Prepare a model with `ResTuningConfig`""" |
|
|
|
|
|
def _forward_seq(self, input, *args, **kwargs): |
|
|
for idx, module in enumerate(self): |
|
|
if idx >= len(self.origin_module_keys): |
|
|
continue |
|
|
input = module(input) |
|
|
return input |
|
|
|
|
|
def _forward_target(self, *args, **kwargs): |
|
|
if self.target_modules_hook == 'input': |
|
|
if isinstance(self.target_hidden_pos, int): |
|
|
args = list(args) |
|
|
_arg = args[self.target_hidden_pos] |
|
|
else: |
|
|
_arg = kwargs[self.target_hidden_pos] |
|
|
args_main = _forward_restuning(self, _arg) |
|
|
if isinstance(self.target_hidden_pos, int): |
|
|
args[self.target_hidden_pos] = args_main |
|
|
else: |
|
|
kwargs[self.target_hidden_pos] = args_main |
|
|
args_main = getattr(self, f'forward_origin_{adapter_name}')(*args, **kwargs) |
|
|
else: |
|
|
_args_main = getattr(self, f'forward_origin_{adapter_name}')(*args, **kwargs) |
|
|
_arg = _args_main[self.target_hidden_pos] if isinstance(_args_main, (tuple, list, dict)) else _args_main |
|
|
args_main = _forward_restuning(self, _arg) |
|
|
if type(_args_main) != type(args_main): |
|
|
_args_main[self.target_hidden_pos] = args_main |
|
|
args_main = _args_main |
|
|
return args_main |
|
|
|
|
|
def _forward_restuning(self, origin_arg): |
|
|
probe_results = [] |
|
|
root_module_ins = self.root_module_ins_list[0] |
|
|
stem_module_ins_list = self.stem_module_ins_list |
|
|
top_module = model.get_submodule('') |
|
|
if root_module_ins: |
|
|
if root_module_ins.root_modules_hook == 'input': |
|
|
probe_results.append(root_module_ins.probe_input_data) |
|
|
else: |
|
|
probe_results.append(root_module_ins.probe_output_data) |
|
|
for i, st_mod in enumerate(stem_module_ins_list): |
|
|
if i == 0 and root_module_ins is None: |
|
|
probe_results.append(st_mod.probe_input_data) |
|
|
if st_mod.stem_modules_hook == 'input': |
|
|
probe_results.append(st_mod.probe_input_data) |
|
|
else: |
|
|
probe_results.append(st_mod.probe_output_data) |
|
|
args_main = getattr(top_module, f'restuning_{adapter_name}')(probe_results, origin_arg) |
|
|
return args_main |
|
|
|
|
|
|
|
|
module_keys = [key for key, _ in model.named_modules()] |
|
|
root_module_ins_list = [] |
|
|
if config.root_modules: |
|
|
for module_key in module_keys: |
|
|
if re.fullmatch(config.root_modules, module_key): |
|
|
root_module = model.get_submodule(module_key) |
|
|
logger.info(f'Matching root module [{module_key}] of type {type(root_module)}') |
|
|
if isinstance(root_module, (nn.ModuleList, nn.ModuleDict)): |
|
|
logger.warning( |
|
|
f'Type of {type(root_module)} may not be supported because of its customized forward') |
|
|
if config.root_modules_hook == 'input': |
|
|
root_module.register_forward_pre_hook(probe_input_pre_hook) |
|
|
else: |
|
|
root_module.register_forward_hook(probe_output_hook) |
|
|
root_module.root_modules_hook = config.root_modules_hook |
|
|
root_module_ins_list.append(root_module) |
|
|
break |
|
|
if len(root_module_ins_list) == 0: |
|
|
logger.error('Cannot match root modules') |
|
|
|
|
|
|
|
|
stem_module_ins_list = [] |
|
|
stem_module_ins_index = [] |
|
|
for module_key in module_keys: |
|
|
if (isinstance(config.stem_modules, str) and re.fullmatch(config.stem_modules, module_key)) or \ |
|
|
(isinstance(config.stem_modules, list) and module_key in config.stem_modules): |
|
|
stem_module = model.get_submodule(module_key) |
|
|
if isinstance(config.stem_modules, list): |
|
|
stem_module_ins_index.append(config.stem_modules.index(module_key)) |
|
|
logger.info(f'Matching stem module [{module_key}] of type {type(stem_module)}') |
|
|
if isinstance(stem_module, (nn.ModuleList, nn.ModuleDict)): |
|
|
logger.warning( |
|
|
f'Type of {type(stem_module)} may not be supported because of its customized forward') |
|
|
if len(root_module_ins_list) == 0 and len(stem_module_ins_list) == 0: |
|
|
stem_module.register_forward_pre_hook(probe_input_pre_hook) |
|
|
if config.stem_modules_hook == 'input': |
|
|
stem_module.register_forward_pre_hook(probe_input_pre_hook) |
|
|
else: |
|
|
stem_module.register_forward_hook(probe_output_hook) |
|
|
stem_module.stem_modules_hook = config.stem_modules_hook |
|
|
stem_module_ins_list.append(stem_module) |
|
|
if isinstance(config.stem_modules, list): |
|
|
stem_module_ins_list = [ |
|
|
stem_module_ins_list[stem_module_ins_index.index(i)] for i in range(len(stem_module_ins_index)) |
|
|
] |
|
|
depth = len(stem_module_ins_list) |
|
|
if len(stem_module_ins_list) == 0: |
|
|
raise Exception('Cannot match source modules') |
|
|
|
|
|
|
|
|
if len(stem_module_ins_list) != 0: |
|
|
top_module = model.get_submodule('') |
|
|
restuning_module = ResTuningBypassModule(config.dims, depth, adapter_name, config.use_upsample, |
|
|
config.upsample_out_channels, config.zero_init_last, |
|
|
config.tuner_cfg) |
|
|
setattr(top_module, f'restuning_{adapter_name}', restuning_module) |
|
|
|
|
|
|
|
|
target_module_ins = None |
|
|
for module_key in module_keys: |
|
|
if re.fullmatch(config.target_modules, module_key): |
|
|
tgt_module = model.get_submodule(module_key) |
|
|
logger.info(f'Matching target module [{module_key}] of type {type(tgt_module)}') |
|
|
if isinstance(tgt_module, (nn.ModuleList, nn.ModuleDict)): |
|
|
raise Exception( |
|
|
f'Type of {type(tgt_module)} may not be supported because of its customized forward') |
|
|
|
|
|
tgt_module.target_modules_hook = config.target_modules_hook |
|
|
tgt_module.target_hidden_pos = config.target_hidden_pos |
|
|
tgt_module.root_module_ins_list = root_module_ins_list |
|
|
tgt_module.stem_module_ins_list = stem_module_ins_list |
|
|
target_module_ins = tgt_module |
|
|
|
|
|
if isinstance(tgt_module, nn.Sequential) and not hasattr(tgt_module, 'origin_module_keys'): |
|
|
tgt_module.origin_module_keys = copy.deepcopy(list(tgt_module._modules.keys())) |
|
|
|
|
|
setattr(tgt_module, f'forward_origin_{adapter_name}', types.MethodType(_forward_seq, tgt_module)) |
|
|
else: |
|
|
setattr(tgt_module, f'forward_origin_{adapter_name}', tgt_module.forward) |
|
|
tgt_module.forward = types.MethodType(_forward_target, tgt_module) |
|
|
if target_module_ins is None: |
|
|
raise Exception('Cannot match target modules') |
|
|
|
|
|
def state_dict_callback(state_dict, adapter_name, **kwargs): |
|
|
return {key: value for key, value in state_dict.items() if f'restuning_{adapter_name}' in key} |
|
|
|
|
|
def mark_trainable_callback(model): |
|
|
return |
|
|
|
|
|
return SwiftOutput( |
|
|
config=config, state_dict_callback=state_dict_callback, mark_trainable_callback=mark_trainable_callback) |
|
|
|
|
|
@staticmethod |
|
|
def activate_adapter(module: torch.nn.Module, adapter_name: str, activate: bool, offload: str = None): |
|
|
modules = find_sub_module(module, f'restuning_{adapter_name}') |
|
|
for _module in modules: |
|
|
_module: ActivationMixin |
|
|
_module: nn.Module |
|
|
_module.set_activation(adapter_name, activate) |
|
|
SwiftAdapter.save_memory(_module, adapter_name, _module.module_key, activate, offload) |
|
|
|
|
|
|
|
|
class ResTuningBypassModule(nn.Module, ActivationMixin): |
|
|
"""The implementation of ResTuningBypass method. |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
dims, |
|
|
depth, |
|
|
adapter_name, |
|
|
use_upsample=False, |
|
|
upsample_out_channels=None, |
|
|
zero_init_last=False, |
|
|
tuner_cfg=None, |
|
|
): |
|
|
super(ResTuningBypassModule, self).__init__() |
|
|
super(nn.Module, self).__init__('') |
|
|
self.adapter_name = adapter_name |
|
|
|
|
|
self.bypass_blocks = nn.Sequential(*[ |
|
|
ResTunerBypassBlock( |
|
|
dim=dims[i] if isinstance(dims, list) else dims, |
|
|
layer_num=i, |
|
|
depth=depth, |
|
|
use_upsample=use_upsample, |
|
|
upsample_out_channels=upsample_out_channels[i] if isinstance(upsample_out_channels, list |
|
|
) else upsample_out_channels, |
|
|
zero_init_last=zero_init_last, |
|
|
tuner_cfg=tuner_cfg[i] if isinstance(tuner_cfg, list) else tuner_cfg) for i in range(depth) |
|
|
]) |
|
|
self.mark_all_sub_modules_as_plugin() |
|
|
|
|
|
def forward(self, x_list, origin_arg, **kwargs): |
|
|
if not self.is_activated(self.adapter_name): |
|
|
return origin_arg |
|
|
x_bypass = detach_tensors(x_list.pop(0)) |
|
|
x_bypass = x_bypass[0] if isinstance(x_bypass, (list, tuple)) else x_bypass |
|
|
x_list = detach_tensors(x_list) |
|
|
x_list = [_x[0] if isinstance(_x, (list, tuple)) else _x for _x in x_list] |
|
|
for i, (bp_blk, x_stem) in enumerate(zip(self.bypass_blocks, x_list)): |
|
|
target_size = x_list[i + 1].shape[2:] if i < len(x_list) - 1 else None |
|
|
x_bypass = bp_blk(x_stem, x_bypass, target_size, **kwargs) |
|
|
return x_bypass |
|
|
|
|
|
|
|
|
class ResTunerBypassBlock(nn.Module): |
|
|
|
|
|
def __init__(self, dim, layer_num=-1, depth=-1, use_upsample=False, zero_init_last=False, tuner_cfg=None, **kwargs): |
|
|
super().__init__() |
|
|
self.layer_num = layer_num |
|
|
self.depth = depth |
|
|
|
|
|
if isinstance(tuner_cfg, str): |
|
|
lateral_cfg = tuner_cfg |
|
|
vertical_cfg = tuner_cfg |
|
|
aux_cfg = 'upsample' if use_upsample and layer_num != depth - 1 else None |
|
|
elif isinstance(tuner_cfg, dict): |
|
|
lateral_cfg = tuner_cfg['lateral_cfg'] if 'lateral_cfg' in tuner_cfg else None |
|
|
vertical_cfg = tuner_cfg['vertical_cfg'] if 'vertical_cfg' in tuner_cfg else None |
|
|
aux_cfg = tuner_cfg['aux_cfg'] if 'aux_cfg' in tuner_cfg else None |
|
|
|
|
|
self.lateral_tuner = ResTuner(dim, layer_num, depth, zero_init_last, 'lateral', lateral_cfg, **kwargs) |
|
|
self.vertical_tuner = ResTuner(dim, layer_num, depth, zero_init_last, 'vertical', vertical_cfg, **kwargs) |
|
|
if aux_cfg and len(aux_cfg) != 0: |
|
|
self.aux_tuner = ResTuner(dim, layer_num, depth, zero_init_last, 'aux', aux_cfg, **kwargs) |
|
|
|
|
|
def forward(self, x_stem, x_bypass, target_size=None, **kwargs): |
|
|
x_lateral = self.lateral_tuner(x_stem) |
|
|
x_vertical = self.vertical_tuner(x_bypass) |
|
|
|
|
|
x_bypass_out = x_lateral + x_vertical |
|
|
if hasattr(self, 'aux_tuner'): |
|
|
x_bypass_out = self.aux_tuner(x_bypass_out, target_size) |
|
|
return x_bypass_out |
|
|
|
