merge_tools/verify_lora_tensorwise.py

#!/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("<Q", prefix)[0]
        header_bytes = handle.read(header_len)
    return 8 + header_len, prefix + header_bytes, json.loads(header_bytes)


def entries(header: dict[str, object]) -> 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()