File size: 3,964 Bytes
274951a | 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 106 107 | """Materialize one PEFT checkpoint for single-model evaluation.
The merged directory is an evaluation cache, not a training checkpoint. It
contains no optimizer state and may be deleted after its generations have been
durably written.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
from pathlib import Path
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--base-model", required=True)
parser.add_argument("--adapter", required=True)
parser.add_argument("--output", required=True)
args = parser.parse_args()
import torch
from peft import PeftModel
from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoModelForImageTextToText,
AutoTokenizer,
)
base = Path(args.base_model).resolve()
adapter = Path(args.adapter).resolve()
output = Path(args.output).resolve()
if not base.is_dir() or not adapter.is_dir():
raise SystemExit("--base-model and --adapter must be existing local directories")
if output.exists() and any(output.iterdir()):
raise SystemExit(f"refusing to overwrite non-empty merged directory: {output}")
output.mkdir(parents=True, exist_ok=True)
adapter_weights = adapter / "adapter_model.safetensors"
if not adapter_weights.is_file():
raise SystemExit(f"adapter weights not found: {adapter_weights}")
tokenizer = AutoTokenizer.from_pretrained(base, trust_remote_code=True, local_files_only=True)
model = AutoModelForCausalLM.from_pretrained(
base,
trust_remote_code=True,
local_files_only=True,
dtype=torch.bfloat16,
device_map="cpu",
)
merged = PeftModel.from_pretrained(model, adapter, local_files_only=True).merge_and_unload()
base_config = AutoConfig.from_pretrained(base, trust_remote_code=True, local_files_only=True)
if base_config.model_type == "qwen3_5" and hasattr(base_config, "vision_config"):
# AutoModelForCausalLM intentionally extracts Qwen3.5's text model. A
# bare Qwen3_5TextConfig cannot currently be loaded by vLLM 0.22.x,
# which still expects the outer config's vision_config even in
# language-model-only mode. Put the merged text tower and LM head back
# into the original full-model container so direct evaluation remains
# loadable while still using the merged adapter weights.
full_model = AutoModelForImageTextToText.from_pretrained(
base,
trust_remote_code=True,
local_files_only=True,
dtype=torch.bfloat16,
device_map="cpu",
)
full_model.model.language_model = merged.model
full_model.lm_head = merged.lm_head
model_to_save = full_model
container = "full_qwen3_5"
else:
model_to_save = merged
container = "causal_lm"
model_to_save.save_pretrained(output, safe_serialization=True, max_shard_size="5GB")
tokenizer.save_pretrained(output)
manifest = {
"schema": "merged_peft_eval_cache_v1",
"base_model": str(base),
"adapter": str(adapter),
"adapter_sha256": _sha256(adapter_weights),
"dtype": "bfloat16",
"container": container,
"optimizer_state_included": False,
}
temporary = output / "merge_manifest.json.tmp"
temporary.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
with temporary.open("r+") as handle:
handle.flush()
os.fsync(handle.fileno())
os.replace(temporary, output / "merge_manifest.json")
print(f"[merge] {adapter} + {base} -> {output}")
if __name__ == "__main__":
main()
|