Blame for tensor_pruning/checkpoint.py

"""Stage Llama-layout query ranges from safetensors and select one global mask."""

from dataclasses import dataclass
import json
from pathlib import Path
import shutil
import struct
import tempfile

import ml_dtypes
import numpy as np

from .files import _input_path, _open_npy, _publish_masks


_DTYPES = {"F16": np.dtype("<f2"), "BF16": np.dtype(ml_dtypes.bfloat16),
           "F32": np.dtype("<f4"), "F64": np.dtype("<f8")}
_CHUNK = 8 * 1024**2
_HEADROOM = 1024**3


def _object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError(f"Duplicate JSON key: {key}")
        result[key] = value
    return result


def _json_object(raw):
    value = json.loads(raw, object_pairs_hook=_object)
    if not isinstance(value, dict):
        raise ValueError("Expected a JSON object")
    return value


def _json_file(path, limit):
    with path.open("rb") as stream:
        raw = stream.read(limit + 1)
    if len(raw) > limit:
        raise ValueError(f"{path.name} exceeds the supported metadata size")
    return _json_object(raw)


def _layout(config):
    layers = config.get("num_hidden_layers")
    heads = config.get("num_attention_heads")
    width = config.get("hidden_size")
    if any(type(n) is not int or n < 1 for n in (layers, heads, width)) or layers > 1024:
        raise ValueError("Expected valid Llama-layout layer, head and hidden sizes")
    channels = config.get("head_dim")
    if channels is None:
        channels, extra = divmod(width, heads)
        if extra:
            raise ValueError("hidden_size must be divisible by num_attention_heads")
    if type(channels) is not int or channels < 1:
        raise ValueError("Invalid head_dim")
    return layers, (heads, channels, width)


def _header(path):
    with path.open("rb") as stream:
        prefix = stream.read(8)
        if len(prefix) != 8:
            raise ValueError(f"Incomplete safetensors header: {path.name}")
        size, = struct.unpack("<Q", prefix)
        if not 2 <= size <= 100_000_000:
            raise ValueError(f"Invalid safetensors header length: {path.name}")
        raw = stream.read(size)
    if len(raw) != size or not raw.startswith(b"{"):
        raise ValueError(f"Invalid safetensors header: {path.name}")
    start = 8 + size
    return _json_object(raw), start, path.stat().st_size - start


@dataclass(frozen=True)
class QueryRange:
    name: str
    shard: Path
    offset: int
    length: int
    shape: tuple[int, int, int]
    dtype: np.dtype


def _queries(checkpoint, max_weights):
    root = Path(checkpoint).resolve(strict=True)
    layers, shape = _layout(_json_file(_input_path(root, "config.json", ".json"), 1024**2))
    heads, channels, width = shape
    total = layers * heads * channels * width
    if total > max_weights:
        raise ValueError("Checkpoint exceeds the configured eligible weight limit")
    names = [f"model.layers.{layer}.self_attn.q_proj.weight" for layer in range(layers)]
    index = root / "model.safetensors.index.json"
    if index.exists():
        mapping = _json_file(_input_path(root, index.name, ".json"), 16 * 1024**2).get("weight_map")
        if not isinstance(mapping, dict) or any(name not in mapping for name in names):
            raise ValueError("Expected one named Llama-layout query tensor for every layer")
    else:
        mapping = dict.fromkeys(names, "model.safetensors")

    headers = {}
    queries = []
    ranges = {}
    for name in names:
        shard = _input_path(root, mapping[name], ".safetensors")
        if shard not in headers:
            headers[shard] = _header(shard)
        header, start, payload_size = headers[shard]
        item = header.get(name)
        if not isinstance(item, dict) or item.get("shape") != [heads * channels, width]:
            raise ValueError(f"Unexpected query shape: {name}")
        dtype = _DTYPES.get(item.get("dtype"))
        if dtype is None:
            raise ValueError(f"Expected dense floating-point query weights: {name}")
        offsets = item.get("data_offsets")
        length = heads * channels * width * dtype.itemsize
        if (not isinstance(offsets, list) or len(offsets) != 2
                or any(type(n) is not int for n in offsets)):
            raise ValueError(f"Invalid safetensors offsets: {name}")
        low, high = offsets
        if not 0 <= low < high <= payload_size or high - low != length:
            raise ValueError(f"Invalid safetensors byte range: {name}")
        ranges.setdefault(shard, []).append((low, high))
        queries.append(QueryRange(name, shard, start + low, length, shape, dtype))
    for shard, spans in ranges.items():
        ordered = sorted(spans)
        if any(left[1] > right[0] for left, right in zip(ordered, ordered[1:])):
            raise ValueError(f"Overlapping query byte ranges: {shard.name}")
    return queries, total


def _copy_query(query, path):
    with query.shard.open("rb") as src, path.open("wb") as dst:
        src.seek(query.offset)
        remaining = query.length
        while remaining:
            block = src.read(min(remaining, _CHUNK))
            if not block:
                raise OSError(f"Shard ended during staging: {query.shard.name}")
            dst.write(block)
            remaining -= len(block)
    return np.memmap(path, dtype=query.dtype, mode="r", shape=query.shape)


def prune_checkpoint(checkpoint, output_dir, *, keep, workspace, max_weights,
                     input_root=None, moments_manifest=None, threads=1, seed=0):
    """Read only the eligible query bytes and write per-layer Boolean masks."""
    queries, total = _queries(checkpoint, max_weights)
    if not 0 <= keep <= total:
        raise ValueError("keep must be between zero and the eligible query slot count")
    names = [query.name for query in queries]
    moments = None
    if moments_manifest is not None:
        root = Path(input_root).resolve(strict=True)
        sources = _json_file(_input_path(root, moments_manifest, ".json"), 1024**2)
        if set(sources) != set(names):
            raise ValueError("Moment manifest must map every eligible query name to an .npy vector")
        moments = [_open_npy(root, sources[name]) for name in names]
    workspace = Path(workspace)
    output_dir = Path(output_dir)
    if workspace.stat().st_dev != output_dir.parent.stat().st_dev:
        raise ValueError("Checkpoint work and output must use the same local filesystem")
    needed = sum(query.length for query in queries) + total + 4096 * len(queries)
    if 0 < keep < total:
        needed += 4 * total
    if shutil.disk_usage(workspace).free < needed + _HEADROOM:
        raise OSError(f"Insufficient local disk: need {needed + _HEADROOM:,} bytes free")
    with tempfile.TemporaryDirectory(prefix="query-staging-", dir=workspace) as stage:
        staged = {}
        for index, query in sorted(enumerate(queries), key=lambda row: (row[1].shard, row[1].offset)):
            staged[index] = _copy_query(query, Path(stage) / f"{index:04d}.bin")
        arrays = [staged[index] for index in range(len(queries))]
        return _publish_masks(names, arrays, moments, keep, output_dir,
                              scratch_dir=workspace, threads=threads, seed=seed)