wehub-resource-sync
62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import os
|
|
|
|
from swift.arguments import ExportArguments
|
|
from swift.model import save_checkpoint
|
|
from swift.tuners import Swift
|
|
from swift.utils import HfConfigFactory, get_logger
|
|
from ..utils import prepare_model_template
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def check_tie_word_embeddings(model):
|
|
config = model.config
|
|
try:
|
|
from peft.utils import ModulesToSaveWrapper
|
|
if not HfConfigFactory.get_config_attr(config, 'tie_word_embeddings'):
|
|
return
|
|
for module in [model.get_input_embeddings(), model.get_output_embeddings()]:
|
|
if not isinstance(module, ModulesToSaveWrapper):
|
|
return
|
|
HfConfigFactory.set_config_attr(config, 'tie_word_embeddings', False)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def merge_lora(args: ExportArguments, device_map=None, replace_if_exists=False) -> None:
|
|
if replace_if_exists:
|
|
logger.info(f'replace_if_exists: {replace_if_exists}')
|
|
output_dir = getattr(args, 'output_dir', None) or f'{args.adapters[0]}-merged'
|
|
if os.path.exists(output_dir) and not replace_if_exists:
|
|
logger.info(f'The weight directory for the merged LoRA already exists in {output_dir}, '
|
|
'skipping the saving process.')
|
|
else:
|
|
# If the model is quantized, perform the merge on the original (unquantized) model.
|
|
# https://github.com/huggingface/peft/issues/2321
|
|
args.quant_method = None
|
|
origin_device_map = args.device_map
|
|
args.device_map = device_map or args.device_map
|
|
logger.info(f'merge_device_map: {device_map}')
|
|
model, template = prepare_model_template(args)
|
|
logger.info('Merge LoRA...')
|
|
check_tie_word_embeddings(model)
|
|
Swift.merge_and_unload(model)
|
|
model = model.model
|
|
logger.info('Saving merged weights...')
|
|
|
|
save_checkpoint(
|
|
model,
|
|
template.processor,
|
|
output_dir,
|
|
safe_serialization=args.safe_serialization,
|
|
model_dirs=args.adapters,
|
|
max_shard_size=args.max_shard_size,
|
|
additional_saved_files=model.model_meta.additional_saved_files)
|
|
logger.info(f'Successfully merged LoRA and saved in `{output_dir}`.')
|
|
args.device_map = origin_device_map
|
|
|
|
args.model = output_dir
|
|
args.model_dir = output_dir
|
|
args.adapters = []
|