Merge fine-tuned checkpoint with base model

if you have fine-tuned an open source model, such as llama or llama2 with lora or qlora, you can only public the fine-tuned checkpoints or merge the checkpoint into the base model, the public the new fine-tuned model. here is the code:

import os
import torch

from peft import PeftModel, PeftConfig
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)

check_point_dir = "/opt/models/sheen/v0.1/checkpoint-56400/"
output_dir = "/opt/models/sheen/v0.1/model"


config = PeftConfig.from_pretrained(check_point_dir)

tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)
model = PeftModel.from_pretrained(model, check_point_dir)

merged_model = model.merge_and_unload()
merged_model.save_pretrained(
    output_dir
    #push_to_hub=True,
    #repo_id=new_model_id,
    #private=True,
    #use_auth_token=os.environ.get("HUGGINGFACE_TOKEN"),
)
tokenizer.save_pretrained(
    output_dir)
    #push_to_hub=True,
    #repo_id=new_model_id,
    #private=True,
    #use_auth_token=os.environ.get("HUGGINGFACE_TOKEN"),

check_point_dir is the directory where your fine tuned checkpoints located, the contents looks as below,

root@ai:~/sheencloud# ls -lh /opt/models/sheen/v0.1/checkpoint-56400/
total 1.2G
-rw-r--r-- 1 root root  514 Oct 20 04:59 adapter_config.json
-rw-r--r-- 1 root root 1.2G Oct 20 03:05 adapter_model.bin
drwxr-xr-x 2 root root  106 Oct 20 03:06 global_step56400
-rw-r--r-- 1 root root   16 Oct 20 03:06 latest
-rw-r--r-- 1 root root  464 Oct 20 03:05 README.md
-rw-r--r-- 1 root root  15K Oct 20 03:06 rng_state.pth
-rw-r--r-- 1 root root  438 Oct 20 03:05 special_tokens_map.json
-rw-r--r-- 1 root root  828 Oct 20 03:05 tokenizer_config.json
-rw-r--r-- 1 root root 489K Oct 20 03:05 tokenizer.model
-rw-r--r-- 1 root root 6.6M Oct 20 03:06 trainer_state.json
-rw-r--r-- 1 root root 5.0K Oct 20 03:05 training_args.bin
-rwxr--r-- 1 root root  24K Oct 20 03:06 zero_to_fp32.py

you can also auto push saved model to huggingface by comment out the parameters on the code snippet above.

Leave a Reply

Your email address will not be published. Required fields are marked *