wehub-resource-sync
73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import torch
|
|
from dataclasses import dataclass, field
|
|
from torch import nn
|
|
|
|
from swift.utils import get_logger
|
|
from .utils import SwiftAdapter, SwiftConfig, SwiftOutput
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class NEFTuneConfig(SwiftConfig):
|
|
"""
|
|
The configuration class for the NEFTune module.
|
|
|
|
NEFTune adds slightly noises to embedding outputs.
|
|
See https://arxiv.org/abs/2310.05914
|
|
|
|
Args:
|
|
noise_alpha(`float`): The noise alpha value used for the NEFTune, default 5.0
|
|
"""
|
|
noise_alpha: float = field(default=5.0, metadata={'help': 'The noise alpha value used for the NEFTune'})
|
|
|
|
def __post_init__(self):
|
|
from .mapping import SwiftTuners
|
|
self.swift_type = SwiftTuners.NEFTUNE
|
|
|
|
|
|
class NEFTune(SwiftAdapter):
|
|
|
|
@staticmethod
|
|
def prepare_model(model: nn.Module, config: NEFTuneConfig, adapter_name: str) -> SwiftOutput:
|
|
"""Prepare a model with `NEFTuneConfig`"""
|
|
for sub_module in model.modules():
|
|
if isinstance(sub_module, torch.nn.Embedding):
|
|
|
|
def neftune_hook(module, args, output):
|
|
if module.training and getattr(module, 'nef_activated'):
|
|
dims = torch.tensor(output.size(-1) * output.size(-2))
|
|
mag_norm = config.noise_alpha / torch.sqrt(dims)
|
|
output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm)
|
|
return output
|
|
|
|
if hasattr(sub_module, 'nef_activated'):
|
|
raise ValueError('NEFTune does not support a second tuner.')
|
|
|
|
sub_module.register_forward_hook(neftune_hook)
|
|
sub_module.nef_activated = True
|
|
|
|
def state_dict_callback(state_dict, adapter_name, **kwargs):
|
|
return state_dict
|
|
|
|
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):
|
|
for sub_module in module.modules():
|
|
if isinstance(sub_module, torch.nn.Embedding):
|
|
sub_module.nef_activated = activate
|
|
|
|
@staticmethod
|
|
def freeze_model():
|
|
return False
|
|
|
|
@staticmethod
|
|
def has_additional_modules():
|
|
return False
|