File size: 1,811 Bytes
cb2428f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from swift.llm import ExportArguments, prepare_model_template, save_checkpoint
from swift.tuners import Swift
from swift.utils import get_logger
logger = get_logger()
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:
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)
quant_method = model.model_info.quant_method
assert quant_method is None, (f'quant_method: {quant_method}, '
'quantized model and does not support merge-lora.')
logger.info('Merge LoRA...')
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 = []
|