Spaces:
Sleeping
Sleeping
| """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 | |
| 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) | |