initial commit

nvfp4gguf/quant.py

@@ -0,0 +1,1000 @@
+"""NVFP4 learned rounding and ggml packing.
+
+The optimizer is convert_to_quant's LearnedNVFP4Converter. This module only
+captures the scales that quantizer used, packs them into llama.cpp's NVFP4
+block layout, and keeps the same hyperparameters convert_to_quant uses.
+Matrix multiplies stay in full FP32. TF32 would make the SVD residual drift
+from a convert_to_quant run on the same tensor.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import logging
+import os
+import re
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+log = logging.getLogger("nvfp4gguf")
+
+QUANT_VERSION = 2
+QK_NVFP4 = 64
+NVFP4_BLOCK_BYTES = 36
+F8_E4M3_MAX = 448.0
+CTQ_NUM_ITER = 4000
+# Joint trains the block scales inside the same loop. Iterative refreshes them
+# only on the original optimizer, so joint is the default that Prodigy, AdamW,
+# and RAdam all honor.
+DEFAULT_SCALE_OPTIMIZATION = "joint"
+DEFAULT_SCALE_REFINEMENT_ROUNDS = 2
+# Patience 1 decays the step on every miss and the lr-floor stop then ends
+# the search while the loss is still the best. Prodigy adapts its own step,
+# so a slower plateau and no floor abort let the budget actually run.
+DEFAULT_LR_PATIENCE = 32
+DEFAULT_EARLY_STOP_LR = 0.0
+# A wobble of 2.000e-7 to 2.002e-7 is not progress, but only once the step
+# size has floored. Before that, any new best has to count or the search
+# stops while the loss is still around 1e-2.
+DEFAULT_LR_THRESHOLD = 0.01
+DEFAULT_EARLY_STOP_STALL = 2000
+# Once the step size is at --lr-min, this many steps without a new best ends
+# the matrix. 64 is about a second at the speed these fits run.
+DEFAULT_EARLY_STOP_FLOOR_STALL = 64
+# ffn_down at 9e-4 already matched a run that reached 3e-7. 1e-5 is below
+# that, so hitting it does not change the packed weights.
+DEFAULT_TARGET_LOSS = 1e-5
+# These projections write the residual or build Q, K, and V. The low-rank
+# objective ignores most of their error, which is where a wrong token comes from.
+_TIGHT_ATTENTION = (
+    "attn_output.weight",
+    "attn_q.weight",
+    "attn_k.weight",
+    "attn_v.weight",
+    "attn_qkv.weight",
+    "o_proj.weight",
+    "q_proj.weight",
+    "k_proj.weight",
+    "v_proj.weight",
+    "q_k_v_proj.weight",
+)
+
+DEFAULT_LLAMA_CPP = Path.home() / "llama.cpp"
+DEFAULT_CONVERT_TO_QUANT = Path("/mnt/d/dev/convert_to_quant")
+
+
+def _ensure_import_path(path: Path, label: str) -> None:
+    path = path.resolve()
+    if not path.exists():
+        raise FileNotFoundError(f"{label} path does not exist: {path}")
+    entry = str(path)
+    if entry not in sys.path:
+        sys.path.insert(0, entry)
+
+
+def load_gguf_modules(llama_cpp: Path):
+    llama_cpp = llama_cpp.resolve()
+    if not (llama_cpp / "convert_hf_to_gguf.py").is_file():
+        raise FileNotFoundError(
+            f"{llama_cpp} does not look like a llama.cpp checkout (missing convert_hf_to_gguf.py)"
+        )
+    _ensure_import_path(llama_cpp, "llama.cpp")
+    _ensure_import_path(llama_cpp / "gguf-py", "llama.cpp gguf-py")
+    os.environ.setdefault("NO_LOCAL_GGUF", "1")
+    import gguf
+    from conversion import ModelBase, ModelType, get_model_architecture, get_model_class
+
+    return gguf, ModelBase, ModelType, get_model_architecture, get_model_class
+
+
+def import_gguf(llama_cpp: Path | None = None):
+    root = Path(llama_cpp) if llama_cpp else DEFAULT_LLAMA_CPP
+    gguf_path = root / "gguf-py"
+    if not gguf_path.is_dir():
+        raise FileNotFoundError(f"llama.cpp gguf-py was not found at {gguf_path}")
+    entry = str(gguf_path)
+    if entry not in sys.path:
+        sys.path.insert(0, entry)
+    import gguf
+    return gguf
+
+
+def load_learned_converter(convert_to_quant_dir: Path | None):
+    try:
+        from convert_to_quant.converters.learned_nvfp4 import LearnedNVFP4Converter
+        from convert_to_quant.utils.float_utils import unpack_uint4
+        return LearnedNVFP4Converter, unpack_uint4
+    except ImportError:
+        pass
+
+    search = []
+    if convert_to_quant_dir is not None:
+        search.append(Path(convert_to_quant_dir))
+    search.append(DEFAULT_CONVERT_TO_QUANT)
+    last_error = None
+    for candidate in search:
+        if not candidate.exists():
+            continue
+        _ensure_import_path(candidate, "convert_to_quant")
+        try:
+            from convert_to_quant.converters.learned_nvfp4 import LearnedNVFP4Converter
+            from convert_to_quant.utils.float_utils import unpack_uint4
+            return LearnedNVFP4Converter, unpack_uint4
+        except ImportError as exc:
+            last_error = exc
+    raise ImportError(
+        "convert_to_quant is not importable. Install it, or pass --convert-to-quant "
+        "pointing at the checkout that provides LearnedNVFP4Converter."
+    ) from last_error
+
+
+def ctq_skips_key(key: str, convert_to_quant_dir: Path | None = None) -> bool:
+    """True when convert_to_quant's default exclusion list would leave this tensor alone."""
+    try:
+        from convert_to_quant.constants import AVOID_KEY_NAMES
+    except ImportError:
+        load_learned_converter(convert_to_quant_dir)
+        from convert_to_quant.constants import AVOID_KEY_NAMES
+    return any(pattern in key for pattern in AVOID_KEY_NAMES)
+
+
+def resolve_optimizer(requested: str) -> str:
+    if requested != "auto":
+        return requested
+    try:
+        import prodigyplus.prodigy_plus_schedulefree  # noqa: F401
+        return "prodigy"
+    except ImportError:
+        log.info(
+            "prodigy-plus-schedule-free is not installed; using AdamW. "
+            "pip install prodigy-plus-schedule-free to match convert_to_quant's default optimizer."
+        )
+        return "adamw"
+
+
+def _opt(args, name: str, default):
+    return getattr(args, name, default)
+
+
+def resolved_lr_min(args) -> float:
+    """Plateau floor from ``--lr-min``. ``--lr`` is only the starting rate."""
+    lr_min = getattr(args, "lr_min", None)
+    if lr_min is None:
+        return 1e-8
+    return float(lr_min)
+
+
+class GgufLearnedNVFP4:
+    """LearnedNVFP4Converter plus the total scale the codes were optimized against."""
+
+    def __init__(self, converter_cls, **kwargs):
+        class _Capturing(converter_cls):
+            def convert(self, *args, **kw):
+                self._captured_total_scale = None
+                self._last_lr = self.lr
+                self._floor_stall = 0
+                self._best_seen = None
+                self._settled = False
+                return super().convert(*args, **kw)
+
+            def _learning_rate_floored(self) -> bool:
+                # lr_min is --lr-min, the plateau floor. --lr is only where the search starts.
+                return getattr(self, "_last_lr", self.lr) <= self.lr_min * 1.01
+
+            def _loss_has_settled(self, loss: float) -> bool:
+                target = float(getattr(self, "target_loss", DEFAULT_TARGET_LOSS) or 0)
+                return target > 0 and 0 < loss <= target
+
+            def _check_improvement(self, current_loss: float, best_loss: float) -> bool:
+                # Sub-1% wiggles stop counting once the fit has settled or the
+                # step has reached --lr-min. Before that they are the descent.
+                settled = getattr(self, "_settled", False)
+                if self.lr_threshold > 0 and (settled or self._learning_rate_floored()):
+                    return super()._check_improvement(current_loss, best_loss)
+                return current_loss < best_loss
+
+            def _check_early_stop(self, current_loss: float, curr_lr: float, worse_loss_counter: int) -> bool:
+                self._last_lr = curr_lr
+                if 0 < current_loss < float("inf"):
+                    best = getattr(self, "_best_seen", None)
+                    self._best_seen = current_loss if best is None else min(best, current_loss)
+                best_seen = getattr(self, "_best_seen", None)
+                self._settled = self._loss_has_settled(best_seen if best_seen is not None else current_loss)
+                limit = int(getattr(self, "early_stop_floor_stall", 0) or 0)
+                # The packed codes are already chosen once the best loss is at
+                # the target. The learning rate does not change that.
+                if self._settled:
+                    log.info("    Loss reached %g. Stopping early.", float(self.target_loss))
+                    return True
+                return super()._check_early_stop(current_loss, curr_lr, worse_loss_counter)
+
+            def _compute_svd_components(self, W_float32, verbose: bool = True):
+                # full_matrices=True builds the square factors and then throws
+                # them away. The loss only uses k directions, and a thin SVD
+                # returns those directly.
+                if not self.full_matrix:
+                    return super()._compute_svd_components(W_float32, verbose)
+                import math
+                rows, cols = W_float32.shape
+                max_rank = min(rows, cols)
+                k = min(self.max_k, max(self.min_k, int(math.floor(self.top_p * max_rank))))
+                k = min(k, max_rank)
+                if verbose:
+                    print(f"    - Tensor shape: [{rows}, {cols}], Max rank: {max_rank}. Using k={k} components.")
+                    print("    - Using torch.linalg.svd with full_matrices=False")
+                left, _, right = torch.linalg.svd(W_float32, full_matrices=False)
+                return left[:, :k], right[:k, :], k
+
+            def _compute_block_scales(self, weight, per_tensor_scale, rows, cols):
+                fp8, total, mask = super()._compute_block_scales(weight, per_tensor_scale, rows, cols)
+                self._captured_total_scale = total.detach()
+                return fp8, total, mask
+
+            def _optimize_nvfp4(self, *args, **kw):
+                codes, total = super()._optimize_nvfp4(*args, **kw)
+                self._captured_total_scale = total.detach()
+                return codes, total
+
+        self.converter = _Capturing(**kwargs)
+
+    def convert(self, weight: torch.Tensor):
+        return self.converter.convert(weight)
+
+
+def pack_nvfp4(codes: np.ndarray, scale_bits: np.ndarray) -> np.ndarray:
+    """Pack E2M1 codes and UE4M3 block scales into ggml NVFP4 bytes.
+
+    ``codes`` is ``(..., K)`` with K a multiple of 64. Each nibble is an E2M1
+    code. ``scale_bits`` is ``(..., K / 16)``. A super-block is 4 scale bytes
+    followed by 32 code bytes, matching ``block_nvfp4``.
+    """
+    codes = np.ascontiguousarray(codes, dtype=np.uint8)
+    scale_bits = np.ascontiguousarray(scale_bits, dtype=np.uint8) & np.uint8(0x7F)
+    if codes.shape[:-1] != scale_bits.shape[:-1]:
+        raise ValueError(f"code shape {codes.shape} does not match scales {scale_bits.shape}")
+    length = codes.shape[-1]
+    if length % QK_NVFP4 != 0:
+        raise ValueError(f"NVFP4 row length {length} is not a multiple of {QK_NVFP4}")
+    if scale_bits.shape[-1] * 16 != length:
+        raise ValueError("expected one scale per 16 values")
+
+    prefix = codes.shape[:-1]
+    n_super = length // QK_NVFP4
+    grouped = codes.reshape(*prefix, n_super, 4, 16)
+    quant_bytes = (grouped[..., :8] | (grouped[..., 8:] << np.uint8(4))).reshape(*prefix, n_super, 32)
+    scales = scale_bits.reshape(*prefix, n_super, 4)
+    return np.ascontiguousarray(
+        np.concatenate([scales, quant_bytes], axis=-1).reshape(*prefix, n_super * NVFP4_BLOCK_BYTES)
+    )
+
+
+def mtp_layer_tensor(name: str, n_blocks: int, n_nextn: int) -> bool:
+    """True for weights stored in the NextN/MTP blocks after the trunk.
+
+    llama.cpp's normal load skips those blocks before it looks for NVFP4
+    ``.scale`` sidecars, so a sidecar on one of them is never claimed and the
+    loader aborts with ``wrong number of tensors``.
+    """
+    if n_nextn <= 0 or n_blocks <= 0 or not name.startswith("blk."):
+        return False
+    layer = name.split(".", 2)[1]
+    if not layer.isdigit():
+        return False
+    return int(layer) >= n_blocks - n_nextn
+
+
+def bake_packed_nvfp4(raw: np.ndarray, per_tensor: float) -> np.ndarray:
+    """Fold a per-tensor scale into the UE4M3 bytes of packed NVFP4 blocks."""
+    raw = np.ascontiguousarray(raw, dtype=np.uint8)
+    if raw.size == 0 or raw.shape[-1] % NVFP4_BLOCK_BYTES != 0:
+        raise ValueError(f"NVFP4 packing has shape {raw.shape}, which is not groups of {NVFP4_BLOCK_BYTES} bytes")
+    if abs(float(per_tensor) - 1.0) <= 1e-6:
+        return raw
+    blocks = raw.reshape(*raw.shape[:-1], -1, NVFP4_BLOCK_BYTES)
+    scales = np.ascontiguousarray(blocks[..., :4])
+    decoded = torch.from_numpy(scales).view(torch.float8_e4m3fn).float().numpy()
+    encoded = ue4m3_bits(decoded * np.float32(per_tensor)).reshape(scales.shape)
+    baked = np.array(blocks, copy=True)
+    baked[..., :4] = encoded
+    return np.ascontiguousarray(baked.reshape(raw.shape))
+
+
+def ue4m3_bits(scale_float: np.ndarray) -> np.ndarray:
+    """Encode positive scales as UE4M3 using torch's E4M3FN bit patterns.
+
+    ggml's own encoder saturates every value from 256 to 448 down to 448.
+    Torch's float8 bits decode back to the right magnitude with ggml's decoder,
+    including the exponent-15 range that NVFP4 block scales actually use.
+    """
+    array = np.ascontiguousarray(np.clip(scale_float, 0.0, F8_E4M3_MAX), dtype=np.float32)
+    bits = torch.from_numpy(array).to(torch.float8_e4m3fn).view(torch.uint8).numpy()
+    return np.ascontiguousarray(bits & np.uint8(0x7F))
+
+
+def _cuda_python_candidates() -> list[Path]:
+    found: list[Path] = []
+    env_root = Path.home() / "anaconda3" / "envs"
+    if env_root.is_dir():
+        for python in sorted(env_root.glob("*/bin/python")):
+            found.append(python)
+    which = shutil.which("python3")
+    if which:
+        found.append(Path(which))
+    return found
+
+
+def _python_sees_cuda(python: Path) -> bool:
+    if not python.is_file():
+        return False
+    probe = subprocess.run(
+        [str(python), "-c", "import torch, sys; sys.exit(0 if torch.cuda.is_available() else 1)"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        check=False,
+    )
+    return probe.returncode == 0
+
+
+def ensure_cuda_interpreter(device: str) -> None:
+    """Re-run under a CUDA torch when this interpreter cannot see the GPU."""
+    wants_cuda = device == "auto" or str(device).startswith("cuda")
+    if not wants_cuda or torch.cuda.is_available() or os.environ.get("HF_NVFP4_REEXEC") == "1":
+        return
+    current = Path(sys.executable).resolve()
+    candidates = _cuda_python_candidates()
+    candidates.sort(key=lambda path: (0 if "envs/torch/" in path.as_posix() else 1, str(path)))
+    for candidate in candidates:
+        try:
+            if candidate.resolve() == current:
+                continue
+        except OSError:
+            continue
+        if not _python_sees_cuda(candidate):
+            continue
+        print(
+            f"This interpreter ({sys.executable}, torch {torch.__version__}) cannot use CUDA. "
+            f"Restarting with {candidate}",
+            flush=True,
+        )
+        os.environ["HF_NVFP4_REEXEC"] = "1"
+        os.execv(str(candidate), [str(candidate), *sys.argv])
+    if device == "auto":
+        return
+    built = torch.version.cuda or "cpu-only"
+    raise RuntimeError(
+        f"{sys.executable} has torch {torch.__version__} (CUDA {built}) and no visible GPU. "
+        "Pass --device cpu, or run this script with a CUDA build of torch."
+    )
+
+
+def configure_matmul(device: str) -> None:
+    """Full FP32 matmul. convert_to_quant does not enable TF32, and TF32 moves the SVD residual."""
+    if not str(device).startswith("cuda"):
+        return
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    torch.set_float32_matmul_precision("highest")
+
+
+class TensorQuantizer:
+    def __init__(self, args, gguf_module, unpack_uint4):
+        self.args = args
+        self.gguf = gguf_module
+        self.unpack_uint4 = unpack_uint4
+        self.nvfp4_cls = getattr(getattr(gguf_module, "quants", None), "NVFP4", None)
+        optimizer = resolve_optimizer(args.optimizer)
+        self.optimizer_name = optimizer
+        device = args.device
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        if str(device).startswith("cuda") and not torch.cuda.is_available():
+            raise RuntimeError(
+                f"{sys.executable} has torch {torch.__version__} (CUDA {torch.version.cuda or 'cpu-only'}) "
+                "and no visible GPU."
+            )
+        self.device = device
+        configure_matmul(device)
+
+        converter_cls, _ = load_learned_converter(getattr(args, "convert_to_quant", None))
+        lr_min = resolved_lr_min(args)
+        self.engine = GgufLearnedNVFP4(
+            converter_cls,
+            optimizer=optimizer,
+            num_iter=args.num_iter,
+            lr=args.lr,
+            top_p=args.top_p,
+            min_k=args.min_k,
+            max_k=args.max_k,
+            full_matrix=args.full_matrix,
+            no_learned_rounding=args.simple,
+            lr_schedule=args.lr_schedule,
+            lr_gamma=_opt(args, "lr_gamma", 0.99),
+            lr_patience=_opt(args, "lr_patience", DEFAULT_LR_PATIENCE),
+            lr_factor=_opt(args, "lr_factor", 0.95),
+            lr_min=lr_min,
+            lr_cooldown=_opt(args, "lr_cooldown", 0),
+            lr_threshold=_opt(args, "lr_threshold", DEFAULT_LR_THRESHOLD),
+            lr_shape_influence=_opt(args, "lr_shape_influence", 1.0),
+            early_stop_loss=args.early_stop_loss,
+            early_stop_lr=_opt(args, "early_stop_lr", DEFAULT_EARLY_STOP_LR),
+            early_stop_stall=_opt(args, "early_stop_stall", DEFAULT_EARLY_STOP_STALL),
+            scale_refinement_rounds=args.scale_refinement_rounds,
+            scale_optimization=args.scale_optimization,
+            use_speed=_opt(args, "use_speed", False),
+            block_size=16,
+            pad_to_16x=True,
+            device=device,
+        )
+        self.engine.converter.early_stop_floor_stall = _opt(
+            args, "early_stop_floor_stall", DEFAULT_EARLY_STOP_FLOOR_STALL,
+        )
+        self.engine.converter.target_loss = float(_opt(args, "target_loss", DEFAULT_TARGET_LOSS))
+        self.calib_cache: dict[int, torch.Tensor] = {}
+        log.info(
+            "NVFP4 quantizer: %s on %s, %s iterations, scale mode %s, lr patience %s, "
+            "stop at loss %g, fp32 matmul",
+            "simple round-to-grid" if args.simple else f"learned rounding ({optimizer})",
+            device,
+            args.num_iter if not args.simple else 0,
+            args.scale_optimization,
+            _opt(args, "lr_patience", DEFAULT_LR_PATIENCE),
+            self.engine.converter.target_loss,
+        )
+        if not args.simple:
+            log.info("Attention projections stay NVFP4 and are fit across every singular direction.")
+
+    def _extract(self, weight: torch.Tensor, converted) -> tuple[np.ndarray, np.ndarray, float, np.ndarray]:
+        packed, _blocked, per_tensor, dequant, _extra = converted
+        rows, cols = weight.shape
+        if not torch.any(packed):
+            codes = np.zeros((rows, cols), dtype=np.uint8)
+            scales = np.zeros((rows, cols // 16), dtype=np.float32)
+            return codes, scales, 1.0, np.zeros((rows, cols), dtype=np.float32)
+
+        codes_t = self.unpack_uint4(packed.detach().contiguous())
+        codes = codes_t[:rows, :cols].to(dtype=torch.uint8).cpu().numpy()
+        total = self.engine.converter._captured_total_scale
+        if total is None:
+            raise RuntimeError("quantizer did not record a block scale")
+        pts = per_tensor.detach().to(device=total.device, dtype=total.dtype).reshape(-1)[0]
+        pts_value = float(pts.float().cpu())
+        if not np.isfinite(pts_value) or pts_value <= 0.0:
+            raise RuntimeError(f"invalid per-tensor scale {pts_value}")
+        block_scale = (total / pts)[:rows, : cols // 16].float().cpu().numpy()
+        dequant_np = dequant.detach().float().cpu().numpy()[:rows, :cols]
+        return codes, block_scale.astype(np.float32, copy=False), pts_value, dequant_np
+
+    def quantize_matrix(self, weight: np.ndarray, name: str = "") -> tuple[np.ndarray, np.ndarray, float, np.ndarray]:
+        tensor = torch.from_numpy(np.ascontiguousarray(weight, dtype=np.float32))
+        converter = self.engine.converter
+        saved = None
+        if _tight_attention(name):
+            # The stock objective keeps the top fifth of the singular spectrum.
+            # A complete basis makes the loss see every direction this matrix has.
+            saved = (converter.top_p, converter.max_k, converter.full_matrix)
+            converter.top_p = 1.0
+            converter.max_k = max(int(converter.max_k), int(min(weight.shape[-2:])), 1)
+            converter.full_matrix = True
+            log.info("    full-rank fit %s", name)
+        converted = None
+        try:
+            with torch.enable_grad():
+                converted = self.engine.convert(tensor)
+            return self._extract(tensor, converted)
+        finally:
+            if saved is not None:
+                converter.top_p, converter.max_k, converter.full_matrix = saved
+            del tensor, converted
+
+    def bake_global_scale(self, codes: np.ndarray, block_scale: np.ndarray, per_tensor: float) -> tuple[np.ndarray, np.ndarray]:
+        """Fold the per-tensor scale into UE4M3 block scales and keep the learned codes."""
+        product = block_scale * np.float32(per_tensor)
+        return codes, ue4m3_bits(product)
+
+    def runtime_dequant(self, raw: np.ndarray, per_tensor: float | None) -> np.ndarray:
+        if self.gguf is None:
+            raise RuntimeError("runtime dequant needs the gguf module")
+        values = self.gguf.quants.dequantize(raw, self.gguf.GGMLQuantizationType.NVFP4)
+        if per_tensor is not None:
+            values = values * np.float32(per_tensor)
+        return values
+
+    def bias_correction(self, original: np.ndarray, reconstructed: np.ndarray) -> np.ndarray:
+        rows, cols = original.shape
+        samples = self.args.calib_samples
+        cached = self.calib_cache.get(cols)
+        if cached is None or cached.shape[0] != samples:
+            generator = torch.Generator(device="cpu")
+            generator.manual_seed(self.args.seed + cols)
+            cached = torch.randn(samples, cols, generator=generator, dtype=torch.float32)
+            self.calib_cache[cols] = cached
+
+        error = torch.from_numpy(np.ascontiguousarray(original - reconstructed, dtype=np.float32))
+        error = error.to(self.device)
+        accumulator = torch.zeros(rows, dtype=torch.float32)
+        batch = 256
+        for start in range(0, samples, batch):
+            batch_x = cached[start:start + batch].to(self.device)
+            accumulator += (batch_x @ error.T).sum(dim=0).float().cpu()
+            del batch_x
+        del error
+        gc.collect()
+        return (accumulator / samples).numpy().astype(np.float32, copy=False)
+
+
+# Feed-forward matmuls in these graphs pass a null scale, so a sidecar
+# `.scale` tensor is loaded and then ignored. Gemma 4 and Gemma embedding 2
+# do apply it. See llama.cpp src/models/gemma3.cpp build_ffn.
+_FFN_SCALE_IGNORED = frozenset({
+    "gemma",
+    "gemma2",
+    "gemma3",
+    "gemma-embedding",
+})
+
+# gemma3n multiplies the NVFP4 scale only on the attention output and the LM head.
+_GEMMA3N_SCALE_USED = frozenset({"attn_output", "output"})
+
+
+def _lfm_ignores_nvfp4_scale(name: str) -> bool:
+    """LFM2 and LFM3 multiply a sidecar scale on Q, K, V, and the attention output.
+
+    The dense feed-forward, the per-head attention gate, and the short-conv
+    projections are called with a null scale. A short-conv sidecar is also
+    never loaded, so writing one makes the GGUF fail to open.
+    """
+    if _is_ffn_matmul(name):
+        return True
+    if ".shortconv." in name:
+        return True
+    return _weight_leaf(name) == "attn_gate"
+
+
+def _architecture_name(architecture: str | bytes | None) -> str:
+    if architecture is None:
+        return ""
+    if isinstance(architecture, bytes):
+        architecture = architecture.decode("utf-8", "replace")
+    return str(architecture).strip().lower()
+
+
+def _is_spark2_5(architecture: str | bytes | None) -> bool:
+    """True for the spark2_5 GGUF arch and the Spark2_5ForCausalLM class name.
+
+    Muse Spark and DeepSeek DSpark stay out of this profile.
+    """
+    name = _architecture_name(architecture).replace("-", "_").replace(".", "_")
+    head = name.split("for", 1)[0]
+    return head in {"spark2_5", "spark25"}
+
+
+class QuantProfile:
+    """Tensors an architecture leaves in the source float type unless --exclude is set."""
+
+    def __init__(self, arch: str, skip: str, label: str):
+        self.arch = arch
+        self.skip = skip
+        self.label = label
+
+
+# GGUF names are blk.N.attn_gate.weight. The Hugging Face checkpoint calls the
+# same matrix self_attn.g_proj.weight.
+_SPARK_PROFILE = QuantProfile(
+    "spark2_5",
+    r"(?:^|\.)(?:attn_gate|g_proj)\.weight$",
+    "attn_gate.weight",
+)
+
+
+def quant_profile(architecture: str | bytes | None) -> QuantProfile | None:
+    if _is_spark2_5(architecture):
+        return _SPARK_PROFILE
+    return None
+
+
+def _tight_attention(name: str) -> bool:
+    return bool(name) and name.endswith(_TIGHT_ATTENTION)
+
+
+def looks_like_spark(names) -> bool:
+    """True when a weight list has Spark's fused QKV and its attention gate."""
+    has_qkv = False
+    has_gate = False
+    for name in names:
+        if name.endswith("q_k_v_proj.weight"):
+            has_qkv = True
+        if name.endswith("g_proj.weight") or name.endswith("attn_gate.weight"):
+            has_gate = True
+        if has_qkv and has_gate:
+            return True
+    return False
+
+
+def bind_exclude(args, architecture: str | bytes | None) -> None:
+    """Compile args.exclude_re, using the architecture profile when --exclude is absent."""
+    if getattr(args, "_exclude_bound", False):
+        return
+    args._exclude_bound = True
+    profile = quant_profile(architecture)
+    user = getattr(args, "exclude", None)
+    pattern = user
+    if profile is not None and not user:
+        pattern = profile.skip
+        args.exclude = pattern
+        log.info("%s profile: leaving %s in the source float type.", profile.arch, profile.label)
+    elif profile is not None:
+        log.info(
+            "%s profile: --exclude replaces the default skip of %s.",
+            profile.arch,
+            profile.label,
+        )
+    args.exclude_re = re.compile(pattern) if pattern else None
+
+
+def quant_settings(args) -> dict:
+    """Flags that change the bytes in a checkpoint. A mismatch forces --fresh."""
+    return {
+        "version": QUANT_VERSION,
+        "simple": bool(getattr(args, "simple", False)),
+        "optimizer": args.optimizer,
+        "num_iter": args.num_iter,
+        "lr": args.lr,
+        "lr_schedule": args.lr_schedule,
+        "scale_optimization": args.scale_optimization,
+        "scale_refinement_rounds": args.scale_refinement_rounds,
+        "full_matrix": bool(getattr(args, "full_matrix", False)),
+        "top_p": args.top_p,
+        "min_k": args.min_k,
+        "max_k": args.max_k,
+        "early_stop_loss": args.early_stop_loss,
+        "early_stop_stall": _opt(args, "early_stop_stall", DEFAULT_EARLY_STOP_STALL),
+        "lr_threshold": _opt(args, "lr_threshold", DEFAULT_LR_THRESHOLD),
+        "target_loss": float(_opt(args, "target_loss", DEFAULT_TARGET_LOSS)),
+        "calib_samples": args.calib_samples,
+        "bias_correction": not getattr(args, "no_bias_correction", False),
+        "seed": args.seed,
+        "quantize_output": bool(getattr(args, "quantize_output", False)),
+        "exclude": getattr(args, "exclude", None),
+        "lr_patience": _opt(args, "lr_patience", DEFAULT_LR_PATIENCE),
+        "early_stop_lr": _opt(args, "early_stop_lr", DEFAULT_EARLY_STOP_LR),
+    }
+
+
+def require_quant_settings(directory: Path, args) -> None:
+    path = directory / "settings.json"
+    payload = quant_settings(args)
+    if path.is_file():
+        previous = json.loads(path.read_text(encoding="utf-8"))
+        if previous != payload:
+            raise SystemExit(
+                f"Checkpoint settings in {path} do not match this run. "
+                "Pass --fresh to recompute, or rerun with the same quantizer flags."
+            )
+        return
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+
+def _weight_leaf(name: str) -> str:
+    base = name[: -len(".weight")] if name.endswith(".weight") else name
+    return base.rsplit(".", 1)[-1]
+
+
+def _is_ffn_matmul(name: str) -> bool:
+    leaf = _weight_leaf(name)
+    if leaf.startswith("ffn_gate_inp"):
+        return False
+    return leaf.startswith(("ffn_up", "ffn_gate", "ffn_down"))
+
+
+def must_bake_nvfp4_scale(architecture: str | bytes | None, name: str, has_ffn_bias: bool = False) -> bool:
+    """Fold the per-tensor scale into the UE4M3 block scales.
+
+    Most llama.cpp graphs multiply ``name.scale`` after the NVFP4 matmul.
+    Gemma 3's feed-forward graph does not, which leaves those layers far too
+    small and the sampler stuck on one token. A dense FFN that also has a
+    bias cannot carry an external scale: llama.cpp adds the bias first and
+    aborts if both are present.
+    """
+    if has_ffn_bias and _is_ffn_matmul(name):
+        return True
+    arch = _architecture_name(architecture)
+    leaf = _weight_leaf(name)
+    if arch == "gemma4-assistant":
+        return True
+    if arch in {"lfm2", "lfm2moe", "lfm3"} and _lfm_ignores_nvfp4_scale(name):
+        return True
+    if arch == "gemma3n":
+        return leaf not in _GEMMA3N_SCALE_USED
+    if arch in _FFN_SCALE_IGNORED and _is_ffn_matmul(name):
+        return True
+    if arch in {"gemma4", "gemma-embedding2"} and (
+        leaf in {"inp_gate", "proj", "per_layer_model_proj"} or leaf.startswith("ffn_gate_inp")
+    ):
+        return True
+    # spark2-5.cpp passes null scales for the feed-forward, the attention
+    # gate, and the LM head. QKV and attn_output do multiply their sidecar.
+    if _is_spark2_5(arch) and (_is_ffn_matmul(name) or leaf in {"attn_gate", "output"}):
+        return True
+    return False
+
+
+def preserve_scale_tensor(architecture: str | bytes | None, weight_name: str) -> bool:
+    """Keep an existing ``.scale`` tensor that is not the NVFP4 weight scale.
+
+    Gemma 4 multiplies ``ffn_gate_inp.scale`` into the router activation
+    before the matmul. The NVFP4 weight scale is baked instead.
+    """
+    arch = _architecture_name(architecture)
+    if arch not in {"gemma4", "gemma-embedding2"}:
+        return False
+    return _weight_leaf(weight_name).startswith("ffn_gate_inp")
+
+
+def scale_bake_notice(architecture: str | bytes | None) -> str | None:
+    arch = _architecture_name(architecture)
+    if arch in _FFN_SCALE_IGNORED:
+        return (
+            f"{arch} feed-forward layers do not multiply an NVFP4 .scale tensor. "
+            "Those per-tensor scales are baked into the block scales."
+        )
+    if arch == "gemma3n":
+        return (
+            "gemma3n multiplies an NVFP4 .scale only on attn_output and output. "
+            "Every other matrix has its per-tensor scale baked into the block scales."
+        )
+    if arch in {"lfm2", "lfm2moe", "lfm3"}:
+        return (
+            f"{arch} does not multiply an NVFP4 .scale on the feed-forward, "
+            "the attention gate, or the short convolutions. "
+            "Those per-tensor scales are baked into the block scales."
+        )
+    if arch == "gemma4-assistant":
+        return (
+            "gemma4-assistant does not multiply NVFP4 .scale tensors. "
+            "Per-tensor scales are baked into the block scales."
+        )
+    if arch in {"gemma4", "gemma-embedding2"}:
+        return (
+            f"{arch} does not multiply an NVFP4 weight scale on the router, "
+            "the per-layer projections, or per_layer_model_proj. "
+            "Those scales are baked into the block scales."
+        )
+    if _is_spark2_5(arch):
+        return (
+            "spark2_5 does not multiply an NVFP4 .scale on the feed-forward, "
+            "the attention gate, or the output. Those per-tensor scales are baked "
+            "into the block scales."
+        )
+    return None
+
+
+def check_scale_policy() -> None:
+    """Guard the architecture table against the llama.cpp graphs it was read from."""
+    bake = must_bake_nvfp4_scale
+    assert bake("gemma3", "blk.0.ffn_up.weight")
+    assert bake("gemma3", "blk.0.ffn_gate.weight")
+    assert bake("gemma3", "blk.0.ffn_down.weight")
+    assert not bake("gemma3", "blk.0.attn_q.weight")
+    assert not bake("gemma3", "blk.0.attn_output.weight")
+    assert not bake("gemma3", "output.weight")
+    assert bake("gemma", "blk.0.ffn_up.weight")
+    assert bake("gemma2", "blk.0.ffn_down.weight")
+    assert bake(b"gemma3", "blk.0.ffn_up.weight")
+    assert bake("llama", "blk.0.ffn_up.weight", has_ffn_bias=True)
+    assert not bake("llama", "blk.0.ffn_up.weight")
+    assert not bake("llama", "blk.0.attn_q.weight")
+    assert bake("gemma3n", "blk.0.attn_q.weight")
+    assert bake("gemma3n", "blk.0.ffn_up.weight")
+    assert bake("gemma3n", "blk.0.laurel_l.weight")
+    assert not bake("gemma3n", "blk.0.attn_output.weight")
+    assert not bake("gemma3n", "output.weight")
+    assert not bake("gemma4", "blk.0.ffn_up.weight")
+    assert not bake("gemma4", "blk.0.attn_q.weight")
+    assert bake("gemma4", "blk.0.ffn_gate_inp.weight")
+    assert bake("gemma4", "blk.0.inp_gate.weight")
+    assert bake("gemma4", "blk.0.proj.weight")
+    assert bake("gemma4", "per_layer_model_proj.weight")
+    assert bake("gemma-embedding2", "per_layer_model_proj.weight")
+    assert not bake("llama", "per_layer_model_proj.weight")
+    assert preserve_scale_tensor("gemma4", "blk.0.ffn_gate_inp.weight")
+    assert not preserve_scale_tensor("gemma3", "blk.0.ffn_up.weight")
+    assert not preserve_scale_tensor("gemma4", "blk.0.ffn_up.weight")
+    assert bake("gemma4-assistant", "blk.0.attn_q.weight")
+    assert bake("gemma4-assistant", "output.weight")
+    assert bake("lfm2", "blk.0.ffn_up.weight")
+    assert bake("lfm2", "blk.0.ffn_gate.weight")
+    assert bake("lfm2", "blk.0.ffn_down.weight")
+    assert bake("lfm2", "blk.0.shortconv.in_proj.weight")
+    assert bake("lfm2", "blk.0.shortconv.out_proj.weight")
+    assert not bake("lfm2", "blk.0.attn_q.weight")
+    assert not bake("lfm2", "blk.0.attn_output.weight")
+    assert bake("lfm2moe", "blk.0.shortconv.in_proj.weight")
+    assert bake("lfm3", "blk.0.ffn_up.weight")
+    assert bake("lfm3", "blk.0.ffn_gate.weight")
+    assert bake("lfm3", "blk.0.ffn_down.weight")
+    assert bake("lfm3", "blk.4.attn_gate.weight")
+    assert bake("lfm3", "blk.0.shortconv.in_proj.weight")
+    assert bake("lfm3", "blk.0.shortconv.out_proj.weight")
+    assert bake("lfm3", "blk.0.shortconv.conv.weight")
+    assert not bake("lfm3", "blk.2.attn_q.weight")
+    assert not bake("lfm3", "blk.2.attn_k.weight")
+    assert not bake("lfm3", "blk.2.attn_v.weight")
+    assert not bake("lfm3", "blk.2.attn_output.weight")
+    assert not bake("lfm3", "output.weight")
+    assert scale_bake_notice("lfm3")
+    assert not bake("gemma-embedding2", "blk.0.ffn_up.weight")
+    assert bake("gemma-embedding2", "blk.0.inp_gate.weight")
+    assert scale_bake_notice("gemma3")
+    assert scale_bake_notice("llama") is None
+    assert bake("spark2_5", "blk.0.ffn_up.weight")
+    assert bake("spark2_5", "blk.0.ffn_gate.weight")
+    assert bake("spark2_5", "blk.0.ffn_down.weight")
+    assert bake("spark2_5", "blk.4.attn_gate.weight")
+    assert bake("spark2_5", "output.weight")
+    assert not bake("spark2_5", "blk.0.attn_qkv.weight")
+    assert not bake("spark2_5", "blk.0.attn_output.weight")
+    assert not bake("spark2_5", "output_norm.weight")
+    assert bake("Spark2_5ForCausalLM", "blk.3.ffn_up.weight")
+    assert bake("spark25", "blk.0.ffn_down.weight")
+    assert not bake("muse-spark", "blk.0.ffn_up.weight")
+    assert not bake("dspark", "output.weight")
+    assert scale_bake_notice("spark2_5")
+    spark = quant_profile("Spark2_5ForCausalLM")
+    assert spark is not None and spark.arch == "spark2_5"
+    spark_skip = re.compile(spark.skip)
+    assert spark_skip.search("blk.12.attn_gate.weight")
+    assert spark_skip.search("model.layers.0.self_attn.g_proj.weight")
+    assert not spark_skip.search("blk.0.attn_output.weight")
+    assert not spark_skip.search("blk.0.ffn_gate.weight")
+    assert not spark_skip.search("blk.0.attn_qkv.weight")
+    assert not spark_skip.search("model.layers.0.self_attn.attn_gate_proj.weight")
+    assert quant_profile("llama") is None
+    assert quant_profile("muse-spark") is None
+    assert quant_profile("dspark") is None
+    kept = argparse.Namespace(exclude=None)
+    bind_exclude(kept, "spark2_5")
+    assert kept.exclude_re.search("blk.1.attn_gate.weight")
+    replaced = argparse.Namespace(exclude=r"output\.weight")
+    bind_exclude(replaced, "spark2_5")
+    assert replaced.exclude == r"output\.weight"
+    assert not replaced.exclude_re.search("blk.1.attn_gate.weight")
+    assert looks_like_spark(["model.layers.0.self_attn.q_k_v_proj.weight", "model.layers.0.self_attn.g_proj.weight"])
+    assert not looks_like_spark(["blk.0.attn_qkv.weight", "blk.0.attn_gate.weight"])
+    assert _tight_attention("blk.0.attn_output.weight")
+    assert _tight_attention("blk.0.attn_qkv.weight")
+    assert _tight_attention("model.layers.0.self_attn.o_proj.weight")
+    assert _tight_attention("model.layers.0.self_attn.q_k_v_proj.weight")
+    assert not _tight_attention("blk.0.ffn_down.weight")
+    assert not _tight_attention("blk.0.attn_gate.weight")
+    assert not _tight_attention("")
+
+
+def quantize_and_pack(
+    quantizer: TensorQuantizer, weight: np.ndarray, bake: bool, name: str = "",
+) -> tuple[np.ndarray, np.ndarray | None]:
+    """Quantize one matrix, or a stack of expert matrices, into ggml NVFP4 bytes."""
+    if weight.ndim == 2:
+        codes, block_scale, per_tensor, _dequant = quantizer.quantize_matrix(weight, name)
+        return _pack_matrix(quantizer, codes, block_scale, per_tensor, bake)
+    raws = []
+    scales = []
+    for index in range(weight.shape[0]):
+        codes, block_scale, per_tensor, _dequant = quantizer.quantize_matrix(weight[index], name)
+        raw, scale = _pack_matrix(quantizer, codes, block_scale, per_tensor, bake)
+        raws.append(raw)
+        scales.append(1.0 if scale is None else float(scale[0]))
+    stacked = np.stack(raws, axis=0)
+    return stacked, None if bake else np.asarray(scales, dtype=np.float32)
+
+
+def _pack_matrix(quantizer: TensorQuantizer, codes, block_scale, per_tensor: float, bake: bool):
+    if bake:
+        baked_codes, bits = quantizer.bake_global_scale(codes, block_scale, per_tensor)
+        return pack_nvfp4(baked_codes, bits), None
+    raw = pack_nvfp4(codes, ue4m3_bits(block_scale))
+    return raw, np.asarray([per_tensor], dtype=np.float32)
+
+
+def report_reconstruction(quantizer: TensorQuantizer, weight: np.ndarray, raw: np.ndarray, scale: np.ndarray | None) -> None:
+    sample_w = weight[0] if weight.ndim == 3 else weight
+    sample_raw = raw[0] if raw.ndim == 3 else raw
+    per_tensor = None if scale is None else float(np.reshape(scale, -1)[0])
+    reconstructed = quantizer.runtime_dequant(sample_raw, per_tensor)
+    error = sample_w.astype(np.float32) - reconstructed
+    denom = float(np.max(np.abs(sample_w))) + 1e-8
+    rms = float(np.sqrt(np.mean(np.square(sample_w)))) + 1e-12
+    log.info(
+        "    rel-max %.3e  rmse %.3e  rmse/std %.4f  scale %s",
+        float(np.max(np.abs(error))) / denom,
+        float(np.sqrt(np.mean(np.square(error)))),
+        float(np.sqrt(np.mean(np.square(error)))) / rms,
+        "baked" if scale is None else f"{per_tensor:.4e}",
+    )
+
+
+class ShapeView:
+    def __init__(self, shape: tuple[int, ...]):
+        self.shape = shape
+        self.ndim = len(shape)
+
+
+def eligible_weight(name: str, shape: tuple[int, ...], args, gguf_module) -> bool:
+    """Llama.cpp weight that this project quantizes instead of leaving as BF16."""
+    if not name.endswith(".weight"):
+        return False
+    if len(shape) not in (2, 3):
+        return False
+    if len(shape) == 3 and "exps" not in name and "expert" not in name:
+        return False
+    length = int(shape[-1])
+    rows = int(shape[-2])
+    if length < QK_NVFP4 or length % QK_NVFP4 != 0 or rows < 1:
+        return False
+    lowered = name.lower()
+    if "norm" in lowered or "token_embd" in lowered:
+        return False
+    if name == "output.weight" and not getattr(args, "quantize_output", False):
+        return False
+    exclude = getattr(args, "exclude_re", None)
+    if exclude is not None and exclude.search(name):
+        return False
+    return True
+
+
+def add_quant_arguments(parser: argparse.ArgumentParser) -> None:
+    """Flags shared with convert_to_quant's learned NVFP4 defaults."""
+    parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or cuda:N")
+    parser.add_argument("--simple", action="store_true", help="Round onto the E2M1 grid without learned rounding")
+    parser.add_argument("--optimizer", default="auto", choices=("auto", "prodigy", "adamw", "radam", "original"),
+                        help="auto uses prodigy when it is installed, otherwise AdamW")
+    parser.add_argument("--num-iter", type=int, default=CTQ_NUM_ITER,
+                        help="Learned-rounding iterations per matrix (convert_to_quant default: 4000)")
+    parser.add_argument("--lr", type=float, default=1.0, help="Initial learning rate")
+    parser.add_argument("--lr-schedule", default="plateau", choices=("plateau", "exponential", "adaptive"))
+    parser.add_argument(
+        "--scale-optimization", default=DEFAULT_SCALE_OPTIMIZATION, choices=("fixed", "iterative", "joint"),
+        help="fixed keeps the initial scale. iterative refreshes it on the original optimizer. "
+             "joint trains the block scales with the codes (default: joint)",
+    )
+    parser.add_argument(
+        "--scale-refinement-rounds", type=int, default=DEFAULT_SCALE_REFINEMENT_ROUNDS,
+        help="Scale refreshes when --scale-optimization is iterative (default: 2). Ignored for fixed and joint.",
+    )
+    parser.add_argument("--full-matrix", action="store_true", help="Use a full SVD instead of a low-rank SVD")
+    parser.add_argument("--top-p", type=float, default=0.2)
+    parser.add_argument("--min-k", type=int, default=128)
+    parser.add_argument("--max-k", type=int, default=1280)
+    parser.add_argument("--early-stop-loss", type=float, default=5e-9)
+    parser.add_argument("--early-stop-lr", type=float, default=1.01e-8)
+    parser.add_argument(
+        "--early-stop-stall", type=int, default=DEFAULT_EARLY_STOP_STALL,
+        help="Stop after this many steps without a new best, at any learning rate (default: 2000)",
+    )
+    parser.add_argument(
+        "--target-loss", type=float, default=DEFAULT_TARGET_LOSS,
+        help="Stop when the best loss reaches this value, at any learning rate (default: 1e-5). 0 disables it.",
+    )
+    parser.add_argument(
+        "--early-stop-floor-stall", type=int, default=DEFAULT_EARLY_STOP_FLOOR_STALL,
+        help="If the loss never reaches --target-loss, stop after this many steps without a new best once it has otherwise flattened. 0 disables it (default: 64)",
+    )
+    parser.add_argument(
+        "--lr-min", type=float, default=1e-8,
+        help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).",
+    )
+    parser.add_argument("--use-speed", action="store_true", help="Enable Prodigy's use_speed")
+    parser.add_argument("--calib-samples", type=int, default=3072)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--no-bias-correction", action="store_true")
+    parser.add_argument("--quantize-output", action="store_true", help="Also quantize output.weight")
+    parser.add_argument("--exclude", default=None, help="Regex of tensor names to leave unquantized")
+    parser.add_argument("--convert-to-quant", type=Path, default=None,
+                        help="convert_to_quant checkout, if it is not installed")