initial commit

This commit is contained in:
zefie committed 2026-11-01 16:31:33 -05:00
commit db7f28d087
31 files changed
+7462

No files matched your search

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