wehub-resource-sync
209 lines
8.6 KiB
Python
209 lines
8.6 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import json
|
|
import torch
|
|
from dataclasses import dataclass
|
|
from torch import nn
|
|
from types import MethodType
|
|
from typing import List, Literal, Optional
|
|
|
|
from swift.utils import get_logger, patch_getattr
|
|
from .utils import SwiftAdapter, SwiftConfig, SwiftOutput
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class ReftConfig(SwiftConfig):
|
|
"""
|
|
Train a model with Reft.
|
|
Paper: https://arxiv.org/pdf/2404.03592
|
|
|
|
Args:
|
|
model_type(`Optional[str]`): The model_type to find down_proj/layers.
|
|
layer_key(`Optional[str]`): Manually specify the layer key, for example `language_model.layers`.
|
|
layers (`Optional[List[int]]`): The layer number to inject.
|
|
r(`int`): The rank of Reft.
|
|
intervention_type (`Literal['NoreftIntervention', 'LoreftIntervention',
|
|
'ConsreftIntervention', 'LobireftIntervention',
|
|
'DireftIntervention', 'NodireftIntervention']`): The intervention type,
|
|
default LoreftIntervention
|
|
args (`Optional[str]`): Other reft_args in json-string format
|
|
"""
|
|
|
|
model_type: Optional[str] = None
|
|
layer_key: Optional[str] = None
|
|
layers: Optional[List[int]] = None
|
|
r: int = 4
|
|
intervention_type: Literal['NoreftIntervention', 'LoreftIntervention', 'ConsreftIntervention',
|
|
'LobireftIntervention', 'DireftIntervention',
|
|
'NodireftIntervention'] = 'LoreftIntervention'
|
|
args: Optional[str] = None
|
|
|
|
def __post_init__(self):
|
|
from .mapping import SwiftTuners
|
|
self.swift_type = SwiftTuners.REFT
|
|
if self.args:
|
|
self.args = json.loads(self.args)
|
|
else:
|
|
self.args = {}
|
|
|
|
|
|
class Reft(SwiftAdapter):
|
|
|
|
@staticmethod
|
|
def prepare_model(model: nn.Module, config: ReftConfig, adapter_name: str):
|
|
from swift.utils.import_utils import is_pyreft_available
|
|
if not is_pyreft_available():
|
|
raise ImportError('Please install pyreft before using ReFT: '
|
|
'`pip install pyreft`')
|
|
|
|
import pyreft
|
|
from pyreft import (ConsreftIntervention, DireftIntervention, LobireftIntervention, LoreftIntervention,
|
|
NodireftIntervention, NoreftIntervention, ReftModel)
|
|
from pyreft.interventions import LowRankRotateLayer
|
|
|
|
intervention_mapping = {
|
|
'NoreftIntervention': NoreftIntervention,
|
|
'LoreftIntervention': LoreftIntervention,
|
|
'ConsreftIntervention': ConsreftIntervention,
|
|
'LobireftIntervention': LobireftIntervention,
|
|
'DireftIntervention': DireftIntervention,
|
|
'NodireftIntervention': NodireftIntervention,
|
|
}
|
|
|
|
patch_getattr(ReftModel, 'model')
|
|
|
|
def forward(self, x):
|
|
self.to(x.device)
|
|
return self.forward_origin(x)
|
|
|
|
def forward2(self, base, source=None, subspaces=None):
|
|
self.to(base.device)
|
|
return self.forward_origin(base, source, subspaces)
|
|
|
|
if not hasattr(LowRankRotateLayer, 'forward_origin'):
|
|
LowRankRotateLayer.forward_origin = LowRankRotateLayer.forward
|
|
LowRankRotateLayer.forward = forward
|
|
NoreftIntervention.forward_origin = NoreftIntervention.forward
|
|
NoreftIntervention.forward = forward2
|
|
LoreftIntervention.forward_origin = LoreftIntervention.forward
|
|
LoreftIntervention.forward = forward2
|
|
ConsreftIntervention.forward_origin = ConsreftIntervention.forward
|
|
ConsreftIntervention.forward = forward2
|
|
LobireftIntervention.forward_origin = LobireftIntervention.forward
|
|
LobireftIntervention.forward = forward2
|
|
DireftIntervention.forward_origin = DireftIntervention.forward
|
|
DireftIntervention.forward = forward2
|
|
NodireftIntervention.forward_origin = NodireftIntervention.forward
|
|
NodireftIntervention.forward = forward2
|
|
|
|
module_list_key = config.layer_key
|
|
if module_list_key is None:
|
|
model_key_mapping = Reft.get_model_key_mapping(config.model_type, config)
|
|
module_list_key = model_key_mapping.module_list
|
|
logger.info(f'Applying Reft to module: {module_list_key}')
|
|
module_list: nn.ModuleList = model.get_submodule(module_list_key)
|
|
representations = []
|
|
for idx, layer in enumerate(module_list):
|
|
if config.layers and idx not in config.layers:
|
|
continue
|
|
intervention_config = {
|
|
'layer':
|
|
idx,
|
|
'component':
|
|
module_list_key + f'[{idx}].output',
|
|
'low_rank_dimension':
|
|
config.r,
|
|
'intervention':
|
|
intervention_mapping[config.intervention_type](
|
|
embed_dim=model.config.hidden_size, low_rank_dimension=config.r, **config.args)
|
|
}
|
|
representations.append(intervention_config)
|
|
|
|
reft_config = pyreft.ReftConfig(representations=representations)
|
|
reft_model = pyreft.get_reft_model(model, reft_config, set_device=False)
|
|
reft_model.reft_config = reft_model.config
|
|
reft_model.config = reft_model.model.config
|
|
|
|
def _pre_forward_hook(module, args, kwargs):
|
|
if 'base' in kwargs:
|
|
return args, kwargs
|
|
|
|
if 'input_ids' not in kwargs:
|
|
raise ValueError('Input does not contain `input_ids`, maybe the model does not support ReFT.')
|
|
# run intervened forward pass
|
|
unit_locations = None
|
|
if 'intervention_locations' in kwargs:
|
|
if kwargs['intervention_locations'].dim() == 3:
|
|
unit_locations = {
|
|
'sources->base': (None, kwargs['intervention_locations'].permute(1, 0, 2).tolist())
|
|
}
|
|
else:
|
|
# this is dummy for lora only baseline
|
|
unit_locations = {'sources->base': (None, 0)}
|
|
kwargs = {
|
|
'base': {
|
|
'input_ids': kwargs['input_ids'],
|
|
'attention_mask': kwargs['attention_mask']
|
|
},
|
|
'unit_locations': unit_locations,
|
|
'labels': kwargs['labels'],
|
|
'subspaces': kwargs['subspaces'].permute(1, 0, 2).tolist() if 'subspaces' in kwargs else None
|
|
}
|
|
return args, kwargs
|
|
|
|
def _post_forward_hook(module, args, kwargs, outputs):
|
|
return outputs[1]
|
|
|
|
def _generate(self, **kwargs):
|
|
# run intervened forward pass
|
|
unit_locations = None
|
|
if 'intervention_locations' in kwargs:
|
|
if kwargs['intervention_locations'].dim() == 3:
|
|
unit_locations = {
|
|
'sources->base': (None, kwargs['intervention_locations'].permute(1, 0, 2).tolist())
|
|
}
|
|
else:
|
|
# this is dummy for lora only baseline
|
|
unit_locations = {'sources->base': (None, 0)}
|
|
|
|
_kwargs = {
|
|
'base': {
|
|
'input_ids': kwargs.pop('input_ids'),
|
|
'attention_mask': kwargs.pop('attention_mask')
|
|
},
|
|
'unit_locations': unit_locations,
|
|
'subspaces': kwargs.pop('subspaces').permute(1, 0, 2).tolist() if 'subspaces' in kwargs else None
|
|
}
|
|
_kwargs = {**_kwargs, **kwargs}
|
|
return self.generate_origin(**_kwargs)[1]
|
|
|
|
reft_model.generate_origin = reft_model.generate
|
|
reft_model.generate = MethodType(_generate, reft_model)
|
|
reft_model.register_forward_pre_hook(_pre_forward_hook, with_kwargs=True)
|
|
reft_model.register_forward_hook(_post_forward_hook, with_kwargs=True)
|
|
|
|
def save_callback(swift_model, model_dir, adapter_name):
|
|
reft_model.save_intervention(save_directory=model_dir, include_model=False)
|
|
|
|
def mark_trainable_callback(model):
|
|
return
|
|
|
|
def load_callback(swift_model, model_dir, adapter_name):
|
|
reft_model.load_intervention(model_dir, include_model=False)
|
|
|
|
return SwiftOutput(
|
|
model=reft_model,
|
|
config=config,
|
|
mark_trainable_callback=mark_trainable_callback,
|
|
save_callback=save_callback,
|
|
load_callback=load_callback)
|
|
|
|
@staticmethod
|
|
def has_additional_modules():
|
|
return True
|
|
|
|
@staticmethod
|
|
def activate_adapter(module: torch.nn.Module, adapter_name: str, activate: bool, offload: str = None):
|
|
assert activate, 'ReFT does not support deactivate'
|