"""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")