merge_tools/merge_lora_tensorwise.py

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


def tensor_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!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()