diff --git a/nebulaforge/qdiff.py b/nebulaforge/qdiff.py new file mode 100644 index 0000000..eabdd77 --- /dev/null +++ b/nebulaforge/qdiff.py @@ -0,0 +1,137 @@ +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"]} diff --git a/tests/test_qdiff.py b/tests/test_qdiff.py new file mode 100644 index 0000000..a74bb51 --- /dev/null +++ b/tests/test_qdiff.py @@ -0,0 +1,32 @@ +from nebulaforge.qdiff import create_qdiff, apply_qdiff, manifest_proofs + + +def test_qdiff_create_and_apply_identity(): + base = [0.0, 1.0, 2.0, -1.0] + new = [0.0, 1.0, 2.0, -1.0] + bundle = create_qdiff(base, new) + out = apply_qdiff(base, bundle) + # exact equality for identical arrays + assert out == new + + +def test_qdiff_quantized_residuals_and_topk(): + base = [0.0, 0.0, 0.0, 0.0] + new = [0.0, 10.0, -5.0, 2.0] + # keep only top-2 residuals + bundle = create_qdiff(base, new, top_k=2, chunk_size=8) + out = apply_qdiff(base, bundle) + # Reconstructed values should be close to original for top-k kept; others may be zero + # Check that at least two elements match closely + matches = sum(1 for a, b in zip(out, new) if abs(a - b) < 1e-2) + assert matches >= 2 + + +def test_manifest_hashes(): + base = [0.0] * 16 + new = [i * 0.1 for i in range(16)] + bundle = create_qdiff(base, new, chunk_size=7) + proofs = manifest_proofs(bundle) + # every manifest entry should have a corresponding proof equal to the key + for k, v in proofs.items(): + assert k == v