1120 lines
48 KiB
Python
1120 lines
48 KiB
Python
#!/usr/bin/env python3
|
|
"""Quantize a Hugging Face model to llama.cpp NVFP4 GGUF.
|
|
|
|
Quantization lives in nvfp4gguf. This script loads the model through llama.cpp
|
|
so the vocabulary, tensor names, and architecture metadata match a normal GGUF.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import re
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from nvfp4gguf 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,
|
|
TensorQuantizer,
|
|
bind_exclude,
|
|
eligible_weight,
|
|
ensure_cuda_interpreter,
|
|
load_gguf_modules,
|
|
load_learned_converter,
|
|
must_bake_nvfp4_scale,
|
|
mtp_layer_tensor,
|
|
bake_packed_nvfp4,
|
|
preserve_scale_tensor,
|
|
quantize_and_pack,
|
|
report_reconstruction,
|
|
require_quant_settings,
|
|
scale_bake_notice,
|
|
resolve_optimizer,
|
|
)
|
|
from nvfp4gguf.bundle import (
|
|
apply_hf_context_length,
|
|
chat_template_has_tools,
|
|
flatten_chat_template,
|
|
normalize_chat_template,
|
|
outtype as _outtype,
|
|
place_mmprojs as _place_mmprojs,
|
|
quote_modelfile_value as _quote_modelfile_value,
|
|
refuse_existing_bundle,
|
|
tool_modelfile_lines,
|
|
warn_missing_gguf_metadata,
|
|
)
|
|
|
|
log = logging.getLogger("hf_to_nvfp4_gguf")
|
|
|
|
|
|
def _checkpoint_path(directory: Path, name: str) -> Path:
|
|
digest = hashlib.sha1(name.encode("utf-8")).hexdigest()[:16]
|
|
safe = re.sub(r"[^A-Za-z0-9._-]+", "_", name)[:80]
|
|
return directory / f"{safe}.{digest}.npz"
|
|
|
|
|
|
def _eligible(name: str, tensor, raw_dtype, args, gguf_module) -> bool:
|
|
if raw_dtype != gguf_module.GGMLQuantizationType.BF16:
|
|
return False
|
|
shape = tuple(int(dim) for dim in tensor.shape)
|
|
return eligible_weight(name, shape, args, gguf_module)
|
|
|
|
|
|
def _count_quant_targets(model, args, gguf_module) -> int:
|
|
"""Count matrices the writer will quantize, using shapes from the file header."""
|
|
total = 0
|
|
for hf_name, lazy in model.get_tensors():
|
|
try:
|
|
shape = tuple(int(dim) for dim in lazy.shape)
|
|
except Exception:
|
|
continue
|
|
if len(shape) < 2 or not str(hf_name).endswith(".weight"):
|
|
continue
|
|
try:
|
|
gguf_name = model.map_tensor_name(hf_name)
|
|
except Exception:
|
|
gguf_name = str(hf_name)
|
|
raw_dtype = gguf_module.GGMLQuantizationType.BF16
|
|
if gguf_name.endswith("_norm.weight"):
|
|
raw_dtype = gguf_module.GGMLQuantizationType.F32
|
|
if raw_dtype == gguf_module.GGMLQuantizationType.BF16 and eligible_weight(gguf_name, shape, args, gguf_module):
|
|
total += 1
|
|
return total
|
|
|
|
|
|
def _materialize(array) -> np.ndarray:
|
|
gguf_lazy = sys.modules.get("gguf")
|
|
lazy_cls = getattr(gguf_lazy, "LazyNumpyTensor", None) if gguf_lazy else None
|
|
if lazy_cls is not None and isinstance(array, lazy_cls):
|
|
array = lazy_cls.to_eager(array)
|
|
array = np.asarray(array)
|
|
if array.dtype != np.float32:
|
|
array = array.astype(np.float32, copy=False)
|
|
return np.ascontiguousarray(array)
|
|
|
|
|
|
def _as_float_tensor(array, gguf_module):
|
|
lazy_cls = getattr(gguf_module, "LazyNumpyTensor", None)
|
|
chunk_cls = getattr(gguf_module, "LazyChunkedTensor", None)
|
|
if chunk_cls is not None and isinstance(array, chunk_cls):
|
|
return array
|
|
if lazy_cls is not None and isinstance(array, lazy_cls):
|
|
return array
|
|
return np.asarray(array)
|
|
|
|
|
|
class Nvfp4Writer:
|
|
def __init__(self, model, quantizer: TensorQuantizer, args, bias_names: set[str]):
|
|
self.model = model
|
|
self.quantizer = quantizer
|
|
self.args = args
|
|
self.gguf = quantizer.gguf
|
|
self.bias_names = set(bias_names)
|
|
self.bias_corrections: dict[str, np.ndarray] = {}
|
|
self.pending_bias: dict[str, tuple] = {}
|
|
self.quantized = 0
|
|
self.baked = 0
|
|
self.corrected = 0
|
|
self.passthrough = 0
|
|
self.planned = _count_quant_targets(model, args, self.gguf)
|
|
self.quant_index = 0
|
|
if self.planned:
|
|
log.info("%d weight matrices queued for NVFP4", self.planned)
|
|
names = getattr(self.gguf, "MODEL_ARCH_NAMES", {})
|
|
self.architecture = names.get(getattr(model, "model_arch", None), "") or ""
|
|
notice = scale_bake_notice(self.architecture)
|
|
if notice:
|
|
log.warning(notice)
|
|
self.ckpt_dir = Path(args.checkpoint_dir) if args.checkpoint_dir else Path(str(args.output) + ".nvfp4-work")
|
|
self.ckpt_dir.mkdir(parents=True, exist_ok=True)
|
|
self._orig_quantize = self.gguf.quants.quantize
|
|
self._orig_add = model.gguf_writer.add_tensor
|
|
self._orig_prepare = model.prepare_tensors
|
|
self._orig_metadata = model.prepare_metadata
|
|
|
|
def install(self) -> None:
|
|
gguf_module = self.gguf
|
|
|
|
def quantize(data, qtype):
|
|
# Keep the source float for matrices we may repack. Everything else
|
|
# still goes through llama.cpp's quantizer.
|
|
if qtype == gguf_module.GGMLQuantizationType.BF16:
|
|
return data
|
|
return self._orig_quantize(data, qtype)
|
|
|
|
self.gguf.quants.quantize = quantize
|
|
self.model.gguf_writer.add_tensor = self._add_tensor
|
|
self.model.prepare_tensors = self._prepare_tensors
|
|
self.model.prepare_metadata = self._prepare_metadata
|
|
|
|
def restore(self) -> None:
|
|
self.gguf.quants.quantize = self._orig_quantize
|
|
self.model.gguf_writer.add_tensor = self._orig_add
|
|
self.model.prepare_tensors = self._orig_prepare
|
|
self.model.prepare_metadata = self._orig_metadata
|
|
|
|
def _prepare_tensors(self) -> None:
|
|
self._orig_prepare()
|
|
self._flush_pending()
|
|
|
|
def _prepare_metadata(self, vocab_only: bool = False) -> None:
|
|
if self.quantized:
|
|
self.model.ftype = self.gguf.LlamaFileType.MOSTLY_NVFP4
|
|
self._orig_metadata(vocab_only)
|
|
|
|
def _write(self, name, tensor, raw_dtype=None) -> None:
|
|
self._orig_add(name, tensor, raw_dtype=raw_dtype)
|
|
|
|
def _add_tensor(self, name, tensor, raw_shape=None, raw_dtype=None, tensor_endianess=None) -> None:
|
|
del raw_shape, tensor_endianess
|
|
if name.endswith(".bias"):
|
|
self._handle_bias(name, tensor, raw_dtype)
|
|
return
|
|
if _eligible(name, tensor, raw_dtype, self.args, self.gguf):
|
|
self._quantize_and_write(name, tensor)
|
|
return
|
|
if raw_dtype == self.gguf.GGMLQuantizationType.BF16:
|
|
tensor = self._orig_quantize(tensor, raw_dtype)
|
|
self._write(name, tensor, raw_dtype=raw_dtype)
|
|
self.passthrough += 1
|
|
|
|
def _handle_bias(self, name: str, tensor, raw_dtype) -> None:
|
|
correction = self.bias_corrections.get(name)
|
|
if correction is None and not self.args.no_bias_correction:
|
|
self.pending_bias[name] = (tensor, raw_dtype)
|
|
return
|
|
self._emit_bias(name, tensor, raw_dtype, correction)
|
|
|
|
def _emit_bias(self, name: str, tensor, raw_dtype, correction: np.ndarray | None) -> None:
|
|
if correction is not None:
|
|
values = _materialize(tensor).astype(np.float32, copy=True)
|
|
flat = values.reshape(-1)
|
|
if flat.shape == correction.shape:
|
|
flat -= correction
|
|
values = flat.reshape(values.shape)
|
|
self.corrected += 1
|
|
log.info(" bias %s adjusted (mean %+.3e)", name, float(correction.mean()))
|
|
else:
|
|
log.warning(" bias %s shape %s does not match correction %s", name, values.shape, correction.shape)
|
|
tensor = values.astype(np.float32, copy=False)
|
|
if raw_dtype == self.gguf.GGMLQuantizationType.BF16:
|
|
tensor = self._orig_quantize(tensor, raw_dtype)
|
|
elif raw_dtype not in (None, self.gguf.GGMLQuantizationType.F32):
|
|
tensor = self._orig_quantize(tensor, raw_dtype)
|
|
elif raw_dtype == self.gguf.GGMLQuantizationType.BF16:
|
|
tensor = self._orig_quantize(tensor, raw_dtype)
|
|
self._write(name, tensor, raw_dtype=raw_dtype)
|
|
self.passthrough += 1
|
|
|
|
def _flush_pending(self) -> None:
|
|
for name, (tensor, raw_dtype) in list(self.pending_bias.items()):
|
|
self._emit_bias(name, tensor, raw_dtype, self.bias_corrections.get(name))
|
|
self.pending_bias.clear()
|
|
|
|
def _load_checkpoint(self, name: str, logical_shape: tuple[int, ...]):
|
|
path = _checkpoint_path(self.ckpt_dir, name)
|
|
if not path.is_file():
|
|
return None
|
|
with np.load(path) as blob:
|
|
if int(blob["version"]) != QUANT_VERSION:
|
|
return None
|
|
if tuple(int(v) for v in blob["shape"]) != tuple(logical_shape):
|
|
return None
|
|
raw = np.array(blob["raw"], copy=True)
|
|
scale = np.array(blob["scale"], copy=True)
|
|
bias = np.array(blob["bias"], copy=True)
|
|
scale_out = None if scale.size == 0 else scale
|
|
bias_out = None if bias.size == 0 else bias
|
|
return raw, scale_out, bias_out
|
|
|
|
def _save_checkpoint(self, name: str, logical_shape, raw, scale, bias) -> None:
|
|
path = _checkpoint_path(self.ckpt_dir, name)
|
|
tmp = path.with_name(path.name + ".tmp")
|
|
with tmp.open("wb") as handle:
|
|
np.savez_compressed(
|
|
handle,
|
|
version=np.int32(QUANT_VERSION),
|
|
shape=np.array(logical_shape, dtype=np.int64),
|
|
raw=raw,
|
|
scale=np.zeros((0,), np.float32) if scale is None else scale.astype(np.float32, copy=False),
|
|
bias=np.zeros((0,), np.float32) if bias is None else bias.astype(np.float32, copy=False),
|
|
)
|
|
tmp.replace(path)
|
|
|
|
def _quantize_matrix_group(self, name: str, weight: np.ndarray) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None]:
|
|
baked = self._must_bake(name)
|
|
raw, scale = quantize_and_pack(self.quantizer, weight, baked, name)
|
|
bias = None
|
|
bias_name = name.replace(".weight", ".bias")
|
|
if weight.ndim == 2 and not self.args.no_bias_correction and bias_name in self.bias_names:
|
|
per_tensor = None if scale is None else float(np.reshape(scale, -1)[0])
|
|
runtime = self.quantizer.runtime_dequant(raw, per_tensor)
|
|
bias = self.quantizer.bias_correction(weight, runtime)
|
|
report_reconstruction(self.quantizer, weight, raw, scale)
|
|
return raw, scale, bias
|
|
|
|
def _must_bake(self, name: str) -> bool:
|
|
bias_name = name.replace(".weight", ".bias")
|
|
has_bias = bias_name in self.bias_names or bias_name in self.pending_bias
|
|
if must_bake_nvfp4_scale(self.architecture, name, has_bias):
|
|
return True
|
|
return mtp_layer_tensor(name, *_mtp_depth(self.model))
|
|
|
|
def _progress(self, name: str, shape: tuple[int, ...], checkpoint: bool) -> None:
|
|
self.quant_index += 1
|
|
if self.planned:
|
|
slot = f"[{self.quant_index}/{self.planned}]"
|
|
else:
|
|
slot = f"[{self.quant_index}]"
|
|
note = " (checkpoint)" if checkpoint else ""
|
|
log.info("NVFP4 %s %s %s%s", slot, name, shape, note)
|
|
|
|
def _quantize_and_write(self, name: str, tensor) -> None:
|
|
weight = _materialize(tensor)
|
|
baked = self._must_bake(name)
|
|
cached = self._load_checkpoint(name, weight.shape)
|
|
if cached is not None and (cached[1] is None) != baked:
|
|
log.info(" checkpoint scale layout does not match %s; recomputing", name)
|
|
cached = None
|
|
self._progress(name, tuple(weight.shape), cached is not None)
|
|
if cached is None:
|
|
raw, scale, bias = self._quantize_matrix_group(name, weight)
|
|
self._save_checkpoint(name, weight.shape, raw, scale, bias)
|
|
else:
|
|
raw, scale, bias = cached
|
|
if bias is not None:
|
|
self.bias_corrections[name.replace(".weight", ".bias")] = bias
|
|
pending = self.pending_bias.pop(name.replace(".weight", ".bias"), None)
|
|
if pending is not None:
|
|
self._emit_bias(name.replace(".weight", ".bias"), pending[0], pending[1], bias)
|
|
if self._must_bake(name) and scale is None:
|
|
self.baked += 1
|
|
self._write(name, raw, raw_dtype=self.gguf.GGMLQuantizationType.NVFP4)
|
|
if scale is not None:
|
|
scale_name = name[: -len(".weight")] + ".scale"
|
|
self._write(scale_name, np.ascontiguousarray(scale, dtype=np.float32))
|
|
self.quantized += 1
|
|
del weight, raw
|
|
|
|
|
|
def _accept_list_extra_special_tokens() -> None:
|
|
"""Let transformers 4.x load a Gemma 4 tokenizer.
|
|
|
|
Gemma 4 stores extra_special_tokens as a list of token strings. The 4.x
|
|
loader calls .keys() on that field, which is what aborts set_vocab after
|
|
the weights have already been packed.
|
|
"""
|
|
try:
|
|
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
|
|
except ImportError:
|
|
return
|
|
original = PreTrainedTokenizerBase._set_model_specific_special_tokens
|
|
if getattr(original, "_accepts_token_list", False):
|
|
return
|
|
|
|
def _set_model_specific_special_tokens(self, special_tokens):
|
|
if isinstance(special_tokens, list):
|
|
log.info(
|
|
"tokenizer extra_special_tokens is a list; treating each entry as its own special token."
|
|
)
|
|
special_tokens = {str(token): str(token) for token in special_tokens}
|
|
return original(self, special_tokens)
|
|
|
|
_set_model_specific_special_tokens._accepts_token_list = True
|
|
PreTrainedTokenizerBase._set_model_specific_special_tokens = _set_model_specific_special_tokens
|
|
|
|
|
|
def _accept_missing_tokenizer_class() -> None:
|
|
"""Load tokenizer.json when tokenizer_class is newer than this transformers.
|
|
|
|
LFM3 names its tokenizer ``TokenizersBackend``. That class arrived in
|
|
transformers 5, and a 4.x install otherwise stops before the vocab is written.
|
|
"""
|
|
try:
|
|
from transformers import AutoTokenizer, PreTrainedTokenizerFast
|
|
except ImportError:
|
|
return
|
|
current = AutoTokenizer.from_pretrained
|
|
if getattr(current, "_accepts_missing_class", False):
|
|
return
|
|
original = current.__func__
|
|
|
|
def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
|
|
try:
|
|
return original(cls, pretrained_model_name_or_path, *args, **kwargs)
|
|
except ValueError as exc:
|
|
if "does not exist or is not currently imported" not in str(exc):
|
|
raise
|
|
log.info(
|
|
"Tokenizer class is not in this transformers build; loading tokenizer.json directly."
|
|
)
|
|
return PreTrainedTokenizerFast.from_pretrained(
|
|
pretrained_model_name_or_path, *args, **kwargs
|
|
)
|
|
|
|
from_pretrained._accepts_missing_class = True
|
|
wrapped = classmethod(from_pretrained)
|
|
wrapped._accepts_missing_class = True
|
|
AutoTokenizer.from_pretrained = wrapped
|
|
|
|
|
|
def _install_tokenizer_fallback(model) -> None:
|
|
"""Use tokenizer.json when sentencepiece is not installed and no tokenizer.model is present."""
|
|
_accept_list_extra_special_tokens()
|
|
_accept_missing_tokenizer_class()
|
|
try:
|
|
import sentencepiece # noqa: F401
|
|
return
|
|
except ImportError:
|
|
pass
|
|
if (model.dir_model / "tokenizer.model").is_file():
|
|
original = model.set_vocab
|
|
|
|
def set_vocab_needs_sentencepiece() -> None:
|
|
try:
|
|
original()
|
|
except ModuleNotFoundError as exc:
|
|
if getattr(exc, "name", None) == "sentencepiece":
|
|
raise ModuleNotFoundError(
|
|
"This model ships tokenizer.model, which needs the sentencepiece package. "
|
|
"pip install sentencepiece"
|
|
) from exc
|
|
raise
|
|
|
|
model.set_vocab = set_vocab_needs_sentencepiece
|
|
return
|
|
if not (model.dir_model / "tokenizer.json").is_file():
|
|
return
|
|
|
|
def set_vocab() -> None:
|
|
config_path = model.dir_model / "tokenizer_config.json"
|
|
if config_path.is_file():
|
|
with config_path.open("r", encoding="utf-8") as handle:
|
|
prefix = json.load(handle).get("add_prefix_space")
|
|
if prefix is not None:
|
|
model.gguf_writer.add_add_space_prefix(prefix)
|
|
log.info("sentencepiece is not installed; loading tokenizer.json")
|
|
model._set_vocab_gpt2()
|
|
|
|
model.set_vocab = set_vocab
|
|
|
|
|
|
def _collect_bias_names(model) -> set[str]:
|
|
names: set[str] = set()
|
|
for hf_name in model.model_tensors:
|
|
if not hf_name.endswith(".bias"):
|
|
continue
|
|
try:
|
|
names.add(model.map_tensor_name(hf_name))
|
|
except Exception:
|
|
continue
|
|
return names
|
|
|
|
|
|
def _float_file_type(gguf, outtype: str):
|
|
return {
|
|
"f32": gguf.LlamaFileType.ALL_F32,
|
|
"f16": gguf.LlamaFileType.MOSTLY_F16,
|
|
"bf16": gguf.LlamaFileType.MOSTLY_BF16,
|
|
"auto": gguf.LlamaFileType.GUESSED,
|
|
}[outtype]
|
|
|
|
|
|
_PACKED_WEIGHT_DTYPES = {"U8", "I8"}
|
|
_PACKED_SCALE_DTYPES = {"U8", "I8", "F8_E4M3", "F8_E5M2"}
|
|
_FP8_DTYPES = {"F8_E4M3", "F8_E5M2"}
|
|
|
|
|
|
def _numel(shape: tuple[int, ...]) -> int:
|
|
count = 1
|
|
for dim in shape:
|
|
count *= int(dim)
|
|
return count
|
|
|
|
|
|
def _is_packed_nvfp4_layout(
|
|
weight_dtype: str,
|
|
weight_shape: tuple[int, ...],
|
|
scale_dtype: str,
|
|
scale_shape: tuple[int, ...],
|
|
) -> bool:
|
|
"""Match llama.cpp's packed-4-bit check using safetensors header dtypes.
|
|
|
|
Packed NVFP4 weights are U8 or I8, two 4-bit values per byte, with one
|
|
scale byte per 16 values (8 stored weight bytes per scale). A weight
|
|
labeled F8_E4M3 is an 8-bit float, not packed 4-bit data.
|
|
"""
|
|
if weight_dtype not in _PACKED_WEIGHT_DTYPES or scale_dtype not in _PACKED_SCALE_DTYPES:
|
|
return False
|
|
if len(weight_shape) < 2 or len(scale_shape) < 2:
|
|
return False
|
|
rows = int(weight_shape[0])
|
|
n_blocks = int(scale_shape[1])
|
|
if rows <= 0 or n_blocks <= 0 or n_blocks % 4 != 0:
|
|
return False
|
|
return _numel(weight_shape) == rows * n_blocks * 8 and _numel(scale_shape) == rows * n_blocks
|
|
|
|
|
|
def _scale_name_for_weight(name: str) -> str | None:
|
|
if name.endswith(".weight_packed"):
|
|
return name.removesuffix("_packed") + "_scale"
|
|
if name.endswith(".weight"):
|
|
return name[: -len(".weight")] + ".weight_scale"
|
|
return None
|
|
|
|
|
|
def _sample_names(names: list[str], limit: int = 4) -> str:
|
|
if len(names) <= limit:
|
|
return ", ".join(names)
|
|
return ", ".join(names[:limit]) + f", ... ({len(names) - limit} more)"
|
|
|
|
|
|
class _TensorLayout:
|
|
def __init__(self) -> None:
|
|
self.packed: list[str] = []
|
|
self.fp8: list[str] = []
|
|
self.unrecognized: list[str] = []
|
|
|
|
@property
|
|
def found(self) -> bool:
|
|
return bool(self.packed or self.fp8 or self.unrecognized)
|
|
|
|
|
|
def _checkpoint_tensor_layout(model_dir: Path) -> _TensorLayout:
|
|
"""Classify safetensors by dtype and shape without reading the bytes."""
|
|
layout = _TensorLayout()
|
|
files = sorted(model_dir.glob("*.safetensors"))
|
|
if not files:
|
|
return layout
|
|
try:
|
|
from safetensors import safe_open
|
|
except ImportError:
|
|
log.warning("safetensors is not installed, so packed 4-bit tensors were not prechecked.")
|
|
return layout
|
|
meta: dict[str, tuple[str, tuple[int, ...]]] = {}
|
|
for path in files:
|
|
with safe_open(path, framework="pt", device="cpu") as handle:
|
|
for key in handle.keys():
|
|
header = handle.get_slice(key)
|
|
meta[key] = (str(header.get_dtype()), tuple(int(dim) for dim in header.get_shape()))
|
|
for name, (dtype, shape) in meta.items():
|
|
scale_name = _scale_name_for_weight(name)
|
|
if scale_name is None:
|
|
continue
|
|
scale = meta.get(scale_name)
|
|
if scale is not None and _is_packed_nvfp4_layout(dtype, shape, scale[0], scale[1]):
|
|
layout.packed.append(name)
|
|
continue
|
|
if dtype in _FP8_DTYPES:
|
|
layout.fp8.append(name)
|
|
elif dtype in {"U8", "I8"} and scale is not None:
|
|
layout.unrecognized.append(name)
|
|
layout.packed.sort()
|
|
layout.fp8.sort()
|
|
layout.unrecognized.sort()
|
|
return layout
|
|
|
|
|
|
def _log_tensor_layout(layout: _TensorLayout) -> None:
|
|
if not layout.found:
|
|
return
|
|
log.info(
|
|
"Safetensors layout: %d packed 4-bit, %d FP8 weights that are not 4-bit, %d other byte weights that are not packed 4-bit.",
|
|
len(layout.packed),
|
|
len(layout.fp8),
|
|
len(layout.unrecognized),
|
|
)
|
|
if layout.packed:
|
|
log.info("Packed 4-bit, repacked as NVFP4: %s", _sample_names(layout.packed))
|
|
if layout.fp8:
|
|
log.info("FP8, stored as float: %s", _sample_names(layout.fp8))
|
|
if layout.unrecognized:
|
|
log.warning(
|
|
"These tensors are U8 or I8 but do not match packed 4-bit NVFP4, so they will not be repacked: %s",
|
|
_sample_names(layout.unrecognized),
|
|
)
|
|
|
|
|
|
def _checkpoint_is_nvfp4(model_dir: Path, hparams: dict) -> bool:
|
|
"""Match llama.cpp's NVFP4 detection in ModelBase.prepare_tensors."""
|
|
quantization_config = hparams.get("quantization_config") or {}
|
|
quant_algo = quantization_config.get("quant_algo")
|
|
quant_method = quantization_config.get("quant_method")
|
|
quant_format = quantization_config.get("format")
|
|
quant_groups = quantization_config.get("config_groups") or {}
|
|
quant_layers = quantization_config.get("quantized_layers") or {}
|
|
quant_config_file = model_dir / "hf_quant_config.json"
|
|
if (not quant_algo or not quant_layers) and quant_config_file.is_file():
|
|
with quant_config_file.open("r", encoding="utf-8") as handle:
|
|
hf_quant_config = json.load(handle)
|
|
quant_config = hf_quant_config.get("quantization") or {}
|
|
producer = hf_quant_config.get("producer") or {}
|
|
if quant_method is None:
|
|
quant_method = (producer.get("name") or "").lower()
|
|
quant_algo = quant_config.get("quant_algo", quant_algo)
|
|
quant_method = quant_config.get("quant_method", quant_method)
|
|
quant_format = quant_config.get("format", quant_format)
|
|
quant_groups = quant_config.get("config_groups", quant_groups) or {}
|
|
quant_layers = quant_config.get("quantized_layers", quant_layers) or {}
|
|
if quant_algo in ("NVFP4", "W4A16_NVFP4"):
|
|
return True
|
|
if quant_method == "compressed-tensors" and (
|
|
quant_format == "nvfp4-pack-quantized"
|
|
or (
|
|
quant_format == "mixed-precision"
|
|
and any(
|
|
isinstance(group, dict) and group.get("format") == "nvfp4-pack-quantized"
|
|
for group in quant_groups.values()
|
|
)
|
|
)
|
|
):
|
|
return True
|
|
return any(
|
|
isinstance(entry, dict) and str(entry.get("quant_algo", "")).endswith("NVFP4")
|
|
for entry in quant_layers.values()
|
|
)
|
|
|
|
|
|
def _token_text(value) -> str | None:
|
|
if isinstance(value, str):
|
|
return value or None
|
|
if isinstance(value, dict):
|
|
content = value.get("content")
|
|
return content if isinstance(content, str) and content else None
|
|
return None
|
|
|
|
|
|
def _hf_stops(model_dir: Path) -> list[str]:
|
|
"""Stop strings Ollama should use for this tokenizer.
|
|
|
|
Llama 3's EOS list includes <|eom_id|>, which ends a tool-call turn inside
|
|
the chat template. Stopping on it clips the call. The working Llama 3.1
|
|
Modelfile stops on the header tokens and <|eot_id|> instead.
|
|
"""
|
|
config_path = model_dir / "config.json"
|
|
config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.is_file() else {}
|
|
tok_cfg: dict = {}
|
|
tokenizer_config = model_dir / "tokenizer_config.json"
|
|
if tokenizer_config.is_file():
|
|
tok_cfg = json.loads(tokenizer_config.read_text(encoding="utf-8"))
|
|
decoder = tok_cfg.get("added_tokens_decoder") or {}
|
|
known = set()
|
|
by_id = {}
|
|
for key, value in decoder.items():
|
|
try:
|
|
token = _token_text(value)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if token:
|
|
known.add(token)
|
|
try:
|
|
by_id[int(key)] = token
|
|
except (TypeError, ValueError):
|
|
continue
|
|
eos_text = _token_text(tok_cfg.get("eos_token"))
|
|
if eos_text:
|
|
known.add(eos_text)
|
|
if "<|eot_id|>" in known or "<|start_header_id|>" in known:
|
|
return [token for token in ("<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>") if token in known]
|
|
stops: list[str] = []
|
|
|
|
def add(token: str | None) -> None:
|
|
if token and token not in stops and token != "<|eom_id|>":
|
|
stops.append(token)
|
|
|
|
add(eos_text)
|
|
eos_ids = config.get("eos_token_id")
|
|
if isinstance(eos_ids, int):
|
|
eos_ids = [eos_ids]
|
|
for token_id in eos_ids or []:
|
|
try:
|
|
add(by_id.get(int(token_id)))
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if str(config.get("model_type") or "").lower().startswith("gemma"):
|
|
if "<end_of_turn>" in known:
|
|
stops = ["<end_of_turn>"] + [token for token in stops if token != "<end_of_turn>"]
|
|
return stops
|
|
|
|
|
|
def _hf_sampler(model_dir: Path) -> list[tuple[str, object]]:
|
|
path = model_dir / "generation_config.json"
|
|
if not path.is_file():
|
|
return []
|
|
config = json.loads(path.read_text(encoding="utf-8"))
|
|
chosen = []
|
|
for key in ("temperature", "top_p", "top_k", "min_p", "repeat_penalty"):
|
|
value = config.get(key)
|
|
if isinstance(value, bool) or value is None:
|
|
continue
|
|
if isinstance(value, (int, float)):
|
|
chosen.append((key, value))
|
|
return chosen
|
|
|
|
|
|
def _hf_chat_template(model_dir: Path) -> str:
|
|
"""The chat template llama.cpp will store, when the checkpoint has one."""
|
|
tok_cfg: dict = {}
|
|
tokenizer_config = model_dir / "tokenizer_config.json"
|
|
if tokenizer_config.is_file():
|
|
tok_cfg = json.loads(tokenizer_config.read_text(encoding="utf-8"))
|
|
if "chat_template" in tok_cfg and tok_cfg.get("chat_template") is not None:
|
|
return flatten_chat_template(tok_cfg.get("chat_template"))
|
|
jinja = model_dir / "chat_template.jinja"
|
|
if jinja.is_file():
|
|
return jinja.read_text(encoding="utf-8")
|
|
return ""
|
|
|
|
|
|
def _write_hf_modelfile(
|
|
model_dir: Path,
|
|
gguf_path: Path,
|
|
overwrite: bool,
|
|
mmprojs: list[Path] | None = None,
|
|
) -> Path:
|
|
"""Write a Modelfile next to the GGUF."""
|
|
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)
|
|
stops = _hf_stops(model_dir)
|
|
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.")
|
|
template = _hf_chat_template(model_dir)
|
|
lines.extend(tool_modelfile_lines(template))
|
|
if chat_template_has_tools(template):
|
|
log.info("Chat template supports tools.")
|
|
else:
|
|
log.info("Chat template has no tool round-trip. Tool requests stay rejected.")
|
|
for stop in stops:
|
|
lines.append(f"PARAMETER stop {_quote_modelfile_value(stop)}")
|
|
for key, value in _hf_sampler(model_dir):
|
|
rendered = str(value).lower() if isinstance(value, bool) else str(value)
|
|
lines.append(f"PARAMETER {key} {rendered}")
|
|
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
|
|
|
|
|
|
def _mtp_depth(model) -> tuple[int, int]:
|
|
n_blocks = int(getattr(model, "block_count", 0) or 0)
|
|
hparams = getattr(model, "hparams", {}) or {}
|
|
n_nextn = int(hparams.get("num_nextn_predict_layers") or 0)
|
|
hidden = hparams.get("num_hidden_layers")
|
|
if n_nextn <= 0 and hidden is not None and n_blocks > int(hidden):
|
|
n_nextn = n_blocks - int(hidden)
|
|
return n_blocks, n_nextn
|
|
|
|
|
|
def _fold_unclaimed_scale(name: str, architecture: str, n_blocks: int, n_nextn: int) -> bool:
|
|
"""True when this sidecar is a per-tensor NVFP4 scale the graph never multiplies."""
|
|
if not name.endswith(".scale"):
|
|
return False
|
|
weight = name[: -len(".scale")] + ".weight"
|
|
if preserve_scale_tensor(architecture, weight):
|
|
return False
|
|
if mtp_layer_tensor(weight, n_blocks, n_nextn):
|
|
return True
|
|
return must_bake_nvfp4_scale(architecture, weight)
|
|
|
|
|
|
def _fold_mtp_scales(model) -> None:
|
|
"""Bake NVFP4 scales the graph never multiplies, and do not write the sidecars."""
|
|
writer = model.gguf_writer
|
|
architecture = str(getattr(writer, "arch", "") or "")
|
|
n_blocks, n_nextn = _mtp_depth(model)
|
|
original = writer.add_tensor
|
|
pending: dict[str, tuple] = {}
|
|
|
|
def add_tensor(name, tensor, raw_shape=None, raw_dtype=None, tensor_endianess=None):
|
|
text = str(name)
|
|
if _fold_unclaimed_scale(text, architecture, n_blocks, n_nextn):
|
|
weight_name = text[: -len(".scale")] + ".weight"
|
|
held = pending.pop(weight_name, None)
|
|
if held is None:
|
|
return
|
|
scale = float(np.asarray(tensor).reshape(-1)[0])
|
|
original(weight_name, bake_packed_nvfp4(held[0], scale), raw_dtype=held[1])
|
|
return
|
|
scale_name = text[: -len(".weight")] + ".scale" if text.endswith(".weight") else ""
|
|
if getattr(raw_dtype, "name", "") == "NVFP4" and _fold_unclaimed_scale(scale_name, architecture, n_blocks, n_nextn):
|
|
pending[text] = (np.array(tensor, copy=True), raw_dtype)
|
|
return
|
|
original(name, tensor, raw_shape=raw_shape, raw_dtype=raw_dtype, tensor_endianess=tensor_endianess)
|
|
|
|
def write_header_to_file(path=None):
|
|
for held_name, (raw, dtype) in pending.items():
|
|
original(held_name, raw, raw_dtype=dtype)
|
|
pending.clear()
|
|
return original_header(path)
|
|
|
|
writer.add_tensor = add_tensor
|
|
original_header = writer.write_header_to_file
|
|
writer.write_header_to_file = write_header_to_file
|
|
|
|
|
|
def _normalize_chat_templates(model) -> None:
|
|
writer = model.gguf_writer
|
|
original = writer.add_chat_template
|
|
|
|
def add_chat_template(template, *args, **kwargs):
|
|
return original(normalize_chat_template(template), *args, **kwargs)
|
|
|
|
writer.add_chat_template = add_chat_template
|
|
|
|
|
|
def _warn_when_header_is_written(model) -> None:
|
|
writer = model.gguf_writer
|
|
original = writer.write_header_to_file
|
|
|
|
def write_header_to_file(path=None):
|
|
warn_missing_gguf_metadata(writer)
|
|
return original(path)
|
|
|
|
writer.write_header_to_file = write_header_to_file
|
|
|
|
|
|
def _assign_output(args) -> None:
|
|
"""Point args.output at the GGUF file inside the bundle directory."""
|
|
model_dir = Path(args.model)
|
|
if args.output is None:
|
|
if args.no_quant and args.outtype == "auto":
|
|
directory = Path.cwd() / model_dir.name
|
|
gguf_name = f"{model_dir.name}-{{FTYPE}}.gguf"
|
|
elif args.no_quant:
|
|
label = args.outtype.upper()
|
|
directory = Path.cwd() / f"{model_dir.name}-{label}"
|
|
gguf_name = f"{directory.name}.gguf"
|
|
else:
|
|
directory = Path.cwd() / f"{model_dir.name}-NVFP4"
|
|
gguf_name = f"{directory.name}.gguf"
|
|
else:
|
|
output = Path(args.output)
|
|
if output.suffix == ".gguf":
|
|
gguf_path = output.resolve()
|
|
gguf_path.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output = str(gguf_path)
|
|
return
|
|
directory = output
|
|
gguf_name = f"{directory.name}.gguf"
|
|
directory = directory.resolve()
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
args.output = str(directory / gguf_name)
|
|
|
|
|
|
def _refuse_existing_output(args) -> None:
|
|
refuse_existing_bundle(
|
|
Path(args.output),
|
|
args.overwrite,
|
|
getattr(args, "mmproj", None),
|
|
skip_gguf=bool(args.modelfile_only),
|
|
)
|
|
|
|
|
|
def export_float(args) -> None:
|
|
"""Convert the Hugging Face model to GGUF without learned rounding.
|
|
|
|
Packed 4-bit tensors are repacked as NVFP4. U8 is the packed weight and
|
|
F8_E4M3 is usually its block scale; a weight whose own dtype is F8_E4M3
|
|
is an 8-bit float and is stored as float. The shape has to be 8 weight
|
|
bytes per scale before a U8 tensor is treated as 4-bit. ``--outtype nvfp4``
|
|
asks for that repack. MOSTLY_NVFP4 is selected only after at least one
|
|
tensor is actually repacked.
|
|
"""
|
|
gguf, _model_base, model_type_enum, get_model_architecture, get_model_class = load_gguf_modules(args.llama_cpp)
|
|
model_dir = Path(args.model).resolve()
|
|
if not model_dir.is_dir():
|
|
raise FileNotFoundError(f"model directory does not exist: {model_dir}")
|
|
passthrough = args.outtype == "nvfp4"
|
|
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
log.info("Loading %s", model_dir)
|
|
hparams = _model_base.load_hparams(model_dir, False)
|
|
layout = _checkpoint_tensor_layout(model_dir)
|
|
_log_tensor_layout(layout)
|
|
config_nvfp4 = _checkpoint_is_nvfp4(model_dir, hparams)
|
|
if passthrough and not config_nvfp4 and not layout.packed:
|
|
if layout.fp8:
|
|
log.info(
|
|
"No packed 4-bit tensors were found. %d FP8 weights will be stored as float.",
|
|
len(layout.fp8),
|
|
)
|
|
else:
|
|
raise SystemExit(
|
|
f"{model_dir} has no packed 4-bit NVFP4 tensors. "
|
|
"Use --outtype f16, bf16, f32, or auto to write float weights, "
|
|
"or omit --outtype nvfp4 to fit NVFP4 from float weights."
|
|
)
|
|
architecture = get_model_architecture(hparams, model_type_enum.TEXT)
|
|
log.info("Architecture: %s", architecture)
|
|
model_cls = get_model_class(architecture)
|
|
file_type = gguf.LlamaFileType.MOSTLY_BF16 if args.outtype == "nvfp4" else _float_file_type(gguf, args.outtype)
|
|
model = model_cls(
|
|
model_dir,
|
|
file_type,
|
|
Path(args.output).resolve(),
|
|
use_temp_file=not args.no_temp_file,
|
|
eager=args.eager,
|
|
model_name=args.model_name,
|
|
)
|
|
quant_config = model.hparams.get("quantization_config")
|
|
if layout.fp8 and not layout.packed and not config_nvfp4 and not (isinstance(quant_config, dict) and quant_config.get("quant_method")):
|
|
merged = dict(quant_config) if isinstance(quant_config, dict) else {}
|
|
merged["quant_method"] = "modelopt"
|
|
model.hparams["quantization_config"] = merged
|
|
log.info("FP8 weights have no quant method in the config. Their scales will be applied and the weights stored as float.")
|
|
_install_tokenizer_fallback(model)
|
|
apply_hf_context_length(model.hparams, args.context_length)
|
|
_warn_when_header_is_written(model)
|
|
_normalize_chat_templates(model)
|
|
_fold_mtp_scales(model)
|
|
if layout.packed or config_nvfp4:
|
|
log.info("Repacking tensors that are packed 4-bit. Other weights are stored as float. Learned rounding is skipped.")
|
|
else:
|
|
stored = "bf16" if args.outtype == "nvfp4" else args.outtype
|
|
log.info("Writing float GGUF (%s). NVFP4 quantization is skipped.", stored)
|
|
model.write()
|
|
_write_hf_modelfile(model_dir, Path(model.fname_out), args.overwrite, args.mmproj)
|
|
log.info("Wrote %s", model.fname_out)
|
|
|
|
|
|
def export_nvfp4(args) -> None:
|
|
if args.modelfile_only:
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
model_dir = Path(args.model).resolve()
|
|
if not model_dir.is_dir():
|
|
raise FileNotFoundError(f"model directory does not exist: {model_dir}")
|
|
log.info("Skipping quantization and GGUF creation; writing Modelfile only")
|
|
_write_hf_modelfile(model_dir, Path(args.output), args.overwrite, args.mmproj)
|
|
return
|
|
_refuse_existing_output(args)
|
|
if args.no_quant or args.outtype == "nvfp4":
|
|
export_float(args)
|
|
return
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
ensure_cuda_interpreter(args.device)
|
|
gguf, _model_base, model_type_enum, get_model_architecture, get_model_class = load_gguf_modules(args.llama_cpp)
|
|
model_dir = Path(args.model).resolve()
|
|
if not model_dir.is_dir():
|
|
raise FileNotFoundError(f"model directory does not exist: {model_dir}")
|
|
|
|
args.optimizer = resolve_optimizer(args.optimizer)
|
|
args.output = str(Path(args.output).resolve())
|
|
log.info("Loading %s", model_dir)
|
|
hparams = _model_base.load_hparams(model_dir, False)
|
|
architecture = get_model_architecture(hparams, model_type_enum.TEXT)
|
|
log.info("Architecture: %s", architecture)
|
|
bind_exclude(args, architecture)
|
|
settings_path_parent = Path(args.checkpoint_dir) if args.checkpoint_dir else Path(args.output + ".nvfp4-work")
|
|
if args.fresh and settings_path_parent.exists():
|
|
shutil.rmtree(settings_path_parent)
|
|
settings_path_parent.mkdir(parents=True, exist_ok=True)
|
|
require_quant_settings(settings_path_parent, args)
|
|
model_cls = get_model_class(architecture)
|
|
output = Path(args.output).resolve()
|
|
model = model_cls(
|
|
model_dir,
|
|
gguf.LlamaFileType.MOSTLY_BF16,
|
|
output,
|
|
use_temp_file=not args.no_temp_file,
|
|
eager=args.eager,
|
|
model_name=args.model_name,
|
|
)
|
|
quantizer = TensorQuantizer(args, gguf, load_learned_converter(args.convert_to_quant)[1])
|
|
_install_tokenizer_fallback(model)
|
|
apply_hf_context_length(model.hparams, args.context_length)
|
|
_warn_when_header_is_written(model)
|
|
_normalize_chat_templates(model)
|
|
writer = Nvfp4Writer(model, quantizer, args, _collect_bias_names(model))
|
|
writer.install()
|
|
try:
|
|
log.info(
|
|
"llama.cpp prints BF16 for matrices before they are repacked. "
|
|
"Selected matrices are stored as NVFP4."
|
|
)
|
|
model.write()
|
|
except KeyboardInterrupt:
|
|
log.warning("Interrupted. Checkpoints are in %s; rerun the same command to resume.", writer.ckpt_dir)
|
|
raise
|
|
finally:
|
|
writer.restore()
|
|
|
|
if not args.keep_checkpoints and writer.quantized:
|
|
shutil.rmtree(writer.ckpt_dir, ignore_errors=True)
|
|
_write_hf_modelfile(model_dir, Path(model.fname_out), args.overwrite, args.mmproj)
|
|
log.info(
|
|
"Wrote %s (%d NVFP4 tensors, %d with scale folded into the blocks, %d biases adjusted, %d other tensors)",
|
|
model.fname_out,
|
|
writer.quantized,
|
|
writer.baked,
|
|
writer.corrected,
|
|
writer.passthrough,
|
|
)
|
|
|
|
|
|
def self_test(args) -> None:
|
|
"""Check that packed blocks dequantize to the same values llama.cpp will load."""
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
from nvfp4gguf import pack_nvfp4, ue4m3_bits
|
|
from nvfp4gguf.quant import check_scale_policy
|
|
check_scale_policy()
|
|
log.info("scale policy ok")
|
|
gguf, *_rest = load_gguf_modules(args.llama_cpp)
|
|
_converter_cls, unpack_uint4 = load_learned_converter(args.convert_to_quant)
|
|
|
|
codes = np.zeros((1, 64), dtype=np.uint8)
|
|
codes[0, :16] = 7 # +6
|
|
codes[0, 16:32] = 15 # -6
|
|
scale_float = np.array([[1.0, 0.5, 256.0, 448.0]], dtype=np.float32)
|
|
bits = ue4m3_bits(scale_float)
|
|
raw = pack_nvfp4(codes, bits)
|
|
decoded = gguf.quants.dequantize(raw, gguf.GGMLQuantizationType.NVFP4)
|
|
expected = np.zeros((1, 64), dtype=np.float32)
|
|
expected[0, :16] = 6.0 * scale_float[0, 0]
|
|
expected[0, 16:32] = -6.0 * scale_float[0, 1]
|
|
if not np.allclose(decoded, expected, rtol=1e-5, atol=1e-5):
|
|
raise AssertionError(f"manual NVFP4 roundtrip failed\n{decoded}\n{expected}")
|
|
log.info("manual block roundtrip ok")
|
|
|
|
def check(quantizer, label: str, shapes: tuple[tuple[int, int], ...]) -> None:
|
|
generator = np.random.default_rng(0)
|
|
for shape in shapes:
|
|
weight = generator.normal(size=shape).astype(np.float32) * 0.05
|
|
codes, block_scale, per_tensor, dequant = quantizer.quantize_matrix(weight)
|
|
runtime = quantizer.runtime_dequant(pack_nvfp4(codes, ue4m3_bits(block_scale)), per_tensor)
|
|
delta = float(np.max(np.abs(runtime - dequant)))
|
|
if delta > 1e-4:
|
|
raise AssertionError(f"{label} {shape} ggml dequant disagrees with the quantizer by {delta}")
|
|
log.info("%s %s roundtrip ok (max abs %.3e)", label, shape, delta)
|
|
baked_codes, baked_bits = quantizer.bake_global_scale(codes, block_scale, per_tensor)
|
|
baked = quantizer.runtime_dequant(pack_nvfp4(baked_codes, baked_bits), None)
|
|
rel = float(np.max(np.abs(baked - dequant))) / (float(np.max(np.abs(dequant))) + 1e-8)
|
|
if rel > 0.15:
|
|
raise AssertionError(f"{label} {shape} baked scale moved the reconstruction by {rel}")
|
|
log.info("%s %s baked-scale relative error %.3e", label, shape, rel)
|
|
|
|
def configured(**overrides):
|
|
values = vars(args).copy()
|
|
values.update(overrides)
|
|
return argparse.Namespace(**values)
|
|
|
|
check(TensorQuantizer(configured(simple=True, num_iter=1, device="cpu", optimizer="adamw"), gguf, unpack_uint4),
|
|
"simple", ((64, 128), (40, 64)))
|
|
check(TensorQuantizer(configured(
|
|
simple=False, num_iter=4, device="cpu", optimizer="original", scale_optimization="fixed",
|
|
), gguf, unpack_uint4), "learned", ((64, 128),))
|
|
log.info("self-test passed")
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Quantize a Hugging Face model to NVFP4 GGUF with convert_to_quant's learned E2M1 rounding.",
|
|
)
|
|
parser.add_argument("model", nargs="?", help="Hugging Face model directory (config.json + safetensors)")
|
|
parser.add_argument("-o", "--output", help="Output directory. Default: ./<model-name>-NVFP4, or ./<model-name> with --no-quant and auto. A path ending in .gguf is the GGUF file itself.")
|
|
parser.add_argument("--modelfile-only", action="store_true", help="Write the Modelfile and skip quantization and GGUF creation")
|
|
parser.add_argument("--overwrite", action="store_true", help="Replace an existing GGUF, Modelfile, or copied projector")
|
|
parser.add_argument("--mmproj", action="append", type=Path, default=None,
|
|
help="Vision projector GGUF. Copied into the output directory under its own name and added as another FROM line. Repeat for more than one projector.")
|
|
parser.add_argument("--context-length", type=int, default=None,
|
|
help="Max context stored in the GGUF. Replaces the checkpoint value, or sets one when the checkpoint has no single context length.")
|
|
parser.add_argument("--no-quant", action="store_true",
|
|
help="Skip learned rounding. U8 tensors are repacked as NVFP4 only when their shape is packed 4-bit. FP8 weights are stored as float.")
|
|
parser.add_argument("--outtype", default="auto", type=_outtype, choices=("auto", "f32", "f16", "bf16", "nvfp4"),
|
|
help="Float type for tensors that are not packed 4-bit. nvfp4 repacks a checkpoint that already has packed 4-bit weights. auto matches the source float weights (default: auto)")
|
|
parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP,
|
|
help="llama.cpp checkout used for architecture conversion and the GGUF writer (default: ~/llama.cpp)")
|
|
parser.add_argument("--convert-to-quant", type=Path, default=None,
|
|
help="convert_to_quant checkout, if it is not installed in the environment")
|
|
parser.add_argument("--model-name", default=None, help="Override the model name stored in the GGUF")
|
|
parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or cuda:N")
|
|
parser.add_argument("--simple", action="store_true", help="Round each value onto the E2M1 grid without learned rounding")
|
|
parser.add_argument("--optimizer", default="auto", choices=("auto", "prodigy", "adamw", "radam", "original"),
|
|
help="Learned-rounding optimizer. 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-patience", type=int, default=5, help="Patience for learning rate decay (default: 5)")
|
|
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("--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-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("--calib-samples", type=int, default=3072, help="Random inputs used to correct a sibling bias")
|
|
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 the output (lm_head) matrix")
|
|
parser.add_argument("--exclude", default=None, help="Regex of GGUF tensor names to leave in BF16")
|
|
parser.add_argument("--checkpoint-dir", type=Path, default=None, help="Per-tensor resume directory")
|
|
parser.add_argument("--fresh", action="store_true", help="Delete the checkpoint directory before starting")
|
|
parser.add_argument("--keep-checkpoints", action="store_true")
|
|
parser.add_argument("--eager", action="store_true", help="Load every tensor up front (uses more RAM)")
|
|
parser.add_argument("--no-temp-file", action="store_true", help="Hold the finished GGUF tensors in RAM instead of a temp file")
|
|
parser.add_argument("--verbose", action="store_true")
|
|
parser.add_argument("--self-test", action="store_true", help="Check NVFP4 packing against llama.cpp's dequantizer and exit")
|
|
args = parser.parse_args(argv)
|
|
if args.self_test:
|
|
if args.device == "auto":
|
|
args.device = "cpu"
|
|
return args
|
|
if not args.model:
|
|
parser.error("model directory is required")
|
|
_assign_output(args)
|
|
_refuse_existing_output(args)
|
|
args.exclude_re = re.compile(args.exclude) if args.exclude else None
|
|
if args.convert_to_quant is None and DEFAULT_CONVERT_TO_QUANT.is_dir():
|
|
args.convert_to_quant = DEFAULT_CONVERT_TO_QUANT
|
|
return args
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
if args.self_test:
|
|
self_test(args)
|
|
return
|
|
export_nvfp4(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|