#!/usr/bin/env python3 """Certify a tensor-wise LoRA merge without loading the model. For every LoRA target, recompute the FP32 LoRA formula in small row tiles and require bit-exact BF16 equality with the merged tensor. For every byte outside the 190 target tensor ranges, require exact equality with the pinned base checkpoint. Optionally require every target to differ from an older merge. """ from __future__ import annotations import argparse import gc import hashlib import json import math import mmap import os import re import struct import time from collections import Counter, defaultdict from datetime import datetime, timezone from pathlib import Path import torch from safetensors import safe_open def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--base-dir", type=Path, required=True) parser.add_argument("--adapter-dir", type=Path, required=True) parser.add_argument("--merged-dir", type=Path, required=True) parser.add_argument("--old-merged-dir", type=Path) parser.add_argument("--tile-rows", type=int, default=32) parser.add_argument("--threads", type=int, default=2) parser.add_argument("--compare-block-mib", type=int, default=8) parser.add_argument("--expected-targets", type=int, default=190) 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_header(path: Path) -> tuple[int, bytes, dict[str, object]]: with path.open("rb") as handle: prefix = handle.read(8) if len(prefix) != 8: raise ValueError(f"invalid safetensors file: {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}") return f"model.language_model.{match.group(0)}.weight" def discover_pairs(adapter_dir: Path) -> tuple[dict[str, dict[str, object]], float]: config = json.loads((adapter_dir / "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_dir / "adapter_model.safetensors", framework="pt", device="cpu") as adapter: keys = set(adapter.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 {b_key}") target = checkpoint_key_for(module) pairs[target] = {"module": module, "a_key": a_key, "b_key": b_key} return pairs, scaling def compare_range( base_handle, merged_handle, offset: int, length: int, block_size: int, ) -> None: base_handle.seek(offset) merged_handle.seek(offset) remaining = length while remaining: wanted = min(block_size, remaining) base_block = base_handle.read(wanted) merged_block = merged_handle.read(wanted) if len(base_block) != wanted or len(merged_block) != wanted: raise RuntimeError(f"short read while comparing at byte {offset + length - remaining}") if base_block != merged_block: start = offset + length - remaining first = next(i for i, (left, right) in enumerate(zip(base_block, merged_block)) if left != right) raise RuntimeError(f"non-target byte changed at absolute offset {start + first}") remaining -= wanted def read_bf16_matrix(mm: mmap.mmap, data_start: int, entry: dict[str, object]) -> torch.Tensor: if entry["dtype"] != "BF16": raise ValueError(f"expected BF16, got {entry['dtype']}") rows, columns = (int(value) for value in entry["shape"]) start, end = (int(value) for value in entry["data_offsets"]) if end - start != rows * columns * 2: raise ValueError("invalid BF16 tensor span") return torch.frombuffer( mm, dtype=torch.bfloat16, count=rows * columns, offset=data_start + start, ).view(rows, columns) def main() -> None: args = parse_args() if min(args.tile_rows, args.threads, args.compare_block_mib) <= 0: raise SystemExit("tile rows, threads, and compare block must be positive") torch.set_num_threads(args.threads) torch.set_num_interop_threads(1) block_size = args.compare_block_mib << 20 base_dir = args.base_dir.resolve() adapter_dir = args.adapter_dir.resolve() merged_dir = args.merged_dir.resolve() old_dir = args.old_merged_dir.resolve() if args.old_merged_dir else None manifest_path = merged_dir / "merge_manifest.json" if not manifest_path.is_file() or not (merged_dir / "MERGE_COMPLETE").is_file(): raise SystemExit("merged artifact lacks merge_manifest.json or MERGE_COMPLETE") if list(merged_dir.glob("*.partial")): raise SystemExit("merged artifact still contains partial shard files") manifest = json.loads(manifest_path.read_text()) adapter_hash = sha256_file(adapter_dir / "adapter_model.safetensors") if adapter_hash != manifest["adapter"]["model_sha256"]: raise SystemExit("adapter hash differs from merge manifest") if sha256_file(base_dir / "model.safetensors.index.json") != manifest["base"]["index_sha256"]: raise SystemExit("base index hash differs from merge manifest") base_index = json.loads((base_dir / "model.safetensors.index.json").read_text()) merged_index = json.loads((merged_dir / "model.safetensors.index.json").read_text()) if merged_index != base_index: raise SystemExit("merged model index is not identical to the pinned base index") weight_map: dict[str, str] = base_index["weight_map"] shard_names = sorted(set(weight_map.values())) if len(shard_names) != 26: raise SystemExit(f"expected 26 base shards, found {len(shard_names)}") pairs, scaling = discover_pairs(adapter_dir) if len(pairs) != args.expected_targets: raise SystemExit(f"expected {args.expected_targets} targets, found {len(pairs)}") if set(pairs) - set(weight_map): raise SystemExit(f"adapter targets missing from base: {sorted(set(pairs)-set(weight_map))[:3]}") targets_by_shard: dict[str, list[str]] = defaultdict(list) for target in pairs: targets_by_shard[weight_map[target]].append(target) target_reports: dict[str, dict[str, object]] = {} unchanged_bytes = 0 started = time.time() with safe_open(adapter_dir / "adapter_model.safetensors", framework="pt", device="cpu") as adapter: for shard_number, shard_name in enumerate(shard_names, start=1): base_path = base_dir / shard_name merged_path = merged_dir / shard_name if not merged_path.is_file() or merged_path.stat().st_size != base_path.stat().st_size: raise RuntimeError(f"missing or wrong-sized merged shard {shard_name}") expected_hash = manifest["shards"][shard_name]["sha256"] actual_hash = sha256_file(merged_path) if actual_hash != expected_hash: raise RuntimeError(f"merged shard hash mismatch: {shard_name}") base_data_start, base_header_bytes, base_header = read_header(base_path) merged_data_start, merged_header_bytes, merged_header = read_header(merged_path) if base_data_start != merged_data_start or base_header_bytes != merged_header_bytes: raise RuntimeError(f"safetensors header changed in {shard_name}") base_entries = entries(base_header) merged_entries = entries(merged_header) if base_entries != merged_entries: raise RuntimeError(f"tensor index changed in {shard_name}") target_ranges = [] for target in targets_by_shard[shard_name]: start, end = (int(value) for value in base_entries[target]["data_offsets"]) target_ranges.append((base_data_start + start, base_data_start + end, target)) target_ranges.sort() cursor = 0 with base_path.open("rb", buffering=0) as base_handle, merged_path.open("rb", buffering=0) as merged_handle: for start, end, _ in target_ranges: compare_range(base_handle, merged_handle, cursor, start - cursor, block_size) unchanged_bytes += start - cursor cursor = end compare_range(base_handle, merged_handle, cursor, base_path.stat().st_size - cursor, block_size) unchanged_bytes += base_path.stat().st_size - cursor for handle in (base_handle, merged_handle): try: os.posix_fadvise(handle.fileno(), 0, 0, os.POSIX_FADV_DONTNEED) except (AttributeError, OSError): pass old_path = old_dir / shard_name if old_dir else None with base_path.open("rb", buffering=0) as base_file, merged_path.open("rb", buffering=0) as merged_file: base_mm = mmap.mmap(base_file.fileno(), 0, access=mmap.ACCESS_READ) merged_mm = mmap.mmap(merged_file.fileno(), 0, access=mmap.ACCESS_READ) old_file = old_path.open("rb", buffering=0) if old_path and old_path.is_file() else None old_mm = mmap.mmap(old_file.fileno(), 0, access=mmap.ACCESS_READ) if old_file else None old_entries = entries(read_header(old_path)[2]) if old_path and old_path.is_file() else None 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"]) base_tensor = read_bf16_matrix(base_mm, base_data_start, base_entries[target]) merged_tensor = read_bf16_matrix(merged_mm, merged_data_start, merged_entries[target]) old_tensor = ( read_bf16_matrix(old_mm, read_header(old_path)[0], old_entries[target]) if old_mm is not None and old_entries is not None else None ) scaled_a = a.float().mul(scaling) changed = 0 differs_old = 0 max_abs = 0.0 sum_sq = 0.0 rows, columns = base_tensor.shape for start in range(0, rows, args.tile_rows): end = min(rows, start + args.tile_rows) expected = ( base_tensor[start:end].float() + b[start:end].float().matmul(scaled_a) ).to(torch.bfloat16) if not torch.equal(expected, merged_tensor[start:end]): mismatches = int(torch.count_nonzero(expected != merged_tensor[start:end]).item()) raise RuntimeError(f"formula mismatch for {target}: {mismatches} elements") changed += int(torch.count_nonzero(merged_tensor[start:end] != base_tensor[start:end]).item()) if old_tensor is not None: differs_old += int(torch.count_nonzero(merged_tensor[start:end] != old_tensor[start:end]).item()) actual = merged_tensor[start:end].float() - base_tensor[start:end].float() max_abs = max(max_abs, float(actual.abs().max().item())) sum_sq += float(actual.double().square().sum().item()) del expected, actual if changed == 0: raise RuntimeError(f"merged target is identical to base: {target}") if old_tensor is not None and differs_old == 0: raise RuntimeError(f"round-2 merged target is identical to round 1: {target}") elements = rows * columns target_reports[target] = { "module": pair["module"], "shard": shard_name, "shape": [rows, columns], "changed_from_base_elements": changed, "changed_from_base_fraction": changed / elements, "changed_from_round1_elements": differs_old if old_tensor is not None else None, "max_abs_effective_bf16_delta": max_abs, "rms_effective_bf16_delta": math.sqrt(sum_sq / elements), "formula_bit_exact": True, } del a, b, base_tensor, merged_tensor, old_tensor, scaled_a gc.collect() print( f"[{shard_number:02d}/{len(shard_names)}] target " f"{target_number:02d}/{len(targets_by_shard[shard_name]):02d} verified {target}", flush=True, ) finally: base_mm.close() merged_mm.close() if old_mm is not None: old_mm.close() if old_file is not None: old_file.close() print( f"[{shard_number:02d}/{len(shard_names)}] {shard_name}: hash, non-target bytes, and formulas verified", flush=True, ) if set(target_reports) != set(pairs): raise RuntimeError(f"verified {len(target_reports)}/{len(pairs)} targets") module_counts = Counter(re.search(r"\.([^.]+)$", value["module"]).group(1) for value in pairs.values()) layer_counts = Counter(int(re.search(r"layers\.(\d+)\.", value["module"]).group(1)) for value in pairs.values()) expected_module_counts = { "in_proj_a": 30, "in_proj_b": 30, "in_proj_qkv": 30, "in_proj_z": 30, "out_proj": 30, "k_proj": 10, "o_proj": 10, "q_proj": 10, "v_proj": 10, } if dict(module_counts) != expected_module_counts: raise RuntimeError(f"unexpected target-module coverage: {dict(module_counts)}") if set(layer_counts) != set(range(40)): raise RuntimeError(f"not all 40 transformer layers are targeted: {dict(layer_counts)}") metadata_loads: dict[str, str] = {} from transformers import AutoConfig, AutoProcessor, AutoTokenizer config = AutoConfig.from_pretrained(merged_dir, local_files_only=True) metadata_loads["config_class"] = type(config).__name__ tokenizer = AutoTokenizer.from_pretrained(merged_dir, local_files_only=True) metadata_loads["tokenizer_class"] = type(tokenizer).__name__ processor = AutoProcessor.from_pretrained(merged_dir, local_files_only=True) metadata_loads["processor_class"] = type(processor).__name__ report = { "schema_version": 1, "verified_at": datetime.now(timezone.utc).isoformat(), "verification_seconds": round(time.time() - started, 3), "base_revision": manifest["base"]["revision"], "adapter_sha256": adapter_hash, "merged_shards": len(shard_names), "base_tensor_count": len(weight_map), "target_tensor_count": len(target_reports), "non_target_tensor_count": len(weight_map) - len(target_reports), "unchanged_bytes_compared": unchanged_bytes, "all_non_target_bytes_identical": True, "all_targets_formula_bit_exact": True, "all_targets_distinct_from_base": True, "all_targets_distinct_from_round1": old_dir is not None, "module_counts": dict(sorted(module_counts.items())), "layer_counts": {str(key): layer_counts[key] for key in sorted(layer_counts)}, "metadata_loads": metadata_loads, "target_reports": target_reports, } atomic_json(merged_dir / "merge_verification.json", report) (merged_dir / "VERIFIED").write_text(datetime.now(timezone.utc).isoformat() + "\n") print(json.dumps({key: value for key, value in report.items() if key != "target_reports"}, indent=2)) print(f"VERIFIED: {merged_dir}", flush=True) if __name__ == "__main__": main()