Upload 14 files

gguf.py

@@ -377,6 +377,7 @@ def write_quantized_gguf(
     limit: int = 0,
     verbose: bool = True,
     progress_callback=None,
+    stage_callback=None,
     stop_event=None,
 ) -> dict[str, float]:
     """Quantize a source GGUF with the V2 engine and emit a valid GGUF.
@@ -469,13 +470,36 @@ def write_quantized_gguf(
             if target > fout.tell():
                 fout.write(b"\x00" * (target - fout.tell()))
             if plan[idx]:
+                has_imat = imat is not None and lookup_imatrix(imat, ti.name) is not None
+                if stage_callback is not None:
+                    stage_callback(
+                        idx + 1,
+                        len(out_infos),
+                        ti,
+                        {
+                            "phase": f"Reading & Dequantizing Source ({ti.type_name})",
+                            "done": 0,
+                            "total": ti.n_elements // QK_K,
+                            "detail": f"dims={ti.dims} elements={ti.n_elements:,} ({ti.n_bytes / 1048576:.2f} MB) imatrix={'yes' if has_imat else 'no'}",
+                        },
+                    )
                 x, row = reader.read_tensor(ti.name)
                 imp = None
                 if imat is not None:
                     raw_imp = lookup_imatrix(imat, ti.name)
                     if raw_imp is not None:
                         imp = expand_importance(raw_imp, x.size, row)
-                res = quantizer.quantize(x, importance=imp, row_width=row)
+                res = quantizer.quantize(
+                    x,
+                    importance=imp,
+                    row_width=row,
+                    status_callback=(
+                        (lambda ev, _i=idx + 1, _ti=ti: stage_callback(_i, len(out_infos), _ti, ev))
+                        if stage_callback is not None
+                        else None
+                    ),
+                    stop_event=stop_event,
+                )
                 payload = blocks_to_bytes(res.scales, res.qs, res.dm)
                 if len(payload) != out_size:
                     raise RuntimeError(f"{ti.name}: Q2_K payload size mismatch")
@@ -498,6 +522,18 @@ def write_quantized_gguf(
                         imp is not None,
                     )
             else:
+                if stage_callback is not None:
+                    stage_callback(
+                        idx + 1,
+                        len(out_infos),
+                        ti,
+                        {
+                            "phase": f"Passthrough Copy ({ti.type_name})",
+                            "done": 1,
+                            "total": 1,
+                            "detail": f"dims={ti.dims} ({ti.n_bytes / 1048576:.2f} MB)",
+                        },
+                    )
                 reader.copy_tensor_raw_to(ti.name, fout)
                 if verbose and idx % 25 == 0:
                     print(f"  copy {ti.name}")

pipeline.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from math import ceil
+import time
 
 import numpy as np
 
@@ -413,7 +414,25 @@ class Q2KQuantizer:
         weights: np.ndarray,
         importance: np.ndarray | None = None,
         row_width: int | None = None,
+        status_callback=None,
+        stop_event=None,
     ) -> QuantizedTensor:
+        def _emit(phase: str, done: int, total: int, detail: str, **extra) -> None:
+            if status_callback is not None:
+                status_callback(
+                    {
+                        "phase": phase,
+                        "done": done,
+                        "total": total,
+                        "detail": detail,
+                        **extra,
+                    }
+                )
+
+        def _check_stop() -> None:
+            if stop_event is not None and stop_event.is_set():
+                raise InterruptedError("Quantization cancelled by user")
+
         cfg = self.config
         x = np.asarray(weights, dtype=np.float32)
         if x.size == 0:
@@ -436,13 +455,34 @@ class Q2KQuantizer:
         else:
             W = None
 
+        sigma = float(np.sqrt(np.mean(x.astype(np.float64) ** 2)))
+        _check_stop()
         clamp_val = seed_clamp_value(x, cfg)
+        _emit(
+            "Phase 1: Kurtosis Clamp & Batched GGML Seeds",
+            0,
+            n_blocks,
+            f"blocks={n_blocks} sigma={sigma:.5f} clamp={clamp_val:.5f} imatrix={'yes' if W is not None else 'no'}",
+            sigma=sigma,
+        )
         seeds = build_seeds(X, W, cfg, clamp_val)
+        _check_stop()
+        _emit(
+            "Phase 2a: Vectorized Superblock WLS",
+            0,
+            n_blocks,
+            f"Solving 5-iteration WLS across {n_blocks} superblocks",
+            sigma=sigma,
+        )
         base_d, base_m = _phase2_wls(X, W, seeds, clamp_val)
 
         candidate_errors = np.full((n_blocks, TOTAL_CANDIDATES), np.inf, dtype=np.float32)
         standalone: list[FactorResult] = []
+        t_cand = time.time()
+        last_emit = 0.0
+        step_report = max(1, min(32, n_blocks // 20 or 1))
         for i in range(n_blocks):
+            _check_stop()
             wi = None if W is None else W[i]
             candidate_errors[i] = factor_candidates(
                 X[i],
@@ -453,13 +493,48 @@ class Q2KQuantizer:
                 base_m[i],
                 cfg,
             )
-            standalone.append(factor_only_solve(X[i], wi, cfg))
+            fac_res = factor_only_solve(X[i], wi, cfg)
+            standalone.append(fac_res)
+            now = time.time()
+            if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
+                last_emit = now
+                elapsed = max(now - t_cand, 1.0e-6)
+                rate = (i + 1) / elapsed
+                eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
+                _emit(
+                    "Phase 2b: Analog-Guided Factored Candidate Field",
+                    i + 1,
+                    n_blocks,
+                    f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
+                    f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} "
+                    f"obj={fac_res.objective:.5e} starts={fac_res.starts}",
+                    rate=rate,
+                    eta=eta,
+                    sigma=sigma,
+                )
 
+        _check_stop()
+        _emit(
+            "Phase 3: 6-State Sieve & 36-State Viterbi DP",
+            0,
+            n_blocks,
+            f"Running Sieve + 36-state Viterbi sequence optimization ({cfg.optimizer.upper()})",
+            sigma=sigma,
+        )
         best, temperature, _ = _phase3_select(candidate_errors, cfg)
         carry = _compute_dc_carry(X, seeds, best, base_d, base_m, cfg, row_width)
+        _emit(
+            "Phase 3.9: Viterbi Path & Rolling DC Carry",
+            n_blocks,
+            n_blocks,
+            f"Boltzmann temp={temperature:.4e}, max |DC carry|={float(np.max(np.abs(carry))):.4e}",
+            temperature=temperature,
+            sigma=sigma,
+        )
 
         at_slot: list[FactorResult] = []
         for i in range(n_blocks):
+            _check_stop()
             wi = None if W is None else W[i]
             d16, m16 = candidate_pair(base_d[i], base_m[i], best[i])
             at_slot.append(
@@ -482,7 +557,11 @@ class Q2KQuantizer:
         blocks: list[Q2KBlock] = []
         factor_floor_used = 0
         candidate_floor_used = 0
+        running_sse = 0.0
+        t_asm = time.time()
+        last_emit = 0.0
         for i in range(n_blocks):
+            _check_stop()
             wi = None if W is None else W[i]
             d, m, ls, lm, codes, obj, sse = _assemble_block(
                 X[i], wi, seeds, base_d[i], base_m[i], best[i], carry[i], cfg
@@ -517,9 +596,42 @@ class Q2KQuantizer:
                         obj, sse = fobj, fsse
                         factor_floor_used += 1
 
-            blocks.append(make_block(d, m, ls, lm, codes))
+            blk_obj = make_block(d, m, ls, lm, codes)
+            blocks.append(blk_obj)
+            running_sse += blk_obj.error(X[i])
+            now = time.time()
+            if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
+                last_emit = now
+                elapsed = max(now - t_asm, 1.0e-6)
+                rate = (i + 1) / elapsed
+                eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
+                cur_rmse = float(np.sqrt(running_sse / ((i + 1) * QK_K)))
+                _emit(
+                    "Phase 4: Sub-block Sieve + ULP + Hadamard Assembly",
+                    i + 1,
+                    n_blocks,
+                    f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
+                    f"rmse={cur_rmse:.6e} rmse/sigma={cur_rmse/max(sigma, 1e-15):.4f} "
+                    f"floors(cand={candidate_floor_used}, fac={factor_floor_used})",
+                    rmse=cur_rmse,
+                    sigma=sigma,
+                    temperature=temperature,
+                    cand_floors=candidate_floor_used,
+                    fac_floors=factor_floor_used,
+                    rate=rate,
+                    eta=eta,
+                )
 
         if cfg.apply_dc_carry and (cfg.dc_lambda > 0.0 or cfg.dc_decay > 0.0):
+            _check_stop()
+            _emit(
+                "Phase 4.8: Sequential Row-Aware DC Residual Nudge",
+                n_blocks,
+                n_blocks,
+                f"Cancelling cumulative row DC residuals (row_width={row_width or QK_K})",
+                sigma=sigma,
+                temperature=temperature,
+            )
             blocks = _apply_dc_nudge(X, W, blocks, cfg, row_width)
 
         total_sse = 0.0
@@ -532,7 +644,6 @@ class Q2KQuantizer:
             else:
                 weighted += float(np.dot(W[i].astype(np.float64), e * e))
 
-        sigma = float(np.sqrt(np.mean(x.astype(np.float64) ** 2)))
         rmse = float(np.sqrt(total_sse / x.size))
         stats = TensorStats(
             n_blocks=n_blocks,
@@ -545,6 +656,19 @@ class Q2KQuantizer:
             factor_floor_used=factor_floor_used,
             candidate_floor_used=candidate_floor_used,
         )
+        _emit(
+            "Phase 5: Tensor Complete",
+            n_blocks,
+            n_blocks,
+            f"Done: blocks={n_blocks} rmse={rmse:.6e} rmse/sigma={stats.rmse_over_sigma:.5f} "
+            f"wSSE={weighted:.5e} temp={temperature:.3e} "
+            f"floors(cand={candidate_floor_used}, fac={factor_floor_used})",
+            rmse=rmse,
+            sigma=sigma,
+            temperature=temperature,
+            cand_floors=candidate_floor_used,
+            fac_floors=factor_floor_used,
+        )
         scales, qs, dm = _stack_blocks(blocks)
         factor_objectives = np.array([r.objective for r in standalone], dtype=np.float64)
         result = QuantizedTensor(