File size: 3,358 Bytes
994182c | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | #!/usr/bin/env python3
"""Merge a selected LoRA adapter into the base model."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any
import torch
import yaml
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--config", required=True)
parser.add_argument("--adapter", help="Adapter checkpoint path. Defaults to run.output_dir.")
parser.add_argument("--output", help="Merged output path. Defaults to run.merged_output_dir.")
parser.add_argument("--dtype", default=None, choices=["bfloat16", "float16", "float32"])
return parser.parse_args()
def read_yaml(path: str | Path) -> dict[str, Any]:
with Path(path).open("r", encoding="utf-8") as fh:
payload = yaml.safe_load(fh) or {}
if not isinstance(payload, dict):
raise TypeError(f"Expected a YAML mapping in {path}")
return payload
def torch_dtype(name: str):
return {
"bfloat16": torch.bfloat16,
"float16": torch.float16,
"float32": torch.float32,
}[name]
def load_base_model(transformers_module, model_name: str, dtype, trust_remote_code: bool):
errors: list[str] = []
for class_name in [
"AutoModelForMultimodalLM",
"AutoModelForImageTextToText",
"AutoModelForVision2Seq",
"AutoModelForCausalLM",
]:
model_cls = getattr(transformers_module, class_name, None)
if model_cls is None:
continue
try:
return model_cls.from_pretrained(
model_name,
torch_dtype=dtype,
trust_remote_code=trust_remote_code,
device_map="auto",
)
except Exception as exc:
errors.append(f"{class_name}: {exc!r}")
raise RuntimeError("Could not load base model:\n" + "\n".join(errors))
def main() -> int:
args = parse_args()
config = read_yaml(args.config)
run_cfg = config["run"]
model_cfg = config["model"]
adapter_path = Path(args.adapter or run_cfg["output_dir"])
output_path = Path(args.output or run_cfg["merged_output_dir"])
dtype_name = args.dtype or str(config.get("merge", {}).get("merged_dtype", model_cfg.get("dtype", "bfloat16")))
import transformers
from peft import PeftModel
from transformers import AutoTokenizer
dtype = torch_dtype(dtype_name)
base = load_base_model(
transformers,
model_cfg["name_or_path"],
dtype=dtype,
trust_remote_code=bool(model_cfg.get("trust_remote_code", True)),
)
model = PeftModel.from_pretrained(base, adapter_path)
merged = model.merge_and_unload()
output_path.mkdir(parents=True, exist_ok=True)
merged.save_pretrained(output_path, safe_serialization=True)
tokenizer = AutoTokenizer.from_pretrained(model_cfg["name_or_path"], trust_remote_code=True)
tokenizer.save_pretrained(output_path)
manifest = {
"base_model": model_cfg["name_or_path"],
"adapter": str(adapter_path),
"output": str(output_path),
"dtype": dtype_name,
}
(output_path / "merge_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
print(json.dumps(manifest, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
|