File size: 5,617 Bytes
d2d310a |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os.path as osp
import shutil
from mmengine.config import Config, DictAction
from mmengine.fileio import PetrelBackend, get_file_backend
from xtuner.configs import cfgs_name_path
from xtuner.model.utils import guess_load_checkpoint
from xtuner.registry import BUILDER
def parse_args():
parser = argparse.ArgumentParser(
description='Convert the pth model to HuggingFace model')
parser.add_argument('config', help='config file name or path.')
parser.add_argument('pth_model', help='pth model file')
parser.add_argument(
'save_dir', help='the directory to save HuggingFace model')
parser.add_argument(
'--fp32',
action='store_true',
help='Save LLM in fp32. If not set, fp16 will be used by default.')
parser.add_argument(
'--max-shard-size',
type=str,
default='2GB',
help='Only applicable for LLM. The maximum size for '
'each sharded checkpoint.')
parser.add_argument(
'--cfg-options',
nargs='+',
action=DictAction,
help='override some settings in the used config, the key-value pair '
'in xxx=yyy format will be merged into config file. If the value to '
'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
'Note that the quotation marks are necessary and that no white space '
'is allowed.')
args = parser.parse_args()
return args
def main():
args = parse_args()
# parse config
if not osp.isfile(args.config):
try:
args.config = cfgs_name_path[args.config]
except KeyError:
raise FileNotFoundError(f'Cannot find {args.config}')
# load config
cfg = Config.fromfile(args.config)
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
model_name = cfg.model.type if isinstance(cfg.model.type,
str) else cfg.model.type.__name__
if 'LLaVAModel' in model_name:
cfg.model.pretrained_pth = None
model = BUILDER.build(cfg.model)
backend = get_file_backend(args.pth_model)
if isinstance(backend, PetrelBackend):
from xtuner.utils.fileio import patch_fileio
with patch_fileio():
state_dict = guess_load_checkpoint(args.pth_model)
else:
state_dict = guess_load_checkpoint(args.pth_model)
model.load_state_dict(state_dict, strict=False)
print(f'Load PTH model from {args.pth_model}')
if 'LLaVAModel' in model_name:
if cfg.model.get('llm') and (not cfg.model.get('freeze_llm', False)
or cfg.model.get('llm_lora')):
if 'PeftModel' in model.llm.__class__.__name__:
llm_path = osp.join(args.save_dir, 'llm_adapter')
print(f'Saving LLM adapter to {llm_path}')
else:
llm_path = args.save_dir
print(f'Saving LLM tokenizer to {llm_path}')
tokenizer = BUILDER.build(cfg.tokenizer)
tokenizer.save_pretrained(llm_path)
print(f'Saving LLM to {llm_path}')
if not args.fp32:
# StarCycle: The llm has been quantinized
# print('Convert LLM to float16')
# model.llm.half()
model.llm.save_pretrained(
llm_path, max_shard_size=args.max_shard_size)
if cfg.model.get('visual_encoder') and (
not cfg.model.get('freeze_visual_encoder', False)
or cfg.model.get('visual_encoder_lora')):
if 'PeftModel' in model.visual_encoder.__class__.__name__:
visual_encoder_path = osp.join(args.save_dir,
'visual_encoder_adapter')
print(
f'Saving visual_encoder adapter to {visual_encoder_path}')
else:
visual_encoder_path = osp.join(args.save_dir, 'visual_encoder')
print('Saving visual_encoder image_processor to'
f'{visual_encoder_path}')
image_processor = BUILDER.build(cfg.image_processor)
image_processor.save_pretrained(visual_encoder_path)
print(f'Saving visual_encoder to {visual_encoder_path}')
model.visual_encoder.save_pretrained(
visual_encoder_path, max_shard_size=args.max_shard_size)
if hasattr(model, 'projector'):
projector_path = osp.join(args.save_dir, 'projector')
print(f'Saving projector to {projector_path}')
model.projector.save_pretrained(
projector_path, max_shard_size=args.max_shard_size)
else:
llm_path = args.save_dir
if 'PeftModel' in model.llm.__class__.__name__:
print(f'Saving adapter to {llm_path}')
else:
print(f'Saving LLM tokenizer to {llm_path}')
tokenizer = BUILDER.build(cfg.tokenizer)
tokenizer.save_pretrained(llm_path)
print(f'Saving LLM to {llm_path}')
if not args.fp32:
print('Convert LLM to float16')
model.llm.half()
model.llm.save_pretrained(
llm_path,
max_shard_size=args.max_shard_size,
safe_serialization=False)
shutil.copyfile(args.config, osp.join(args.save_dir, 'xtuner_config.py'))
print('All done!')
if __name__ == '__main__':
main()
|