mirror of
https://github.com/Dynamis-Labs/spectralquant.git
synced 2026-10-21 09:49:45 -05:00
Ashwin GopinathandAnirudh B. Vangara
19ebecbc4d
Upgrades SpectralQuant from uniform bit allocation to greedy water-filling inside the semantic subspace, with per-dimension Lloyd-Max codebooks fit to the empirical scalar marginal. Adds full three-way evaluation harness (perplexity, generation, latency, LongBench) on Modal H200 GPUs, comprehensive documentation, and the consolidated arxiv-ready manuscript. v1 is preserved on the 'v1' branch. Co-Authored-By: Anirudh B. Vangara <anirudh@sentra.app>
554 lines
20 KiB
Python
554 lines
20 KiB
Python
"""Shared utilities for the v2 evaluation harnesses.
|
|
|
|
The four next-stage evidence families (perplexity, LongBench, real
|
|
generation quality, end-to-end latency) share the same operational
|
|
discipline established by ``experiments/run_three_way.py``:
|
|
|
|
* ``--dry-run`` / ``--synthetic-smoke`` / ``--inline-corpus-smoke``
|
|
modes that need *no* heavy model download.
|
|
* Atomic JSON writes with schema validation before the rename.
|
|
* Status artifact emission via :class:`experiments.run_status.StatusWriter`.
|
|
* Provenance metadata: repo slug, git commit (with ``SPECTRALQUANT_GIT_COMMIT``
|
|
override for Modal containers), software/hardware blocks, sanitized
|
|
command string.
|
|
* Token-name redaction in logs and stdout/stderr tails.
|
|
|
|
This module factors that common surface into one place so the family
|
|
harnesses below stay short and testable. None of these helpers import
|
|
torch / transformers / datasets; lazy imports stay inside the harness
|
|
that needs them.
|
|
|
|
There are NO secret reads in this module. Token env-var **names** are
|
|
referenced for redaction only.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import platform
|
|
import re
|
|
import shlex
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
|
|
|
|
EVAL_COMMON_DIR = Path(__file__).resolve().parent
|
|
REPO_ROOT = EVAL_COMMON_DIR.parent
|
|
SRC_DIR = REPO_ROOT / "src"
|
|
SCHEMAS_DIR = REPO_ROOT / "schemas"
|
|
|
|
REPO_SLUG = "niashwin/spectralquant-full"
|
|
GIT_COMMIT_ENV = "SPECTRALQUANT_GIT_COMMIT"
|
|
|
|
# Schema version this module emits. Bump when payload shape changes in a
|
|
# backward-incompatible way.
|
|
SCHEMA_VERSION = "1"
|
|
|
|
#: Modes a harness can execute in. ``full`` is the only mode that produces
|
|
#: paper-valid evidence; everything else is harness-validation only.
|
|
MODE_FULL = "full"
|
|
MODE_SYNTHETIC_SMOKE = "synthetic_smoke"
|
|
MODE_INLINE_CORPUS_SMOKE = "inline_corpus_smoke"
|
|
MODE_DRY_RUN = "dry_run"
|
|
|
|
ALL_MODES = (MODE_FULL, MODE_SYNTHETIC_SMOKE, MODE_INLINE_CORPUS_SMOKE, MODE_DRY_RUN)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Provenance helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def utc_now_iso() -> str:
|
|
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
|
|
_SHA_RE = re.compile(r"^[0-9a-fA-F]+$")
|
|
|
|
|
|
def git_commit(
|
|
repo_root: Path = REPO_ROOT,
|
|
override: Optional[str] = None,
|
|
) -> str:
|
|
"""Return the git SHA to record in result metadata.
|
|
|
|
Resolution order matches ``experiments/run_three_way.py::_git_commit``:
|
|
explicit ``override``, ``SPECTRALQUANT_GIT_COMMIT`` env var, then
|
|
``git rev-parse HEAD``. Returns ``"0000000"`` if nothing usable is
|
|
available.
|
|
"""
|
|
candidates: List[str] = []
|
|
if override:
|
|
candidates.append(override)
|
|
env_val = os.environ.get(GIT_COMMIT_ENV, "")
|
|
if env_val:
|
|
candidates.append(env_val)
|
|
for sha in candidates:
|
|
sha = sha.strip()
|
|
if len(sha) >= 7 and _SHA_RE.fullmatch(sha):
|
|
return sha
|
|
try:
|
|
out = subprocess.check_output(
|
|
["git", "-C", str(repo_root), "rev-parse", "HEAD"],
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
sha = out.decode("utf-8").strip()
|
|
if len(sha) >= 7 and _SHA_RE.fullmatch(sha):
|
|
return sha
|
|
except Exception:
|
|
pass
|
|
return "0000000"
|
|
|
|
|
|
def safe_command_string(argv: Sequence[str]) -> str:
|
|
"""Reconstruct argv into a single-line command string.
|
|
|
|
No env values are inlined — only the argv is included so accidental
|
|
leakage of token values cannot happen here. Mirrors
|
|
``run_three_way._safe_command_string``.
|
|
"""
|
|
return " ".join(shlex.quote(a) for a in argv)
|
|
|
|
|
|
def model_short(model_name: str) -> str:
|
|
base = model_name.split("/")[-1]
|
|
return re.sub(r"[^A-Za-z0-9._-]+", "_", base).strip("._-") or "model"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Hardware / software metadata
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def build_software_block(extra: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
|
|
pyver = ".".join(str(x) for x in sys.version_info[:3])
|
|
|
|
def _try_version(mod: str) -> str:
|
|
try:
|
|
return __import__(mod).__version__ # type: ignore[attr-defined]
|
|
except Exception:
|
|
return "unavailable"
|
|
|
|
block: Dict[str, Any] = {
|
|
"python": pyver,
|
|
"torch": _try_version("torch"),
|
|
"transformers": _try_version("transformers"),
|
|
"datasets": _try_version("datasets"),
|
|
}
|
|
if extra:
|
|
block.update(extra)
|
|
return block
|
|
|
|
|
|
def build_hardware_block(device: str = "cpu") -> Dict[str, Any]:
|
|
block: Dict[str, Any] = {
|
|
"gpu": "none" if device == "cpu" else "unknown",
|
|
"cuda": "n/a",
|
|
"host": socket.gethostname(),
|
|
"platform": platform.platform(),
|
|
}
|
|
if device == "cuda":
|
|
try:
|
|
import torch # type: ignore[import-not-found]
|
|
|
|
if torch.cuda.is_available():
|
|
block["gpu"] = torch.cuda.get_device_name(0)
|
|
block["cuda"] = (
|
|
torch.version.cuda or "unknown" # type: ignore[attr-defined]
|
|
)
|
|
try:
|
|
props = torch.cuda.get_device_properties(0)
|
|
block["gpu_total_memory_mb"] = round(
|
|
props.total_memory / (1024 * 1024), 2
|
|
)
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
return block
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Atomic write + schema validation (shared with run_three_way pattern)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _validate_payload(payload: Dict[str, Any], schema_path: Path) -> None:
|
|
try:
|
|
import jsonschema # type: ignore[import-not-found]
|
|
from jsonschema import Draft202012Validator
|
|
except ImportError:
|
|
return
|
|
|
|
schema = json.loads(schema_path.read_text(encoding="utf-8"))
|
|
validator = Draft202012Validator(schema)
|
|
errors = sorted(validator.iter_errors(payload),
|
|
key=lambda e: list(e.absolute_path))
|
|
if errors:
|
|
msg = "\n".join(
|
|
f" - at {list(e.absolute_path) or '<root>'}: {e.message}"
|
|
for e in errors
|
|
)
|
|
raise ValueError(f"payload does not validate against {schema_path.name}:\n{msg}")
|
|
|
|
|
|
def atomic_write_json(
|
|
path: Path,
|
|
payload: Dict[str, Any],
|
|
schema_path: Optional[Path] = None,
|
|
) -> None:
|
|
"""Atomically write ``payload`` to ``path``.
|
|
|
|
1. Validate against schema (if jsonschema importable).
|
|
2. Write to a tempfile in same dir, fsync, ``os.replace``.
|
|
|
|
A half-written or schema-invalid file never appears at the canonical
|
|
path. Mirrors ``experiments/run_three_way.atomic_write_json``.
|
|
"""
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
if schema_path is not None and schema_path.exists():
|
|
_validate_payload(payload, schema_path)
|
|
fd, tmp_name = tempfile.mkstemp(prefix=".tmp_", dir=str(path.parent))
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
|
json.dump(payload, fh, indent=2, sort_keys=True)
|
|
fh.flush()
|
|
os.fsync(fh.fileno())
|
|
os.replace(tmp_name, path)
|
|
finally:
|
|
if os.path.exists(tmp_name):
|
|
try:
|
|
os.unlink(tmp_name)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Output path / run id
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def derive_run_id(
|
|
family: str,
|
|
model: str,
|
|
*,
|
|
suffix: str,
|
|
mode: str,
|
|
) -> str:
|
|
"""Build a deterministic ``run_id`` for an eval-family run.
|
|
|
|
``suffix`` is a family-specific tail (e.g. ``b3_seed42_wikitext-103``)
|
|
so two configs of the same family on the same model never collide.
|
|
Smoke / dry-run modes get a stable prefix so they never overwrite
|
|
full-mode artifacts.
|
|
"""
|
|
short = model_short(model)
|
|
rid = f"{family}__{short}__{suffix}"
|
|
if mode == MODE_SYNTHETIC_SMOKE:
|
|
rid = "synthetic_smoke__" + rid
|
|
elif mode == MODE_INLINE_CORPUS_SMOKE:
|
|
rid = "inline_corpus_smoke__" + rid
|
|
elif mode == MODE_DRY_RUN:
|
|
rid = "dryrun__" + rid
|
|
return rid
|
|
|
|
|
|
def output_path(output_dir: Path, run_id: str) -> Path:
|
|
return Path(output_dir) / f"{run_id}.json"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Method label discipline
|
|
# ---------------------------------------------------------------------------
|
|
|
|
#: Method keys used across all four eval families. Methods present in any
|
|
#: given run depend on what the harness can produce; we only require >= 1.
|
|
KNOWN_METHOD_KEYS: Tuple[str, ...] = (
|
|
"fp16",
|
|
"turboquant",
|
|
"spectralquant_v1",
|
|
"spectralquant_v2",
|
|
"official_turboquant",
|
|
)
|
|
|
|
|
|
def assert_method_keys(methods: Dict[str, Any]) -> None:
|
|
"""Light sanity check: every method key is recognized.
|
|
|
|
Unknown keys are a hint that the harness wrote something downstream
|
|
code won't expect; we surface this as a hard error to keep the
|
|
catalog discipline tight.
|
|
"""
|
|
unknown = sorted(set(methods.keys()) - set(KNOWN_METHOD_KEYS))
|
|
if unknown:
|
|
raise ValueError(
|
|
f"Unknown method keys {unknown}; permitted: {KNOWN_METHOD_KEYS}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Inline corpus (matches run_three_way's INLINE_SMOKE_CORPUS spirit)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
INLINE_CORPUS: Tuple[str, ...] = (
|
|
"Quantization research aims to reduce the memory footprint of large "
|
|
"language models without significantly degrading downstream task quality. "
|
|
"Common approaches include weight quantization, key-value cache "
|
|
"compression, and post-training calibration on a small held-out corpus.",
|
|
"Eigenspectral calibration estimates a per-head rotation that aligns the "
|
|
"key-value subspace with the dominant eigendirections of the empirical "
|
|
"covariance matrix observed during a forward pass over short text "
|
|
"samples drawn from the calibration corpus.",
|
|
"Water-filling allocates a non-uniform per-dimension bit budget so that "
|
|
"dimensions carrying more signal receive more bits. This typically "
|
|
"improves attention-output cosine similarity at low average bitrates "
|
|
"compared to uniform semantic quantizers.",
|
|
"The three-way comparison records cosine similarity between FP16 "
|
|
"reference attention weights and the weights produced by TurboQuant, "
|
|
"SpectralQuant v1, and SpectralQuant v2 across a sampled set of layers.",
|
|
"Rotary position embeddings interleave sine and cosine components into "
|
|
"the query and key projections; pre-RoPE quantization keeps the rotation "
|
|
"math identical between the reference path and the quantized path.",
|
|
"Calibration tokens are clipped to a configurable cap so that very long "
|
|
"examples do not dominate the eigenspectrum estimate. Default caps in "
|
|
"the spec sit between 256 and 512 tokens per layer per sample.",
|
|
"Group-query attention factors the key-value heads through a smaller "
|
|
"set of projections shared across query-head groups, which reduces the "
|
|
"memory bandwidth required to serve long context windows in production.",
|
|
"Synthetic smoke runs verify the result-plumbing end to end on small "
|
|
"Q/K/V tensors. Inline-corpus smoke runs additionally exercise the full "
|
|
"HuggingFace model load and hook path without downloading any dataset.",
|
|
)
|
|
|
|
|
|
def build_inline_corpus(n_total: int) -> List[str]:
|
|
"""Return ``n_total`` deterministic strings; cycles INLINE_CORPUS."""
|
|
if n_total < 1:
|
|
raise ValueError(f"n_total must be >= 1, got {n_total}")
|
|
base = list(INLINE_CORPUS)
|
|
out: List[str] = []
|
|
for i in range(n_total):
|
|
text = base[i % len(base)]
|
|
out.append(f"{text} [inline_smoke#{i:03d}]")
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Deterministic prompt set for run_generation.py synthetic / inline modes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
DEFAULT_GENERATION_PROMPTS: Tuple[Dict[str, str], ...] = (
|
|
{
|
|
"id": "summary-001",
|
|
"category": "summarization",
|
|
"prompt": (
|
|
"Summarize the following passage in two sentences. "
|
|
"Passage: SpectralQuant exploits a structural property of the "
|
|
"key-value cache in transformer attention to compress without "
|
|
"destroying downstream attention-output quality."
|
|
),
|
|
},
|
|
{
|
|
"id": "qa-001",
|
|
"category": "qa",
|
|
"prompt": (
|
|
"Question: What is the main motivation for non-uniform per-dimension "
|
|
"bit allocation in KV cache compression? Answer briefly."
|
|
),
|
|
},
|
|
{
|
|
"id": "code-001",
|
|
"category": "code",
|
|
"prompt": (
|
|
"Write a short Python function that computes the participation "
|
|
"ratio of a vector of eigenvalues. The function should return "
|
|
"(sum(eig)**2) / sum(eig**2)."
|
|
),
|
|
},
|
|
{
|
|
"id": "complete-001",
|
|
"category": "completion",
|
|
"prompt": (
|
|
"The participation ratio of a spectrum is defined as"
|
|
),
|
|
},
|
|
{
|
|
"id": "instruct-001",
|
|
"category": "instruction",
|
|
"prompt": (
|
|
"List three reasons why low-rank approximation typically works "
|
|
"for keys but not for values in transformer attention."
|
|
),
|
|
},
|
|
{
|
|
"id": "translate-001",
|
|
"category": "translation",
|
|
"prompt": (
|
|
"Translate to French: 'The calibration step takes a few seconds "
|
|
"and is performed once per model.'"
|
|
),
|
|
},
|
|
{
|
|
"id": "longctx-001",
|
|
"category": "long_context",
|
|
"prompt": (
|
|
"Given the context that water-filling assigns more bits to "
|
|
"dimensions with larger eigenvalues, explain in one paragraph "
|
|
"why this improves attention quality at low average bitrates."
|
|
),
|
|
},
|
|
{
|
|
"id": "factual-001",
|
|
"category": "factual",
|
|
"prompt": (
|
|
"What does the acronym KV cache stand for in the context of "
|
|
"transformer inference?"
|
|
),
|
|
},
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Token-name redaction (mirrors launcher / run_status sanitize_text)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
SECRET_NAME_PATTERNS: Tuple[str, ...] = (
|
|
"TOKEN", "SECRET", "API_KEY", "APIKEY",
|
|
"PASSWORD", "PASSWD", "PRIVATE_KEY",
|
|
)
|
|
|
|
|
|
_TOKEN_LITERAL_PATTERNS = (
|
|
re.compile(r"\bhf_[A-Za-z0-9]{20,}\b"),
|
|
re.compile(r"\bghp_[A-Za-z0-9]{20,}\b"),
|
|
re.compile(r"\bgithub_pat_[A-Za-z0-9_]{20,}\b"),
|
|
re.compile(r"\bAKIA[0-9A-Z]{16}\b"),
|
|
re.compile(r"\bak-[A-Za-z0-9_-]{16,}\b"),
|
|
re.compile(r"\bas-[A-Za-z0-9_-]{16,}\b"),
|
|
)
|
|
|
|
|
|
def sanitize_text(text: str, *, extra_env: Optional[Iterable[str]] = None) -> str:
|
|
if not text:
|
|
return text
|
|
redacted = text
|
|
seen: set = set()
|
|
env_names: List[str] = list(extra_env or ())
|
|
for k in os.environ:
|
|
upper = k.upper()
|
|
if any(p in upper for p in SECRET_NAME_PATTERNS):
|
|
env_names.append(k)
|
|
for name in env_names:
|
|
val = os.environ.get(name, "")
|
|
if not val or len(val) < 8 or val in seen:
|
|
continue
|
|
seen.add(val)
|
|
redacted = redacted.replace(val, f"[REDACTED:{name}]")
|
|
for pat in _TOKEN_LITERAL_PATTERNS:
|
|
redacted = pat.sub("[REDACTED:token-literal]", redacted)
|
|
return redacted
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Common payload skeleton
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def base_payload(
|
|
*,
|
|
family: str,
|
|
run_id: str,
|
|
argv: Sequence[str],
|
|
model_name: str,
|
|
mode: str,
|
|
paper_valid: bool,
|
|
device: str = "cpu",
|
|
extra_software: Optional[Dict[str, str]] = None,
|
|
git_commit_override: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Build the common top-level fields for an eval-family payload.
|
|
|
|
Family-specific fields (``data``, ``decoding``, ``operating_points``,
|
|
``timing``, ``methods``, ``evidence_ids``, ``caveats``) are added by
|
|
the caller.
|
|
"""
|
|
return {
|
|
"schema_version": SCHEMA_VERSION,
|
|
"family": family,
|
|
"run_id": run_id,
|
|
"timestamp": utc_now_iso(),
|
|
"repo": REPO_SLUG,
|
|
"commit": git_commit(override=git_commit_override),
|
|
"command": safe_command_string(argv),
|
|
"mode": mode,
|
|
"paper_valid": bool(paper_valid),
|
|
"model": {"name": model_name},
|
|
"hardware": build_hardware_block(device),
|
|
"software": build_software_block(extra_software),
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def add_common_eval_args(parser: Any) -> None:
|
|
"""Add the ``--dry-run / --synthetic-smoke / --inline-corpus-smoke``
|
|
triad and friends to an argparse parser. Mirrors ``run_three_way``.
|
|
"""
|
|
parser.add_argument("--model", required=True, type=str,
|
|
help="HuggingFace model id, e.g. mistralai/Mistral-7B-v0.3")
|
|
parser.add_argument("--output-dir", required=True, type=str,
|
|
dest="output_dir")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--device", type=str, default="cpu",
|
|
choices=("cpu", "cuda"))
|
|
parser.add_argument("--dtype", type=str, default="float32",
|
|
choices=("float32", "float16", "bfloat16"))
|
|
parser.add_argument("--dry-run", action="store_true", dest="dry_run")
|
|
parser.add_argument("--synthetic-smoke", action="store_true",
|
|
dest="synthetic_smoke")
|
|
parser.add_argument("--inline-corpus-smoke", action="store_true",
|
|
dest="inline_corpus_smoke")
|
|
parser.add_argument("--force", action="store_true", dest="force")
|
|
parser.add_argument("--skip-if-exists", action="store_true", default=True,
|
|
dest="skip_if_exists")
|
|
parser.add_argument("--no-skip-if-exists", action="store_false",
|
|
dest="skip_if_exists")
|
|
parser.add_argument("--git-commit", type=str, default=None,
|
|
dest="git_commit_override")
|
|
parser.add_argument("--status-dir", type=str, default=None,
|
|
dest="status_dir")
|
|
|
|
|
|
def resolve_mode(ns: Any) -> str:
|
|
"""Translate the parsed namespace's flag triad into a single mode."""
|
|
if getattr(ns, "dry_run", False):
|
|
return MODE_DRY_RUN
|
|
if getattr(ns, "synthetic_smoke", False) and getattr(ns, "inline_corpus_smoke", False):
|
|
raise SystemExit(
|
|
"--synthetic-smoke and --inline-corpus-smoke are mutually exclusive"
|
|
)
|
|
if getattr(ns, "synthetic_smoke", False):
|
|
return MODE_SYNTHETIC_SMOKE
|
|
if getattr(ns, "inline_corpus_smoke", False):
|
|
return MODE_INLINE_CORPUS_SMOKE
|
|
return MODE_FULL
|
|
|
|
|
|
def install_repo_paths_into_sys_path() -> None:
|
|
"""Make ``import spectralquant`` and ``from experiments import ...``
|
|
work without ``pip install -e .``. Mirrors run_three_way's bootstrap.
|
|
"""
|
|
if str(SRC_DIR) not in sys.path:
|
|
sys.path.insert(0, str(SRC_DIR))
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|