| |
| """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()) |
|
|
