zefie/nvfp4_gguf_quant
1
0 Fork 0
Code Actions Activity
Files
baa9c09e0e3b7eb6dc25cc80a5b586a1db2fc204
nvfp4_gguf_quant/nvfp4gguf/quant.py
T
zefie baa9c09e0e improve hf_to_nvfp4_safetensors
2026-11-04 14:30:11 -05:00

1492 lines
64 KiB
Python
Raw Blame History

"""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 = 3
QK_NVFP4 = 64
NVFP4_BLOCK_BYTES = 36
F8_E4M3_MAX = 448.0
# The 15 non-zero representable values in FP4 E2M1 (symmetric, unsigned half + zero).
# Codes 0-7 encode [0, 0.5, 1, 1.5, 2, 3, 4, 6]; codes 8-15 are negated.
_FP4_E2M1_VALUES = np.array(
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
dtype=np.float32,
)
FP4_TABLE = _FP4_E2M1_VALUES.astype(np.float32)
FP4_LEVELS = FP4_TABLE
FP4_LEVELS_ORDER = np.argsort(FP4_LEVELS)
FP4_LEVELS_SORTED = FP4_LEVELS[FP4_LEVELS_ORDER]
FP4_LEVEL_MAP = FP4_LEVELS_ORDER.astype(np.uint8)
# Number of candidate scales to try in outlier-aware block scale selection.
_SCALE_CANDIDATES = 8
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 = 5
DEFAULT_EARLY_STOP_LR = 0.0
DEFAULT_LR_FLOOR = 1e-8
# 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.home() / "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 DEFAULT_LR_FLOOR
return float(lr_min)
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 _ue4m3_roundtrip(scale_float: np.ndarray) -> np.ndarray:
"""Encode to UE4M3 and decode back, returning the actual runtime scale values."""
bits = ue4m3_bits(scale_float)
return torch.from_numpy(bits).view(torch.float8_e4m3fn).float().numpy()
def ue4m3_bits_torch(scale_float: torch.Tensor) -> torch.Tensor:
"""Encode positive scales as UE4M3 using torch's float8_e4m3fn directly on tensor/device."""
clipped = torch.clamp(scale_float, 0.0, F8_E4M3_MAX).to(torch.float32)
bits = clipped.to(torch.float8_e4m3fn).view(torch.uint8)
return bits & 0x7F
def _ue4m3_roundtrip_torch(scale_float: torch.Tensor) -> torch.Tensor:
"""Encode to UE4M3 and decode back directly in Torch without CPU/NumPy roundtrips."""
bits = ue4m3_bits_torch(scale_float)
return bits.view(torch.float8_e4m3fn).to(torch.float32)
_UE4M3_LUT: torch.Tensor | None = None
def get_ue4m3_lut(device: str | torch.device | None = None) -> torch.Tensor:
"""Cached 128-element lookup table for all positive UE4M3 float values."""
global _UE4M3_LUT
if _UE4M3_LUT is None:
bits = torch.arange(128, dtype=torch.uint8)
vals = bits.view(torch.float8_e4m3fn).to(torch.float32)
vals[127] = F8_E4M3_MAX # 0x7F is NaN in float8_e4m3fn; saturate to 448.0
_UE4M3_LUT = vals
if device is not None:
target_device = torch.device(device) if isinstance(device, str) else device
if _UE4M3_LUT.device != target_device:
return _UE4M3_LUT.to(device=target_device)
return _UE4M3_LUT
_FP4_TORCH_CACHE: dict[torch.device, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
def get_fp4_torch_constants(device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Pre-loaded FP4 levels, argmin map, and reconstruction table on the target device."""
target_device = torch.device(device) if isinstance(device, str) else device
if target_device not in _FP4_TORCH_CACHE:
levels_sorted = torch.from_numpy(FP4_LEVELS_SORTED).to(device=target_device, dtype=torch.float32)
level_map = torch.from_numpy(FP4_LEVEL_MAP).to(device=target_device, dtype=torch.long)
table = torch.from_numpy(FP4_TABLE).to(device=target_device, dtype=torch.float32)
_FP4_TORCH_CACHE[target_device] = (levels_sorted, level_map, table)
return _FP4_TORCH_CACHE[target_device]
def snap_to_fp4_codes_torch(
scaled: torch.Tensor,
fp4_levels_sorted: torch.Tensor,
fp4_level_map: torch.Tensor,
) -> torch.Tensor:
"""Vectorized O(N) snap of pre-scaled weights to nearest FP4 nibble codes."""
idx = torch.searchsorted(fp4_levels_sorted, scaled)
idx = idx.clamp(1, 15)
left = fp4_levels_sorted[idx - 1]
right = fp4_levels_sorted[idx]
choose_right = (scaled - right).abs() < (scaled - left).abs()
nearest = torch.where(choose_right, idx, idx - 1)
return fp4_level_map[nearest].to(torch.uint8)
def snap_to_fp4_torch(
scaled: torch.Tensor,
fp4_levels_sorted: torch.Tensor,
fp4_level_map: torch.Tensor,
fp4_table: torch.Tensor,
) -> torch.Tensor:
"""Vectorized O(N) snap of pre-scaled weights to reconstructed FP4 float values."""
codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map)
return fp4_table[codes.long()]
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
if limit > 0 and self._learning_rate_floored() and worse_loss_counter >= limit:
log.info(" Loss stalled at lr-min (%d steps). Stopping early.", worse_loss_counter)
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.asarray(codes, dtype=np.uint8)
if not codes.flags.c_contiguous:
codes = np.ascontiguousarray(codes)
scale_bits = np.asarray(scale_bits, dtype=np.uint8) & np.uint8(0x7F)
if not scale_bits.flags.c_contiguous:
scale_bits = np.ascontiguousarray(scale_bits)
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 refine_codes_for_scale_torch(
weight: torch.Tensor,
codes: torch.Tensor,
block_scale: torch.Tensor,
per_tensor: float,
bake: bool,
device: str | torch.device | None = None,
) -> torch.Tensor:
"""Re-select FP4 codes against the actual UE4M3-rounded runtime scale directly on GPU/Torch."""
target_device = device if device is not None else weight.device
w_t = weight.to(device=target_device, dtype=torch.float32)
c_t = codes.to(device=target_device, dtype=torch.uint8)
bs_t = block_scale.to(device=target_device, dtype=torch.float32)
fp4_levels_sorted, fp4_level_map, fp4_table = get_fp4_torch_constants(target_device)
lut = get_ue4m3_lut(device=target_device)
rows, cols = w_t.shape
n_blocks_per_row = cols // 16
pts = max(float(per_tensor), 1e-30)
if bake:
scale_bits = ue4m3_bits_torch(bs_t * pts)
runtime_scale = lut[scale_bits.long()]
else:
scale_bits = ue4m3_bits_torch(bs_t)
runtime_scale = lut[scale_bits.long()] * pts
w_blocks = w_t.reshape(rows, n_blocks_per_row, 16)
c_blocks = c_t.reshape(rows, n_blocks_per_row, 16)
rs_3d = runtime_scale.unsqueeze(-1)
safe_rs = rs_3d.clamp_min(1e-30)
scaled = w_blocks / safe_rs
best_codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map)
old_recon = fp4_table[c_blocks.long()] * rs_3d
new_recon = fp4_table[best_codes.long()] * rs_3d
old_err = (w_blocks - old_recon).square()
new_err = (w_blocks - new_recon).square()
improved = new_err < old_err
refined = torch.where(improved, best_codes, c_blocks)
n_changed = int(improved.sum().item())
if n_changed > 0:
old_mse = float(old_err.mean().item())
new_mixed_err = torch.where(improved, new_err, old_err)
new_mse = float(new_mixed_err.mean().item())
log.info(
" code refinement: %d/%d codes changed, block MSE %.3e -> %.3e",
n_changed, int(refined.numel()), old_mse, new_mse,
)
return refined.reshape(rows, cols)
def optimize_block_scales_torch(
weight: torch.Tensor,
codes: torch.Tensor,
block_scale: torch.Tensor,
per_tensor: float,
bake: bool,
device: str | torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Optimize block scales using windowed discrete UE4M3 local search and outlier mitigation (Optimizations 3 & 5)."""
target_device = device if device is not None else weight.device
w_t = weight.to(device=target_device, dtype=torch.float32)
c_t = codes.to(device=target_device, dtype=torch.uint8)
bs_t = block_scale.to(device=target_device, dtype=torch.float32)
rows, cols = w_t.shape
n_blocks = cols // 16
w_blocks = w_t.reshape(rows, n_blocks, 16)
c_blocks = c_t.reshape(rows, n_blocks, 16)
abs_blocks = w_blocks.abs()
fp4_levels_sorted, fp4_level_map, fp4_table = get_fp4_torch_constants(target_device)
lut = get_ue4m3_lut(device=target_device)
pts = max(float(per_tensor), 1e-30)
if bake:
base_bits = ue4m3_bits_torch(bs_t * pts)
else:
base_bits = ue4m3_bits_torch(bs_t)
sorted_abs, _ = torch.sort(abs_blocks, dim=-1)
scale_denom = 6.0 * pts
# Percentile-based candidate bit patterns
b_max = ue4m3_bits_torch(sorted_abs[:, :, 15] / 6.0 if bake else sorted_abs[:, :, 15] / scale_denom)
b_sub = ue4m3_bits_torch(sorted_abs[:, :, 14] / 6.0 if bake else sorted_abs[:, :, 14] / scale_denom)
b_3rd = ue4m3_bits_torch(sorted_abs[:, :, 13] / 6.0 if bake else sorted_abs[:, :, 13] / scale_denom)
# Outlier-dampened candidates (mitigates isolated spikes stretching scale)
b_damp1 = ue4m3_bits_torch((0.85 * sorted_abs[:, :, 15]) / 6.0 if bake else (0.85 * sorted_abs[:, :, 15]) / scale_denom)
b_damp2 = ue4m3_bits_torch((0.75 * sorted_abs[:, :, 15]) / 6.0 if bake else (0.75 * sorted_abs[:, :, 15]) / scale_denom)
candidate_bit_list: list[torch.Tensor] = []
# Window around current scale (+/- 4 discrete steps in UE4M3 grid)
for delta in (-4, -3, -2, -1, 1, 2, 3, 4):
candidate_bit_list.append(torch.clamp(base_bits.to(torch.int16) + delta, 1, 126).to(torch.uint8))
# Max-derived candidates (the natural block scale from the largest element)
for delta in (-1, 0, 1):
candidate_bit_list.append(torch.clamp(b_max.to(torch.int16) + delta, 1, 126).to(torch.uint8))
# Outlier-clipped scales (2nd and 3rd max)
for delta in (-1, 0, 1):
candidate_bit_list.append(torch.clamp(b_sub.to(torch.int16) + delta, 1, 126).to(torch.uint8))
candidate_bit_list.append(torch.clamp(b_3rd, 1, 126))
# Outlier dampening (85% and 75% of max — clips isolated spikes)
candidate_bit_list.append(torch.clamp(b_damp1, 1, 126))
candidate_bit_list.append(torch.clamp(b_damp2, 1, 126))
# Baseline MSE from original (codes, scale) so we never regress
if bake:
orig_rt = lut[base_bits.long()]
else:
orig_rt = lut[base_bits.long()] * pts
orig_recon = fp4_table[c_blocks.long()] * orig_rt.unsqueeze(-1)
orig_mse = (w_blocks - orig_recon).square().mean(dim=-1)
best_codes = c_blocks.clone()
best_scale = bs_t.clone()
best_mse = orig_mse.clone()
for cand_b in candidate_bit_list:
cand_scale_val = lut[cand_b.long()]
if bake:
cand_rt = cand_scale_val
cand_raw = cand_scale_val / pts
else:
cand_rt = cand_scale_val * pts
cand_raw = cand_scale_val
valid = cand_rt > 0
rs_3d = cand_rt.unsqueeze(-1)
safe_rs = rs_3d.clamp_min(1e-30)
scaled = w_blocks / safe_rs
opt_codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map)
recon = fp4_table[opt_codes.long()] * rs_3d
mse = (w_blocks - recon).square().mean(dim=-1)
accept = (mse < best_mse) & valid
if torch.any(accept):
best_codes = torch.where(accept.unsqueeze(-1), opt_codes, best_codes)
best_scale = torch.where(accept, cand_raw, best_scale)
best_mse = torch.where(accept, mse, best_mse)
n_changed = int((best_scale != bs_t).sum().item())
if n_changed > 0:
orig_total_mse = float(orig_mse.mean().item())
new_total_mse = float(best_mse.mean().item())
log.info(
" scale optimization: %d/%d blocks rescaled, MSE %.3e -> %.3e",
n_changed, int(rows * n_blocks), orig_total_mse, new_total_mse,
)
return best_codes.reshape(rows, cols), best_scale
def coordinate_descent_refine_torch(
weight: torch.Tensor,
codes: torch.Tensor,
block_scale: torch.Tensor,
per_tensor: float,
bake: bool,
rounds: int = DEFAULT_SCALE_REFINEMENT_ROUNDS,
device: str | torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Multi-round coordinate descent alternating block scale optimization and code refinement (Optimization 2)."""
target_device = device if device is not None else weight.device
w_t = weight.to(device=target_device, dtype=torch.float32)
c_t = codes.to(device=target_device, dtype=torch.uint8)
bs_t = block_scale.to(device=target_device, dtype=torch.float32)
# Initial code refinement against the post-UE4M3 runtime scale before
# entering the scale optimization loop. The learned rounding chose codes
# against the full-precision scale, but the runtime scale is quantized to
# UE4M3. Snapping codes first gives the scale optimizer a better baseline.
c_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device)
total_rounds = max(int(rounds), 1)
for r in range(total_rounds):
prev_bs = bs_t
c_t, bs_t = optimize_block_scales_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device)
scale_diff = int((bs_t != prev_bs).sum().item())
if r > 0 and scale_diff == 0:
log.info(" coordinate descent converged at round %d/%d", r + 1, total_rounds)
break
if scale_diff > 0:
c_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device)
return c_t, bs_t
def sweep_per_tensor_scale_torch(
weight: torch.Tensor,
codes: torch.Tensor,
block_scale: torch.Tensor,
per_tensor: float,
bake: bool,
device: str | torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor, float]:
"""Sweep per-tensor scale over fine multipliers to align block scales with UE4M3 mantissa bins (Optimization 4)."""
if bake:
return codes, block_scale, per_tensor
target_device = device if device is not None else weight.device
w_t = weight.to(device=target_device, dtype=torch.float32)
c_t = codes.to(device=target_device, dtype=torch.uint8)
bs_t = block_scale.to(device=target_device, dtype=torch.float32)
rows, cols = w_t.shape
n_blocks = cols // 16
w_blocks = w_t.reshape(rows, n_blocks, 16)
c_blocks = c_t.reshape(rows, n_blocks, 16)
fp4_levels_sorted, fp4_level_map, fp4_table = get_fp4_torch_constants(target_device)
lut = get_ue4m3_lut(device=target_device)
orig_bits = ue4m3_bits_torch(bs_t)
orig_rt = lut[orig_bits.long()] * float(per_tensor)
orig_recon = fp4_table[c_blocks.long()] * orig_rt.unsqueeze(-1)
best_total_mse = float((w_blocks - orig_recon).square().mean().item())
best_pt = per_tensor
best_bs = bs_t
best_codes = c_t
multipliers = [0.93, 0.95, 0.97, 0.98, 0.99, 1.01, 1.02, 1.03, 1.05, 1.07]
for m in multipliers:
cand_pt = per_tensor * m
cand_bs = bs_t / m
cand_bits = ue4m3_bits_torch(cand_bs)
cand_rt = lut[cand_bits.long()] * float(cand_pt)
cand_rs_3d = cand_rt.unsqueeze(-1)
safe_rs = cand_rs_3d.clamp_min(1e-30)
scaled = w_blocks / safe_rs
opt_codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map)
recon = fp4_table[opt_codes.long()] * cand_rs_3d
cand_mse = float((w_blocks - recon).square().mean().item())
if cand_mse < best_total_mse:
best_total_mse = cand_mse
best_pt = cand_pt
best_bs = cand_bs
best_codes = opt_codes.reshape(rows, cols)
# Fine refinement around the best multiplier from the coarse pass
if best_pt != per_tensor:
coarse_m = best_pt / per_tensor
for fine_delta in (-0.005, -0.003, -0.001, 0.001, 0.003, 0.005):
m = coarse_m + fine_delta
cand_pt = per_tensor * m
cand_bs = bs_t / m
cand_bits = ue4m3_bits_torch(cand_bs)
cand_rt = lut[cand_bits.long()] * float(cand_pt)
cand_rs_3d = cand_rt.unsqueeze(-1)
safe_rs = cand_rs_3d.clamp_min(1e-30)
scaled = w_blocks / safe_rs
opt_codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map)
recon = fp4_table[opt_codes.long()] * cand_rs_3d
cand_mse = float((w_blocks - recon).square().mean().item())
if cand_mse < best_total_mse:
best_total_mse = cand_mse
best_pt = cand_pt
best_bs = cand_bs
best_codes = opt_codes.reshape(rows, cols)
if best_pt != per_tensor:
orig_mse_val = float((w_blocks - orig_recon).square().mean().item())
log.info(
" per-tensor scale sweep: adjusted by %+.2f%% (MSE %.3e -> %.3e)",
(best_pt / per_tensor - 1.0) * 100.0,
orig_mse_val,
best_total_mse,
)
return best_codes, best_bs, best_pt
def refine_codes_for_scale(
weight: np.ndarray | torch.Tensor,
codes: np.ndarray | torch.Tensor,
block_scale: np.ndarray | torch.Tensor,
per_tensor: float,
bake: bool,
device: str | torch.device | None = None,
) -> np.ndarray | torch.Tensor:
"""Re-select FP4 codes against the actual UE4M3-rounded runtime scale."""
if torch.is_tensor(weight):
return refine_codes_for_scale_torch(weight, codes, block_scale, per_tensor, bake, device=device)
dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
w_t = torch.from_numpy(weight).to(dev)
c_t = torch.from_numpy(codes).to(dev)
bs_t = torch.from_numpy(block_scale).to(dev)
refined_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev)
return refined_t.cpu().numpy()
def optimize_block_scales(
weight: np.ndarray | torch.Tensor,
codes: np.ndarray | torch.Tensor,
block_scale: np.ndarray | torch.Tensor,
per_tensor: float,
bake: bool,
device: str | torch.device | None = None,
) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]:
"""Try multiple candidate scales per block and pick the lowest-MSE pair."""
if torch.is_tensor(weight):
return optimize_block_scales_torch(weight, codes, block_scale, per_tensor, bake, device=device)
dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
w_t = torch.from_numpy(weight).to(dev)
c_t = torch.from_numpy(codes).to(dev)
bs_t = torch.from_numpy(block_scale).to(dev)
best_codes_t, best_scale_t = optimize_block_scales_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev)
return best_codes_t.cpu().numpy(), best_scale_t.cpu().numpy()
def coordinate_descent_refine(
weight: np.ndarray | torch.Tensor,
codes: np.ndarray | torch.Tensor,
block_scale: np.ndarray | torch.Tensor,
per_tensor: float,
bake: bool,
rounds: int = DEFAULT_SCALE_REFINEMENT_ROUNDS,
device: str | torch.device | None = None,
) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]:
"""Multi-round coordinate descent alternating block scale optimization and code refinement (Optimization 2)."""
if torch.is_tensor(weight):
return coordinate_descent_refine_torch(weight, codes, block_scale, per_tensor, bake, rounds=rounds, device=device)
dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
w_t = torch.from_numpy(weight).to(dev)
c_t = torch.from_numpy(codes).to(dev)
bs_t = torch.from_numpy(block_scale).to(dev)
best_c_t, best_bs_t = coordinate_descent_refine_torch(w_t, c_t, bs_t, per_tensor, bake, rounds=rounds, device=dev)
return best_c_t.cpu().numpy(), best_bs_t.cpu().numpy()
def sweep_per_tensor_scale(
weight: np.ndarray | torch.Tensor,
codes: np.ndarray | torch.Tensor,
block_scale: np.ndarray | torch.Tensor,
per_tensor: float,
bake: bool,
device: str | torch.device | None = None,
) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor, float]:
"""Sweep per-tensor scale over fine multipliers to minimize global reconstruction error (Optimization 4)."""
if torch.is_tensor(weight):
return sweep_per_tensor_scale_torch(weight, codes, block_scale, per_tensor, bake, device=device)
dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
w_t = torch.from_numpy(weight).to(dev)
c_t = torch.from_numpy(codes).to(dev)
bs_t = torch.from_numpy(block_scale).to(dev)
best_c_t, best_bs_t, best_pt = sweep_per_tensor_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev)
return best_c_t.cpu().numpy(), best_bs_t.cpu().numpy(), best_pt
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
cached_device = cached.to(self.device)
error = torch.from_numpy(np.ascontiguousarray(original - reconstructed, dtype=np.float32)).to(self.device)
accumulator = torch.zeros(rows, dtype=torch.float32, device=self.device)
batch = 256
for start in range(0, samples, batch):
batch_x = cached_device[start:start + batch]
accumulator += (batch_x @ error.T).sum(dim=0)
del error, cached_device
if str(self.device).startswith("cuda"):
torch.cuda.empty_cache()
gc.collect()
return (accumulator / samples).cpu().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."""
rounds = getattr(quantizer.args, "scale_refinement_rounds", DEFAULT_SCALE_REFINEMENT_ROUNDS)
if weight.ndim == 2:
codes, block_scale, per_tensor, _dequant = quantizer.quantize_matrix(weight, name)
codes, block_scale = coordinate_descent_refine(
weight, codes, block_scale, per_tensor, bake, rounds=rounds, device=quantizer.device,
)
codes, block_scale, per_tensor = sweep_per_tensor_scale(
weight, codes, block_scale, per_tensor, bake, device=quantizer.device,
)
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)
codes, block_scale = coordinate_descent_refine(
weight[index], codes, block_scale, per_tensor, bake, rounds=rounds, device=quantizer.device,
)
codes, block_scale, per_tensor = sweep_per_tensor_scale(
weight[index], codes, block_scale, per_tensor, bake, device=quantizer.device,
)
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=DEFAULT_LR_FLOOR,
help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).",
)
parser.add_argument(
"--lr-patience", type=int, default=DEFAULT_LR_PATIENCE,
help="Patience for learning rate decay (default: 5)",
)
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")
View Git Blame Copy Permalink