merge_aes_lora.py

#!/usr/bin/env python3
"""
Merge the AES LoRA adapter into the all-merged base model to create
a final standalone AES-merged model.

Base: /workspace/elinnos/merged_models/elinnos_all_merged_final
      (Qwen2.5-7B + V1+V2+V3+V4+SRAM+I2CS all baked in)
Adapter: /workspace/elinnos/elinnos-qwen2.5-7b-aes-lora

Output: /workspace/elinnos/merged_models/elinnos_aes_merged_final

Usage:
  python3 merge_aes_lora.py
  python3 merge_aes_lora.py --output /path/to/output
"""
import argparse
import gc
import json
import shutil
import time
from pathlib import Path

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

WORKSPACE = Path("/workspace/elinnos")
DEFAULT_BASE = WORKSPACE / "merged_models" / "elinnos_all_merged_final"
DEFAULT_ADAPTER = WORKSPACE / "elinnos-qwen2.5-7b-aes-lora"
DEFAULT_OUTPUT = WORKSPACE / "merged_models" / "elinnos_aes_merged_final"


def parse_args():
    p = argparse.ArgumentParser(description="Merge AES LoRA adapter into all-merged base model")
    p.add_argument("--base", type=str, default=str(DEFAULT_BASE))
    p.add_argument("--adapter", type=str, default=str(DEFAULT_ADAPTER))
    p.add_argument("--output", type=str, default=str(DEFAULT_OUTPUT))
    p.add_argument("--dtype", type=str, default="bfloat16", choices=["float16", "bfloat16", "float32"])
    return p.parse_args()


def get_dtype(s):
    return {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}[s]


def main():
    args = parse_args()
    base = Path(args.base).resolve()
    adapter = Path(args.adapter).resolve()
    output = Path(args.output).resolve()
    dtype = get_dtype(args.dtype)

    print("=" * 70)
    print("  MERGE AES LoRA → FINAL STANDALONE MODEL")
    print("=" * 70)
    print(f"  Base:    {base}")
    print(f"  Adapter: {adapter}")
    print(f"  Output:  {output}")
    print(f"  Dtype:   {args.dtype}")

    if torch.cuda.is_available():
        print(f"  GPU:     {torch.cuda.get_device_name(0)}")
        print(f"  VRAM:    {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")

    # Pre-flight checks
    assert (base / "config.json").is_file(), f"Base model not found: {base}"
    assert (adapter / "adapter_config.json").is_file(), f"Adapter not found: {adapter}"
    print("\n  Pre-flight checks: OK")

    output.mkdir(parents=True, exist_ok=True)
    t0 = time.time()

    # Load base model
    print(f"\n  Loading base model ({args.dtype})...")
    model = AutoModelForCausalLM.from_pretrained(
        str(base), torch_dtype=dtype, device_map="auto",
        trust_remote_code=True, low_cpu_mem_usage=True,
    )
    print(f"  Base loaded. Device map: {getattr(model, 'hf_device_map', 'N/A')}")

    # Load adapter config
    with open(adapter / "adapter_config.json") as f:
        cfg = json.load(f)
    print(f"  AES LoRA config: r={cfg['r']}, alpha={cfg['lora_alpha']}, targets={cfg['target_modules']}")

    # Apply AES LoRA adapter
    print(f"\n  Applying AES LoRA adapter...")
    model = PeftModel.from_pretrained(model, str(adapter))

    # Merge and unload
    print(f"  Merging AES LoRA weights into model...")
    model = model.merge_and_unload()
    print(f"  Merge complete.")

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Save merged model
    print(f"\n  Saving merged model to {output}...")
    model.save_pretrained(str(output), safe_serialization=True, max_shard_size="5GB")
    print(f"  Model saved.")

    # Copy tokenizer from base (already has chat template etc.)
    tokenizer = AutoTokenizer.from_pretrained(str(base), trust_remote_code=True)
    tokenizer.save_pretrained(str(output))
    print(f"  Tokenizer saved.")

    # Copy chat template
    chat_template = base / "chat_template.jinja"
    if not chat_template.exists():
        chat_template = adapter / "chat_template.jinja"
    if chat_template.exists():
        shutil.copy2(str(chat_template), str(output / "chat_template.jinja"))
        print(f"  Chat template copied.")

    # Copy generation config
    gen_config = base / "generation_config.json"
    if gen_config.exists():
        shutil.copy2(str(gen_config), str(output / "generation_config.json"))
        print(f"  Generation config copied.")

    elapsed = time.time() - t0
    print("\n" + "=" * 70)
    print("  AES LORA MERGE COMPLETE")
    print("=" * 70)
    saved = sorted(output.glob("*"))
    total_size = sum(f.stat().st_size for f in saved if f.is_file())
    for f in saved:
        if f.is_file():
            print(f"  {f.name:45s}  {f.stat().st_size / (1024*1024):>10.2f} MB")
    print(f"  {'TOTAL':45s}  {total_size / (1024*1024):>10.2f} MB")
    print(f"\n  Merge time: {elapsed:.1f}s ({elapsed/60:.1f} min)")
    print(f"  Ready at: {output}")


if __name__ == "__main__":
    main()