forked from THUDM/CogVLM
-
Notifications
You must be signed in to change notification settings - Fork 0
/
merge_model.py
42 lines (38 loc) · 1.55 KB
/
merge_model.py
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
# -*- encoding: utf-8 -*-
import os, sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import argparse
from models.cogvlm_model import FineTuneTestCogVLMModel
from sat.training.model_io import save_checkpoint
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--version", type=str, default="base", help='version to interact with')
parser.add_argument("--from_pretrained", type=str, default="checkpoints/merged_lora", help='pretrained ckpt')
parser.add_argument("--fp16", action="store_true")
parser.add_argument("--bf16", action="store_true")
args = parser.parse_args()
rank = int(os.environ.get('RANK', 0))
world_size = int(os.environ.get('WORLD_SIZE', 1))
parser = FineTuneTestCogVLMModel.add_model_specific_args(parser)
args = parser.parse_args()
# load model
model, model_args = FineTuneTestCogVLMModel.from_pretrained(
args.from_pretrained,
args=argparse.Namespace(
deepspeed=None,
local_rank=rank,
rank=rank,
world_size=world_size,
model_parallel_size=world_size,
mode='inference',
skip_init=True,
use_gpu_initialization=True if torch.cuda.is_available() else False,
device='cuda',
**vars(args)
), url='local', overwrite_args={'model_parallel_size': 1})
model = model.eval()
model_args.save = './checkpoints/merged_model_{}'.format(model_args.eva_args["image_size"][0])
save_checkpoint(1, model, None, None, model_args)
if __name__ == "__main__":
main()