wehub-resource-sync
189 lines
7.5 KiB
Python
189 lines
7.5 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import inspect
|
|
import re
|
|
import torch
|
|
import types
|
|
from dataclasses import dataclass, field
|
|
from torch import nn
|
|
from transformers.activations import ACT2CLS
|
|
from typing import List, Union
|
|
|
|
from swift.utils import find_sub_module, get_logger
|
|
from .utils import ActivationMixin, SwiftAdapter, SwiftConfig, SwiftOutput
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class AdapterConfig(SwiftConfig):
|
|
"""
|
|
The configuration class for the adapter module.
|
|
|
|
Adapters project input tokens by an MLP layer.
|
|
'Parameter-Efficient Transfer Learning for NLP' by Houlsby et al.(2019)
|
|
See http://arxiv.org/abs/1902.00751
|
|
|
|
Args:
|
|
dim(`int`): The dimension of the hidden states
|
|
target_modules(`Union[str, List[str]]`): The feedforward module to be replaced.
|
|
in regex format if this argument is str, else will match with `end with` if List[str].
|
|
hidden_pos(`Union[str, int]`): The position of the hidden state to be passed into the adapter,
|
|
can be int (args) or str (kwargs)
|
|
method_name(`str`): The method to be replaced, default is `forward`
|
|
adapter_length: The length of the adapter length (intermediate length)
|
|
act_layer: The activation layer of the adapter
|
|
"""
|
|
|
|
dim: int = field(default=None, metadata={'help': 'The dimension of the hidden states'})
|
|
|
|
target_modules: Union[str, List[str]] = field(
|
|
default=None,
|
|
metadata={
|
|
'help':
|
|
'The feedforward module to be replaced. in regex format if this argument is str, '
|
|
'else will match with `end with` if List[str].'
|
|
})
|
|
|
|
hidden_pos: Union[str, int] = field(
|
|
default=None,
|
|
metadata={
|
|
'help': 'The position of the hidden state to be passed into the adapter, can be int (args) or str (kwargs)'
|
|
})
|
|
|
|
method_name: str = field(default='forward', metadata={'help': 'The method to be replaced, default is `forward`'})
|
|
|
|
adapter_length: int = field(
|
|
default=128, metadata={'help': 'The length of the adapter length (intermediate length)'})
|
|
|
|
act_layer: str = field(default='gelu', metadata={'help': 'The activation layer of the adapter'})
|
|
|
|
def __post_init__(self):
|
|
from .mapping import SwiftTuners
|
|
self.swift_type = SwiftTuners.ADAPTER
|
|
|
|
|
|
class Adapter(SwiftAdapter):
|
|
|
|
@staticmethod
|
|
def prepare_model(model: nn.Module, config: AdapterConfig, adapter_name: str) -> SwiftOutput:
|
|
"""Prepare a model with `AdapterConfig`"""
|
|
module_keys = [key for key, _ in model.named_modules()]
|
|
|
|
for module_key in module_keys:
|
|
if isinstance(config.target_modules, str):
|
|
target_module_found = re.fullmatch(config.target_modules, module_key)
|
|
else:
|
|
target_module_found = any(module_key.endswith(target_key) for target_key in config.target_modules)
|
|
|
|
if target_module_found: # noqa
|
|
module = model.get_submodule(module_key)
|
|
|
|
def _forward(self, *args, **kwargs):
|
|
args = getattr(self, f'forward_origin_{adapter_name}')(*args, **kwargs)
|
|
if isinstance(args, (tuple, list, dict)):
|
|
if isinstance(config.hidden_pos, int):
|
|
_type = type(args)
|
|
args = list(args)
|
|
args[config.hidden_pos] = getattr(self, f'adapter_{adapter_name}')(args[config.hidden_pos])
|
|
args = _type(args)
|
|
else:
|
|
args[config.hidden_pos] = getattr(self, f'adapter_{adapter_name}')(args[config.hidden_pos])
|
|
elif isinstance(args, torch.Tensor):
|
|
args = getattr(self, f'adapter_{adapter_name}')(args)
|
|
return args
|
|
|
|
def _feed_forward_chunk(self, attention_output):
|
|
return _forward(self, attention_output)
|
|
|
|
# TODO The `config.method_name` method should not be replaced twice.
|
|
|
|
setattr(module, f'forward_origin_{adapter_name}', getattr(module, config.method_name))
|
|
num_args_in_forward_chunk_fn = len(
|
|
inspect.signature(getattr(module, f'forward_origin_{adapter_name}')).parameters)
|
|
if config.method_name == 'feed_forward_chunk' and num_args_in_forward_chunk_fn == 1:
|
|
setattr(module, config.method_name, types.MethodType(_feed_forward_chunk, module))
|
|
else:
|
|
setattr(module, config.method_name, types.MethodType(_forward, module))
|
|
adapter_module = AdapterModule(config.dim, adapter_name, module_key, config.adapter_length,
|
|
ACT2CLS[config.act_layer])
|
|
setattr(module, f'adapter_{adapter_name}', adapter_module)
|
|
logger.info(f'Adapter modules(module_key): {module_key}.adapter_{adapter_name}')
|
|
|
|
def state_dict_callback(state_dict, adapter_name: str, **kwargs):
|
|
return {key: value for key, value in state_dict.items() if f'adapter_{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'adapter_{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 AdapterModule(nn.Module, ActivationMixin):
|
|
"""The implementation of adapter tuning method.
|
|
|
|
Adapters project input tokens by an MLP layer.
|
|
'Parameter-Efficient Transfer Learning for NLP' by Houlsby et al.(2019)
|
|
See http://arxiv.org/abs/1902.00751
|
|
|
|
Args:
|
|
dim: An integer indicating the embedding dimension.
|
|
adapter_length: An integer indicating the length of adapter tuning.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
adapter_name,
|
|
module_key,
|
|
adapter_length=None,
|
|
act_layer=nn.GELU,
|
|
):
|
|
super(AdapterModule, self).__init__()
|
|
super(nn.Module, self).__init__(module_key)
|
|
self.dim = dim
|
|
self.adapter_name = adapter_name
|
|
self.adapter_length = adapter_length
|
|
self.linear1 = nn.Linear(dim, adapter_length)
|
|
self.act = act_layer()
|
|
self.linear2 = nn.Linear(adapter_length, dim)
|
|
self.init_weights()
|
|
self._prepared = False
|
|
self.mark_all_sub_modules_as_plugin()
|
|
|
|
def init_weights(self):
|
|
|
|
def _init_weights(m):
|
|
if isinstance(m, nn.Linear):
|
|
nn.init.xavier_uniform_(m.weight)
|
|
nn.init.normal_(m.bias, std=1e-6)
|
|
|
|
self.apply(_init_weights)
|
|
|
|
def forward(self, x, identity=None):
|
|
if not self.is_activated(self.adapter_name):
|
|
return x
|
|
if not self._prepared:
|
|
self.linear1.to(x.device)
|
|
self.act.to(x.device)
|
|
self.linear2.to(x.device)
|
|
self._prepared = True
|
|
|
|
x_dtype = x.dtype
|
|
x = x.to(self.linear1.weight.dtype)
|
|
out = self.linear2(self.act(self.linear1(x)))
|
|
if identity is None:
|
|
identity = x
|
|
identity = identity.to(out.dtype)
|
|
out = identity + out
|
|
return out.to(x_dtype)
|