DPO-Training/merge_dpo_adapters.py

#!/usr/bin/env python3
"""
Merge LoRA adapters from downloaded DPO model with base Qwen3-0.6B.
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import os

BASE_MODEL = "Qwen/Qwen3-0.6B"
LORA_PATH = "/home/ma/models/Qwen3-0.6B-DPO"
OUTPUT_PATH = "/home/ma/models/Qwen3-0.6B-DPO-merged"

print("=" * 60)
print("Merging DPO LoRA Adapters")
print("=" * 60)

# Load tokenizer
print("\n📥 Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)

# Load base model
print("📥 Loading base model (this may take a few minutes on CPU)...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float32,
    device_map="cpu",
    trust_remote_code=True,
)

# Load LoRA adapters
print("📥 Loading DPO LoRA adapters...")
model = PeftModel.from_pretrained(base_model, LORA_PATH)

# Merge adapters
print("🔧 Merging adapters...")
model = model.merge_and_unload()

# Save merged model
print(f"💾 Saving merged model to {OUTPUT_PATH}...")
os.makedirs(OUTPUT_PATH, exist_ok=True)
model.save_pretrained(OUTPUT_PATH)
tokenizer.save_pretrained(OUTPUT_PATH)

print("\n" + "=" * 60)
print("✅ Merge Complete!")
print("=" * 60)
print(f"Merged model saved to: {OUTPUT_PATH}")
print(
    f"Size: {sum(os.path.getsize(os.path.join(OUTPUT_PATH, f)) for f in os.listdir(OUTPUT_PATH) if os.path.isfile(os.path.join(OUTPUT_PATH, f))) / (1024 * 1024):.1f} MB"
)