Repack a split GGUF so the qwen4exp PLE/engram table sits alone in the final shard (llama.cpp PR 27742)

repack-ple-shard.py

#!/usr/bin/env python3
# Repack a split GGUF so one tensor (default: the qwen4exp engram table
# per_layer_token_embd.weight) sits ALONE in the final shard, with no other tensors
# sharing that file. Motivation: on Metal, llama.cpp wires the mmap'd regions that
# back GPU tensors; a tensor interleaved with them in the same file gets wired along
# for the ride (measured +24 GiB on Qwen3.8-Flash-Next). A file containing only
# CPU-side tensors keeps its own mapping and stays pageable.
#
# Tensor bytes are copied verbatim (no requantization). Output layout for an N-file
# input: shard 1 = byte-for-byte copy of input shard 1 (metadata-only; split.count and
# split.tensors.count are unchanged so it stays valid), shards 2..N-1 = every other
# tensor in original global order, shard N = the isolated tensor alone. Data shards
# carry exactly the three split.* KVs, with the same GGUF value types the official
# llama-gguf-split writes (split.no/count UINT16, split.tensors.count INT32).
#
# Run with the gguf-py matching the model's llama.cpp tree:
#   uv run --no-project --with ~/AI/llama.cpp-pr27742/gguf-py python repack-ple-shard.py \
#       ~/AI/models/Qwen3.8-Flash-Next-Unsloth/Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf
#   ... --verify   (after a repack: byte-compares every tensor against the input)

import argparse
import hashlib
import os
import re
import shutil
import sys
from pathlib import Path

import numpy as np
from gguf import GGUFReader, GGUFWriter
from gguf.constants import GGUFValueType

SPLIT_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
SPLIT_KV_TYPES = {
    "split.no": GGUFValueType.UINT16,
    "split.tensors.count": GGUFValueType.INT32,
    "split.count": GGUFValueType.UINT16,
}


def shard_paths(first_shard: Path) -> tuple[str, list[Path]]:
    m = SPLIT_RE.match(first_shard.name)
    if not m or int(m.group(2)) != 1:
        sys.exit(f"ERROR: {first_shard.name} is not a '-00001-of-NNNNN.gguf' first shard")
    stem, count = m.group(1), int(m.group(3))
    paths = [first_shard.parent / f"{stem}-{i:05d}-of-{count:05d}.gguf" for i in range(1, count + 1)]
    for p in paths:
        if not p.is_file():
            sys.exit(f"ERROR: missing input shard {p}")
    return stem, paths


def out_paths(first_shard: Path, stem: str, suffix: str, count: int) -> list[Path]:
    return [first_shard.parent / f"{stem}-{suffix}-{i:05d}-of-{count:05d}.gguf" for i in range(1, count + 1)]


def read_all(paths: list[Path]) -> list[GGUFReader]:
    return [GGUFReader(p) for p in paths]


def kv_int(reader: GGUFReader, key: str) -> int:
    field = reader.fields.get(key)
    if field is None:
        sys.exit(f"ERROR: {key} missing from {reader.data.filename if hasattr(reader.data, 'filename') else 'shard'} — not a split GGUF?")
    return int(field.contents())


def file_sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(1 << 22):
            h.update(chunk)
    return h.hexdigest()


def bytes_equal(a: np.ndarray, b: np.ndarray, chunk: int = 1 << 28) -> bool:
    # Compare as raw bytes, chunked: a whole-tensor `a == b` materializes a bool array
    # the size of the tensor (26.8 GiB for the PLE), and value comparison on floats
    # treats +0.0 == -0.0 and NaN != NaN — byte view avoids both.
    av = a.reshape(-1).view(np.uint8)
    bv = b.reshape(-1).view(np.uint8)
    if av.nbytes != bv.nbytes:
        return False
    return all(np.array_equal(av[off:off + chunk], bv[off:off + chunk]) for off in range(0, av.nbytes, chunk))


def plan_shards(readers: list[GGUFReader], isolate: str, n_data_shards: int):
    # Global tensor order = file order across data shards; that order is preserved in
    # the output so nothing but file membership changes.
    all_tensors = [t for r in readers for t in r.tensors]
    iso = [t for t in all_tensors if t.name == isolate]
    if len(iso) != 1:
        sys.exit(f"ERROR: expected exactly one tensor named {isolate!r}, found {len(iso)}")
    rest = [t for t in all_tensors if t.name != isolate]

    def padded(t):  # on-disk footprint, 32-byte GGUF alignment
        return (int(t.n_bytes) + 31) // 32 * 32

    total = sum(padded(t) for t in rest)
    shards: list[list] = [[] for _ in range(n_data_shards)]
    target = total / n_data_shards
    idx, acc = 0, 0
    for t in rest:
        # advance once the current shard is at target; the last shard takes the tail
        if acc >= target and idx < n_data_shards - 1 and shards[idx]:
            idx += 1
            acc = 0
        shards[idx].append(t)
        acc += padded(t)
    if any(not s for s in shards):
        sys.exit("ERROR: shard packing produced an empty data shard — lower the shard count")
    return all_tensors, shards, iso[0]


def write_data_shard(path: Path, tensors, split_no: int, split_count: int, split_tensors_total: int) -> None:
    # arch is a constructor formality: add_architecture() is never called, so the only
    # KVs in the file are the three split.* keys, exactly like llama-gguf-split output.
    w = GGUFWriter(path=None, arch="")
    w.add_key_value("split.no", split_no, SPLIT_KV_TYPES["split.no"])
    w.add_key_value("split.tensors.count", split_tensors_total, SPLIT_KV_TYPES["split.tensors.count"])
    w.add_key_value("split.count", split_count, SPLIT_KV_TYPES["split.count"])
    for t in tensors:
        # Contract with gguf-py: ReaderTensor.shape is GGUF file order (ne); the writer
        # serializes reversed(shape), so pass reversed dims. The dummy non-uint8 dtype
        # skips the writer's byte-shape-to-quant-shape conversion (we already have
        # element dims), and raw_dtype carries the real quant type through untouched.
        w.add_tensor_info(
            t.name,
            tuple(reversed(t.shape.tolist())),
            np.dtype(np.float32),
            int(t.n_bytes),
            raw_dtype=t.tensor_type,
        )
    w.write_header_to_file(path=path)
    w.write_kv_data_to_file()
    w.write_ti_data_to_file()
    for i, t in enumerate(tensors):
        w.write_tensor_data(t.data)  # memmap-backed: streams, no full-tensor RAM copy
        if (i + 1) % 200 == 0 or i + 1 == len(tensors):
            print(f"    {path.name}: {i + 1}/{len(tensors)} tensors", flush=True)
    w.close()


def do_repack(args, first: Path, stem: str, in_paths: list[Path]) -> None:
    readers = read_all(in_paths)
    split_count = kv_int(readers[0], "split.count")
    split_total = kv_int(readers[0], "split.tensors.count")
    if split_count != len(in_paths):
        sys.exit(f"ERROR: split.count={split_count} but {len(in_paths)} files found")
    if readers[0].tensors:
        sys.exit("ERROR: input shard 1 is not metadata-only; this script relies on copying it verbatim")

    n_data_shards = len(in_paths) - 2  # everything but the metadata shard and the isolate shard
    if n_data_shards < 1:
        sys.exit("ERROR: need at least 3 input shards (metadata + 1 data + isolate)")
    all_tensors, shards, iso = plan_shards(readers[1:], args.isolate, n_data_shards)
    if len(all_tensors) != split_total:
        sys.exit(f"ERROR: read {len(all_tensors)} tensors but split.tensors.count={split_total}")

    outs = out_paths(first, stem, args.suffix, len(in_paths))
    existing = [p for p in outs if p.exists()]
    if existing and not args.force:
        sys.exit(f"ERROR: output exists (use --force to overwrite): {existing[0]}")

    need = sum(int(t.n_bytes) for t in all_tensors) + in_paths[0].stat().st_size
    free = shutil.disk_usage(first.parent).free
    if free < need * 1.02:
        sys.exit(f"ERROR: needs ~{need / 2**30:.1f} GiB free in {first.parent}, only {free / 2**30:.1f} GiB available")

    print(f"Plan: {len(all_tensors)} tensors -> {len(outs)} files (suffix -{args.suffix})")
    for i, s in enumerate(shards):
        print(f"  shard {i + 2}: {len(s)} tensors, {sum(int(t.n_bytes) for t in s) / 2**30:.2f} GiB")
    print(f"  shard {len(outs)}: 1 tensor ({args.isolate}), {int(iso.n_bytes) / 2**30:.2f} GiB  <- isolated")
    if args.dry_run:
        return

    # Stage to .tmp names and publish only when every shard is complete: a failure
    # mid-write must neither leave partial final-named files nor (with --force)
    # destroy a previously good output set.
    tmps = [p.with_name(p.name + ".tmp") for p in outs]
    try:
        print(f"  shard 1: verbatim copy of {in_paths[0].name}")
        shutil.copyfile(in_paths[0], tmps[0])
        for i, s in enumerate(shards):
            write_data_shard(tmps[i + 1], s, split_no=i + 1, split_count=split_count, split_tensors_total=split_total)
        write_data_shard(tmps[-1], [iso], split_no=len(outs) - 1, split_count=split_count, split_tensors_total=split_total)
    except BaseException:
        for t in tmps:
            t.unlink(missing_ok=True)
        raise
    for t, p in zip(tmps, outs):
        os.replace(t, p)
    print("Repack written. Run with --verify before using it.")


def do_verify(args, first: Path, stem: str, in_paths: list[Path]) -> None:
    outs = out_paths(first, stem, args.suffix, len(in_paths))
    for p in outs:
        if not p.is_file():
            sys.exit(f"ERROR: missing output shard {p}")
    fails = 0
    if file_sha256(in_paths[0]) != file_sha256(outs[0]):
        print("FAIL: shard 1 differs from input shard 1")
        fails += 1
    in_readers = read_all(in_paths[1:])
    out_readers = read_all(outs[1:])

    for i, r in enumerate(out_readers):
        want_no = i + 1
        if kv_int(r, "split.no") != want_no or kv_int(r, "split.count") != len(outs):
            print(f"FAIL: bad split KVs in {outs[i + 1].name}")
            fails += 1
        for key, want_type in SPLIT_KV_TYPES.items():
            # llama.cpp reads these with fixed types; the right value in the wrong
            # encoding still fails to load
            types = r.fields[key].types if key in r.fields else []
            if types != [want_type]:
                print(f"FAIL: {key} in {outs[i + 1].name} has types {types}, want [{want_type}]")
                fails += 1
        # llama.cpp requires offsets equal to the cumulative 32-byte-padded sizes; a
        # file with relocated data would read back fine in Python but be rejected there
        prev_end = None
        for t in r.tensors:
            off = int(t.data_offset)
            if off % 32 != 0 or (prev_end is not None and off != prev_end):
                print(f"FAIL: non-canonical data offset for {t.name} in {outs[i + 1].name}")
                fails += 1
                break
            prev_end = (off + int(t.n_bytes) + 31) // 32 * 32
    iso_shard = out_readers[-1]
    if len(iso_shard.tensors) != 1 or iso_shard.tensors[0].name != args.isolate:
        print(f"FAIL: final shard does not contain exactly [{args.isolate}]")
        fails += 1

    src = {t.name: t for r in in_readers for t in r.tensors}
    dst = {}
    for r in out_readers:
        for t in r.tensors:
            if t.name in dst:
                print(f"FAIL: duplicate tensor {t.name} across output shards")
                fails += 1
            dst[t.name] = t
    split_total = kv_int(out_readers[0], "split.tensors.count")
    if len(dst) != split_total:
        print(f"FAIL: {len(dst)} tensors across output shards, split.tensors.count says {split_total}")
        fails += 1
    if src.keys() != dst.keys():
        print(f"FAIL: tensor sets differ (in-only: {sorted(src.keys() - dst.keys())[:3]}, out-only: {sorted(dst.keys() - src.keys())[:3]})")
        fails += 1
    checked = 0
    for name in src.keys() & dst.keys():
        a, b = src[name], dst[name]
        if a.tensor_type != b.tensor_type or a.shape.tolist() != b.shape.tolist() or int(a.n_bytes) != int(b.n_bytes):
            print(f"FAIL: meta mismatch on {name}")
            fails += 1
            continue
        if not bytes_equal(a.data, b.data):
            print(f"FAIL: byte mismatch on {name}")
            fails += 1
        checked += 1
        if checked % 200 == 0:
            print(f"    verified {checked}/{len(src)} tensors", flush=True)
    if fails:
        sys.exit(f"VERIFY FAILED: {fails} problem(s)")
    print(f"VERIFY PASSED: {checked} tensors byte-identical, shard 1 identical, split KVs correct.")


def main() -> None:
    ap = argparse.ArgumentParser(description="Isolate one tensor into its own GGUF split shard (byte-exact repack)")
    ap.add_argument("first_shard", type=Path, help="path to the -00001-of-NNNNN.gguf input shard")
    ap.add_argument("--isolate", default="per_layer_token_embd.weight")
    ap.add_argument("--suffix", default="PLESHARD", help="inserted into output filenames before the shard numbering")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--verify", action="store_true", help="verify a previous repack instead of writing one")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    first = args.first_shard.resolve()
    stem, in_paths = shard_paths(first)
    if args.verify:
        do_verify(args, first, stem, in_paths)
    else:
        do_repack(args, first, stem, in_paths)


if __name__ == "__main__":
    main()