#!/usr/bin/env python3 """Merge a PEFT LoRA into sharded safetensors with bounded host RAM. The model is never instantiated. Each base shard is copied byte-for-byte to a temporary output file, then each LoRA target tensor is patched in-place, one tensor at a time and in small row tiles. Completed shards are atomically renamed, so ``--resume`` can safely continue after interruption. """ from __future__ import annotations import argparse import gc import hashlib import json import math import mmap import os import platform import re import shutil import struct import sys import time from collections import Counter from datetime import datetime, timezone from pathlib import Path import safetensors import torch from safetensors import safe_open DTYPES = { "BF16": (torch.bfloat16, 2), "F16": (torch.float16, 2), "F32": (torch.float32, 4), } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--base-dir", type=Path, required=True) parser.add_argument("--base-repo", required=True) parser.add_argument("--base-revision", required=True) parser.add_argument("--adapter-dir", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--tile-rows", type=int, default=32) parser.add_argument("--threads", type=int, default=2) parser.add_argument("--expected-targets", type=int, default=190) parser.add_argument("--resume", action="store_true") return parser.parse_args() def sha256_file(path: Path, block_size: int = 8 << 20) -> str: digest = hashlib.sha256() with path.open("rb", buffering=0) as handle: while True: block = handle.read(block_size) if not block: break digest.update(block) try: os.posix_fadvise(handle.fileno(), 0, 0, os.POSIX_FADV_DONTNEED) except (AttributeError, OSError): pass return digest.hexdigest() def atomic_json(path: Path, value: object) -> None: temporary = path.with_name(path.name + ".tmp") temporary.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n") os.replace(temporary, path) def read_safetensors_header(path: Path) -> tuple[int, dict[str, object]]: with path.open("rb") as handle: raw = handle.read(8) if len(raw) != 8: raise ValueError(f"invalid safetensors header in {path}") header_len = struct.unpack(" dict[str, dict[str, object]]: return {key: value for key, value in header.items() if key != "__metadata__"} def checkpoint_key_for(module_path: str) -> str: match = re.search(r"layers\.\d+\..+$", module_path) if match is None: raise ValueError(f"cannot map adapter module {module_path!r}") return f"model.language_model.{match.group(0)}.weight" def adapter_pairs(adapter_path: Path) -> tuple[dict[str, dict[str, object]], dict[str, object]]: config = json.loads((adapter_path.parent / "adapter_config.json").read_text()) rank = int(config["r"]) alpha = float(config["lora_alpha"]) scaling = alpha / math.sqrt(rank) if config.get("use_rslora") else alpha / rank pairs: dict[str, dict[str, object]] = {} with safe_open(adapter_path, framework="pt", device="cpu") as handle: keys = set(handle.keys()) for a_key in sorted(key for key in keys if key.endswith(".lora_A.weight")): module = a_key[: -len(".lora_A.weight")] b_key = module + ".lora_B.weight" if b_key not in keys: raise ValueError(f"missing pair tensor {b_key}") target = checkpoint_key_for(module) if target in pairs: raise ValueError(f"duplicate target {target}") a_slice = handle.get_slice(a_key) b_slice = handle.get_slice(b_key) pairs[target] = { "module": module, "a_key": a_key, "b_key": b_key, "a_shape": list(a_slice.get_shape()), "b_shape": list(b_slice.get_shape()), "a_dtype": str(a_slice.get_dtype()), "b_dtype": str(b_slice.get_dtype()), "scaling": scaling, } return pairs, config def validate_preflight( base_dir: Path, adapter_path: Path, expected_targets: int, ) -> tuple[ dict[str, dict[str, object]], dict[str, object], dict[str, str], dict[str, tuple[int, dict[str, object]]], ]: index_path = base_dir / "model.safetensors.index.json" index = json.loads(index_path.read_text()) weight_map: dict[str, str] = index["weight_map"] pairs, adapter_config = adapter_pairs(adapter_path) if len(pairs) != expected_targets: raise ValueError(f"expected {expected_targets} LoRA targets, found {len(pairs)}") headers: dict[str, tuple[int, dict[str, object]]] = {} shard_names = sorted(set(weight_map.values())) for shard_name in shard_names: headers[shard_name] = read_safetensors_header(base_dir / shard_name) for target, pair in pairs.items(): if target not in weight_map: raise ValueError(f"adapter target missing from base index: {target}") shard_name = weight_map[target] entry = tensor_entries(headers[shard_name][1])[target] base_shape = list(entry["shape"]) a_shape = pair["a_shape"] b_shape = pair["b_shape"] if entry["dtype"] != "BF16": raise ValueError(f"target {target} is {entry['dtype']}, expected BF16") if len(base_shape) != 2 or len(a_shape) != 2 or len(b_shape) != 2: raise ValueError(f"target {target} is not a 2-D matrix") if b_shape[0] != base_shape[0] or a_shape[1] != base_shape[1] or b_shape[1] != a_shape[0]: raise ValueError( f"shape mismatch for {target}: base={base_shape}, A={a_shape}, B={b_shape}" ) return pairs, adapter_config, weight_map, headers def patch_tensor( mm: mmap.mmap, data_start: int, entry: dict[str, object], a: torch.Tensor, b: torch.Tensor, scaling: float, tile_rows: int, ) -> dict[str, object]: dtype_name = str(entry["dtype"]) if dtype_name not in DTYPES: raise ValueError(f"unsupported target dtype {dtype_name}") dtype, item_size = DTYPES[dtype_name] rows, cols = (int(value) for value in entry["shape"]) offset_start, offset_end = (int(value) for value in entry["data_offsets"]) if offset_end - offset_start != rows * cols * item_size: raise ValueError("safetensors byte span does not match target shape") raw = torch.frombuffer( mm, dtype=dtype, count=rows * cols, offset=data_start + offset_start, ).view(rows, cols) scaled_a = a.float().mul(float(scaling)) changed = 0 sum_sq = 0.0 max_abs = 0.0 nonzero_formula = 0 for start in range(0, rows, tile_rows): end = min(rows, start + tile_rows) original = raw[start:end].float() formula_delta = b[start:end].float().matmul(scaled_a) merged = (original + formula_delta).to(dtype) changed += int(torch.count_nonzero(merged != raw[start:end]).item()) nonzero_formula += int(torch.count_nonzero(formula_delta).item()) actual_delta = merged.float().sub(original) if actual_delta.numel(): max_abs = max(max_abs, float(actual_delta.abs().max().item())) sum_sq += float(actual_delta.double().square().sum().item()) raw[start:end].copy_(merged) del original, formula_delta, merged, actual_delta elements = rows * cols del raw, scaled_a gc.collect() return { "elements": elements, "changed_elements": changed, "changed_fraction": changed / elements, "formula_nonzero_elements": nonzero_formula, "max_abs_effective_bf16_delta": max_abs, "rms_effective_bf16_delta": math.sqrt(sum_sq / elements), } def main() -> None: args = parse_args() if args.tile_rows <= 0 or args.threads <= 0: raise SystemExit("--tile-rows and --threads must be positive") torch.set_num_threads(args.threads) torch.set_num_interop_threads(1) base_dir = args.base_dir.resolve() adapter_dir = args.adapter_dir.resolve() output_dir = args.output_dir.resolve() adapter_path = adapter_dir / "adapter_model.safetensors" for required in (adapter_path, adapter_dir / "adapter_config.json", base_dir / "model.safetensors.index.json"): if not required.is_file(): raise SystemExit(f"missing required file: {required}") output_dir.mkdir(parents=True, exist_ok=True) progress_path = output_dir / ".merge_progress.json" existing = [path for path in output_dir.iterdir() if path.name != ".merge_progress.json"] if existing and not args.resume: raise SystemExit(f"output directory is nonempty; use --resume or a fresh path: {output_dir}") progress = json.loads(progress_path.read_text()) if progress_path.exists() else {"shards": {}} pairs, adapter_config, weight_map, headers = validate_preflight( base_dir, adapter_path, args.expected_targets ) shard_names = sorted(set(weight_map.values())) targets_by_shard: dict[str, list[str]] = {name: [] for name in shard_names} for target in pairs: targets_by_shard[weight_map[target]].append(target) for shard_name in shard_names: entries = tensor_entries(headers[shard_name][1]) targets_by_shard[shard_name].sort(key=lambda key: int(entries[key]["data_offsets"][0])) run_identity = { "base_dir": str(base_dir), "base_repo": args.base_repo, "base_revision": args.base_revision, "base_index_sha256": sha256_file(base_dir / "model.safetensors.index.json"), "adapter_dir": str(adapter_dir), "adapter_sha256": sha256_file(adapter_path), "adapter_config_sha256": sha256_file(adapter_dir / "adapter_config.json"), "targets": len(pairs), "tile_rows": args.tile_rows, } if progress.get("run_identity") not in (None, run_identity): raise SystemExit("resume identity differs from the existing merge progress") progress["run_identity"] = run_identity progress.setdefault("started_at", datetime.now(timezone.utc).isoformat()) progress.setdefault("shards", {}) atomic_json(progress_path, progress) target_stats: dict[str, dict[str, object]] = {} started = time.time() with safe_open(adapter_path, framework="pt", device="cpu") as adapter: for shard_number, shard_name in enumerate(shard_names, start=1): source = base_dir / shard_name destination = output_dir / shard_name recorded = progress["shards"].get(shard_name) if args.resume and recorded and destination.is_file(): if destination.stat().st_size == recorded["size"] and sha256_file(destination) == recorded["sha256"]: print(f"[{shard_number:02d}/{len(shard_names)}] verified completed {shard_name}", flush=True) for target, stats in recorded.get("targets", {}).items(): target_stats[target] = stats continue temporary = output_dir / (shard_name + ".partial") temporary.unlink(missing_ok=True) print( f"[{shard_number:02d}/{len(shard_names)}] copying {shard_name} " f"({source.stat().st_size / (1 << 30):.2f} GiB; {len(targets_by_shard[shard_name])} targets)", flush=True, ) shutil.copyfile(source, temporary) shutil.copystat(source, temporary) data_start, header = read_safetensors_header(temporary) entries = tensor_entries(header) shard_target_stats: dict[str, dict[str, object]] = {} with temporary.open("r+b", buffering=0) as handle: mm = mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_WRITE) try: for target_number, target in enumerate(targets_by_shard[shard_name], start=1): pair = pairs[target] a = adapter.get_tensor(pair["a_key"]) b = adapter.get_tensor(pair["b_key"]) stats = patch_tensor( mm, data_start, entries[target], a, b, float(pair["scaling"]), args.tile_rows, ) if stats["changed_elements"] == 0: raise RuntimeError(f"merge rounded to no change for {target}") stats.update( { "module": pair["module"], "base_shape": list(entries[target]["shape"]), "a_shape": pair["a_shape"], "b_shape": pair["b_shape"], "scaling": pair["scaling"], "shard": shard_name, } ) shard_target_stats[target] = stats target_stats[target] = stats del a, b gc.collect() mm.flush() print( f" target {target_number:02d}/{len(targets_by_shard[shard_name]):02d} " f"{target} changed={stats['changed_fraction']:.3%}", flush=True, ) mm.flush() os.fsync(handle.fileno()) finally: mm.close() try: os.posix_fadvise(handle.fileno(), 0, 0, os.POSIX_FADV_DONTNEED) except (AttributeError, OSError): pass os.replace(temporary, destination) shard_hash = sha256_file(destination) progress["shards"][shard_name] = { "sha256": shard_hash, "size": destination.stat().st_size, "targets": shard_target_stats, "completed_at": datetime.now(timezone.utc).isoformat(), } atomic_json(progress_path, progress) print(f" committed sha256={shard_hash[:16]}...", flush=True) if len(target_stats) != len(pairs): missing = sorted(set(pairs) - set(target_stats)) raise RuntimeError(f"only merged {len(target_stats)}/{len(pairs)} targets: {missing[:3]}") for source in base_dir.iterdir(): if source.is_file() and not source.name.endswith(".safetensors"): shutil.copy2(source, output_dir / source.name) module_counts = Counter(re.search(r"\.([^.]+)$", str(value["module"])).group(1) for value in pairs.values()) layer_counts = Counter( int(re.search(r"layers\.(\d+)\.", str(value["module"])).group(1)) for value in pairs.values() ) manifest = { "schema_version": 1, "created_at": datetime.now(timezone.utc).isoformat(), "merge_seconds": round(time.time() - started, 3), "base": { "repo": args.base_repo, "revision": args.base_revision, "local_snapshot": str(base_dir), "index_sha256": run_identity["base_index_sha256"], "shards": len(shard_names), }, "adapter": { "local_dir": str(adapter_dir), "model_sha256": run_identity["adapter_sha256"], "config_sha256": run_identity["adapter_config_sha256"], "training_args_sha256": ( sha256_file(adapter_dir / "training_args.json") if (adapter_dir / "training_args.json").is_file() else None ), "config": adapter_config, }, "merge": { "formula": "BF16(W_FP32 + B_FP32 @ (A_FP32 * (lora_alpha / r)))", "strategy": "one target tensor at a time, row-tiled, in-place in a copied shard", "tile_rows": args.tile_rows, "threads": args.threads, "target_count": len(pairs), "module_counts": dict(sorted(module_counts.items())), "layer_counts": {str(key): layer_counts[key] for key in sorted(layer_counts)}, }, "target_stats": target_stats, "shards": { name: {"sha256": value["sha256"], "size": value["size"]} for name, value in progress["shards"].items() }, "environment": { "python": sys.version, "platform": platform.platform(), "torch": torch.__version__, "safetensors": safetensors.__version__, "script_sha256": sha256_file(Path(__file__).resolve()), }, } atomic_json(output_dir / "merge_manifest.json", manifest) (output_dir / "MERGE_COMPLETE").write_text(datetime.now(timezone.utc).isoformat() + "\n") progress_path.unlink(missing_ok=True) print( f"MERGE COMPLETE: {len(target_stats)} targets, {len(shard_names)} shards -> {output_dir}", flush=True, ) if __name__ == "__main__": main()