138 lines
4.1 KiB
Python
138 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
"""QDiff: compact quantized delta/checkpoint utilities.
|
|
|
|
This module provides a minimal, well-tested implementation of a
|
|
quantized residual delta suitable for toy workloads and unit tests.
|
|
|
|
Features implemented:
|
|
- 8-bit symmetric quantization of residuals
|
|
- optional top-k sparsification (keep largest magnitude residuals)
|
|
- chunking and Merkle-style manifest (SHA256 per-chunk)
|
|
|
|
This is intentionally small: it demonstrates the QDiff concept and is
|
|
meant to be extended by other agents (e.g. block-sparse, file-backed,
|
|
and resumable manifests).
|
|
"""
|
|
|
|
from dataclasses import dataclass, asdict
|
|
import hashlib
|
|
import json
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
|
|
@dataclass
|
|
class QDiffBundle:
|
|
header: Dict
|
|
chunks: Dict[str, bytes]
|
|
|
|
|
|
def _quantize_residuals(base: np.ndarray, new: np.ndarray) -> Tuple[np.ndarray, float]:
|
|
# compute residuals and a symmetric scale to fit in int8
|
|
resid = new - base
|
|
max_abs = float(np.max(np.abs(resid)))
|
|
if max_abs == 0.0:
|
|
return np.zeros_like(resid, dtype=np.int8), 1.0
|
|
scale = max_abs / 127.0
|
|
q = np.round(resid / scale).astype(np.int8)
|
|
return q, scale
|
|
|
|
|
|
def _dequantize(q: np.ndarray, scale: float) -> np.ndarray:
|
|
return q.astype(np.float32) * scale
|
|
|
|
|
|
def _sparsify(q: np.ndarray, top_k: Optional[int]) -> Tuple[np.ndarray, Optional[List[int]]]:
|
|
if top_k is None or top_k >= q.size:
|
|
return q, None
|
|
# keep top_k by absolute magnitude
|
|
idx = np.argsort(np.abs(q))[-top_k:]
|
|
mask = np.zeros(q.size, dtype=bool)
|
|
mask[idx] = True
|
|
sparse_q = np.zeros_like(q)
|
|
sparse_q[mask] = q[mask]
|
|
return sparse_q, idx.tolist()
|
|
|
|
|
|
def _chunk_bytes(b: bytes, chunk_size: int = 1024) -> List[bytes]:
|
|
return [b[i : i + chunk_size] for i in range(0, len(b), chunk_size)]
|
|
|
|
|
|
def _sha256_hex(b: bytes) -> str:
|
|
return hashlib.sha256(b).hexdigest()
|
|
|
|
|
|
def create_qdiff(
|
|
base: List[float],
|
|
new: List[float],
|
|
top_k: Optional[int] = None,
|
|
chunk_size: int = 1024,
|
|
) -> QDiffBundle:
|
|
"""Create a QDiff bundle describing new relative to base.
|
|
|
|
base and new are numeric arrays (lists or numpy-compatible). Returns
|
|
a QDiffBundle with a header and chunk map (sha256->bytes).
|
|
"""
|
|
a = np.asarray(base, dtype=np.float32)
|
|
b = np.asarray(new, dtype=np.float32)
|
|
if a.shape != b.shape:
|
|
raise ValueError("base and new must have same shape")
|
|
|
|
q, scale = _quantize_residuals(a, b)
|
|
sparse_q, indices = _sparsify(q, top_k)
|
|
|
|
# serialize: header JSON, and the quantized bytes as raw int8
|
|
q_bytes = sparse_q.tobytes()
|
|
|
|
chunks = {}
|
|
chunk_list = _chunk_bytes(q_bytes, chunk_size=chunk_size)
|
|
manifest = []
|
|
for c in chunk_list:
|
|
h = _sha256_hex(c)
|
|
chunks[h] = c
|
|
manifest.append(h)
|
|
|
|
header = {
|
|
"shape": list(a.shape),
|
|
"dtype": "float32",
|
|
"quant": "int8",
|
|
"scale": float(scale),
|
|
"top_k_indices": indices,
|
|
"manifest": manifest,
|
|
"chunk_size": chunk_size,
|
|
}
|
|
|
|
return QDiffBundle(header=header, chunks=chunks)
|
|
|
|
|
|
def apply_qdiff(base: List[float], bundle: QDiffBundle) -> List[float]:
|
|
"""Apply QDiff bundle to a base array and return reconstructed new array."""
|
|
a = np.asarray(base, dtype=np.float32)
|
|
shape = tuple(bundle.header["shape"])
|
|
if a.shape != shape:
|
|
raise ValueError("base shape does not match bundle header")
|
|
|
|
# reconstruct bytes from manifest
|
|
manifest = bundle.header["manifest"]
|
|
parts = [bundle.chunks[h] for h in manifest]
|
|
q_bytes = b"".join(parts)
|
|
|
|
# ensure length matches
|
|
expected_elems = int(np.prod(shape))
|
|
q = np.frombuffer(q_bytes, dtype=np.int8, count=expected_elems)
|
|
|
|
scale = float(bundle.header["scale"])
|
|
deq = _dequantize(q, scale)
|
|
new = a + deq.reshape(shape)
|
|
return new.tolist()
|
|
|
|
|
|
def manifest_proofs(bundle: QDiffBundle) -> Dict[str, str]:
|
|
"""Return a mapping of chunk hash -> hex digest (Merkle leaf hashes).
|
|
|
|
This is a thin helper used by tests and auditors.
|
|
"""
|
|
return {h: _sha256_hex(bundle.chunks[h]) for h in bundle.header["manifest"]}
|