File size: 2,024 Bytes
7c50656 | 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 | """
Step 1: Merge DFlash-LoRA adapter into base model.
Usage:
conda activate sglang
python3 merge_lora.py
python3 merge_lora.py --ckpt epoch_2_step_15000 # 测其他 checkpoint
"""
import argparse
import os
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
BASE_MODEL = "/workspace/models/Qwen3-8B"
OUTPUT_ROOT = "/workspace/hanrui/syxin_old/Specforge/outputs/qwen3-8b-sft-32gpu-v2"
MERGE_ROOT = "/workspace/hanrui/syxin_old/Specforge/outputs/qwen3-8b-sft-32gpu-v2-merged"
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--ckpt", default="epoch_0_step_3000",
help="Checkpoint folder name under OUTPUT_ROOT")
p.add_argument("--merged-path", default=MERGE_ROOT,
help="Where to save the merged model")
return p.parse_args()
def main():
args = parse_args()
adapter_path = os.path.join(OUTPUT_ROOT, args.ckpt)
merged_path = args.merged_path
if os.path.exists(merged_path):
print(f"[skip] Merged model already exists: {merged_path}")
return
assert os.path.isdir(adapter_path), f"Adapter not found: {adapter_path}"
print(f"Base model : {BASE_MODEL}")
print(f"Adapter : {adapter_path}")
print(f"Output : {merged_path}")
print()
print("[1/4] Loading base model to CPU ...")
model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
torch_dtype=torch.bfloat16,
device_map="cpu",
)
print("[2/4] Loading LoRA adapter ...")
model = PeftModel.from_pretrained(model, adapter_path)
print("[3/4] Merging weights ...")
model = model.merge_and_unload()
print("[4/4] Saving merged model ...")
os.makedirs(merged_path, exist_ok=True)
model.save_pretrained(merged_path, safe_serialization=True)
AutoTokenizer.from_pretrained(BASE_MODEL).save_pretrained(merged_path)
print(f"\nDone. Merged model saved to: {merged_path}")
if __name__ == "__main__":
main()
|