HPC-Quantize-V2 / config.py
CompressedGemma's picture
Revert to keep stock llama compatibility
dfb9420 verified
Raw History Blame Contribute Delete
7.91 kB
"""Configuration for the HPC-Quantize V2 Q2_K engine.
Every tunable that was a compile-time macro, environment variable or hidden
global in the C engine is collected here as an explicit, validated field.
The defaults are the C engine's production defaults unless marked
``V2`` in the docstring, in which case they are deliberate improvements.
"""
from __future__ import annotations
from dataclasses import dataclass, field
@dataclass(slots=True)
class HPCConfig:
# ---- block geometry -------------------------------------------------
qk_k: int = 256
# ---- candidate field ------------------------------------------------
# V2 improvement: the default is the 6x6 symbolic lattice (36 exact
# factored solves) instead of the 24x24 physical grid (576 solves).
# Set to 24/24 to reproduce the full C candidate field.
n_cand_d: int = 6
n_cand_m: int = 6
factor_pin: int = 3
factor_iters: int = 1
factor_1d_iters: int = 6
factor_1d_refine_sweeps: int = 1
factor_standalone_iters: int = 8
factor_starts: int = 2
factor_ulp_radius: int = 1
factor_refine_1d: bool = True
factor_portfolio: bool = True
factor_spectral_refine: bool = True
factor_exact_emit: bool = True
use_weight_analogs: bool = True
factor_coupled_q2k: bool = True
factor_joint_subscales: bool = True
# ---- objective ------------------------------------------------------
use_spectral: bool = True
dc_lambda: float = 1.0
vw_lambda: float = 1.0
fold_depth: int = 1
fold_gamma: float = 2.0
dc_decay: float = 0.85
carry_cumulative: bool = True
sse_budget: float = 5.0e-4
greedy_override_ratio: float = 0.995
# ---- sieve ----------------------------------------------------------
same_penalty: float = 0.85
log_slack: float = 1.5
quhits_per_block: int = 2
graph_max_blocks: int = 8192
# ---- Viterbi --------------------------------------------------------
optimizer: str = "hpc"
viterbi_beta: float = 0.25
viterbi_alpha: float = 0.08
temperature_sample_blocks: int = 128
# ---- assembly -------------------------------------------------------
refine_passes: int = 3
ulp_radius: int = 4
subblock_states: int = 6
subblock_sieve: bool = True
use_factor_floor: bool = True
use_candidate_floor: bool = True
apply_dc_carry: bool = True
quant_format: str = "mixed"
custom_quant_map: dict[str, str] = field(default_factory=dict)
enable_q4_fallback: bool = True
quantize_embd_q4: bool = True
promote_attn_q4: bool = False
# ---- thermal refinement (pair heat-bath annealer, opt-in) ------------
# V2-only knob (0 = off, bit-identical to the deterministic pipeline).
# Runs a seeded checkerboard Gibbs sampler over the vesica-pair lattice
# after Phase 4; late passes are monotone so output never regresses.
therm_passes: int = 0
therm_sweeps: int = 2
therm_t0: float = 0.25
therm_t_decay: float = 0.35
therm_seed: int = 271828
# Zero-residual anneal: multiplier >= 1 applied to the spectral lambdas
# (dc_lambda, vw_lambda) while annealing. At large tilts (~1e2-1e3) the
# exact-conditional Gibbs chain spends weighted SSE to drain the *coherent*
# residual (row DC + vesica pairs) toward the lattice floor, leaving an
# (almost) spectrally white error. 1.0 = off.
therm_lambda_tilt: float = 1.0
# Zero-DC pin (the "Q2_K free lunch"): after assembly (and any thermal
# refinement) drive each block's residual DC toward 0. Q2_K re-solves
# (d, m) on the exact dc=0 constraint line — a joint least-squares
# translation whose closed form *lowers* SSE (free). Q4_0 / Q3_K carry
# a single scale dof, so the pin there is a deterministic code-flip
# exchange: it spends at most ``pin_budget`` of each block's SSE to buy
# the coherent-residual drop (Q3_K lands at ~-85% DC RMS for ~0.9% SSE;
# Q4_0 is bounded by the narrower [0, 15] grid's DC slack).
pin_dc: bool = False
# Fraction of per-block SSE the Q4_0 / Q3_K code-flip pin may spend
# (the Q2_K closed-form pin never spends any). 0 disables the expense.
pin_budget: float = 0.05
# ---- seed clamp -----------------------------------------------------
clamp_kurtosis_off: float = 6.0
clamp_kurtosis_soft: float = 20.0
clamp_sigma_soft: float = 6.0
clamp_sigma_hard: float = 4.0
temper_imatrix: bool = True
# ---- reporting & parallelism ----------------------------------------
device: str = "auto"
workers: int = 0
verbose: bool = False
def validate(self) -> "HPCConfig":
if self.qk_k != 256:
raise ValueError("Q2_K requires qk_k == 256")
if not (1 <= self.n_cand_d <= 24 and 1 <= self.n_cand_m <= 24):
raise ValueError("candidate counts must lie in [1, 24]")
if self.optimizer not in ("hpc", "mse"):
raise ValueError("optimizer must be 'hpc' or 'mse'")
if not 0 <= self.factor_pin <= 3:
raise ValueError("factor_pin must be a 2-bit mask")
if self.fold_depth < 0:
raise ValueError("fold_depth must be non-negative")
if not 0.0 <= self.dc_decay <= 1.0:
raise ValueError("dc_decay must lie in [0, 1]")
if self.ulp_radius < 0 or self.factor_ulp_radius < 0:
raise ValueError("ULP radii must be non-negative")
if self.refine_passes < 1:
raise ValueError("refine_passes must be >= 1")
if not 1 <= self.subblock_states <= 6:
raise ValueError("subblock_states must lie in [1, 6]")
if self.factor_1d_iters < 1 or self.factor_standalone_iters < 1:
raise ValueError("factor iteration budgets must be >= 1")
if self.factor_starts < 1:
raise ValueError("factor_starts must be >= 1")
if self.sse_budget < 0.0:
raise ValueError("sse_budget must be non-negative")
if self.pin_budget < 0.0:
raise ValueError("pin_budget must be non-negative")
if not 0.0 < self.same_penalty <= 1.0:
raise ValueError("same_penalty must lie in (0, 1]")
if self.log_slack < 0.0:
raise ValueError("log_slack must be non-negative")
if self.viterbi_alpha < 0.0 or self.viterbi_beta < 0.0:
raise ValueError("Viterbi weights must be non-negative")
if self.workers < 0:
raise ValueError("workers must be >= 0 (0 = auto CPU count)")
if self.therm_passes < 0:
raise ValueError("therm_passes must be >= 0 (0 disables thermal refinement)")
if self.therm_sweeps < 1:
raise ValueError("therm_sweeps must be >= 1")
if self.therm_t0 <= 0.0:
raise ValueError("therm_t0 must be positive")
if not 0.0 < self.therm_t_decay < 1.0:
raise ValueError("therm_t_decay must lie in (0, 1)")
if self.therm_lambda_tilt < 1.0:
raise ValueError("therm_lambda_tilt must be >= 1.0 (1.0 = off)")
if self.device not in ("auto", "cuda", "cpu") and not self.device.startswith("cuda:"):
raise ValueError("device must be 'auto', 'cuda', 'cuda:N', or 'cpu'")
if self.quant_format not in ("mixed", "mixed_q1", "custom", "q1_0", "q3_k", "q4_0", "q2_k"):
raise ValueError(
"quant_format must be 'mixed', 'mixed_q1', 'custom', 'q1_0', 'q3_k', 'q4_0', or 'q2_k'"
)
allowed_custom = {"q1_0", "q2_k", "q3_k", "q4_0", "copy", "auto"}
for tensor_type, quant in self.custom_quant_map.items():
if str(quant).strip().lower() not in allowed_custom:
raise ValueError(
f"custom_quant_map[{tensor_type!r}] must be one of "
f"Q1_0, Q2_K, Q3_K, Q4_0, Copy, Auto, got {quant!r}"
)
return self