669 lines
29 KiB
Python
669 lines
29 KiB
Python
#!/usr/bin/env python3
|
|
"""Quantize a GGUF to llama.cpp NVFP4, or repack one that is already NVFP4.
|
|
|
|
The learned-rounding flags match hf_to_nvfp4_gguf.py. Vocabulary, tensor names,
|
|
and architecture metadata are copied from the source GGUF. --outtype nvfp4
|
|
copies an existing NVFP4 file without fitting it again. --no-quant with
|
|
f32, f16, bf16, or auto rewrites the tensors as float, or leaves them as stored.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import logging
|
|
import re
|
|
import shutil
|
|
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,
|
|
import_gguf,
|
|
load_learned_converter,
|
|
must_bake_nvfp4_scale,
|
|
mtp_layer_tensor,
|
|
bake_packed_nvfp4,
|
|
preserve_scale_tensor,
|
|
quantize_and_pack,
|
|
report_reconstruction,
|
|
require_quant_settings,
|
|
resolve_optimizer,
|
|
scale_bake_notice,
|
|
)
|
|
from nvfp4gguf.bundle import (
|
|
flatten_chat_template,
|
|
normalize_chat_template,
|
|
is_model_context_key,
|
|
outtype as _outtype,
|
|
refuse_existing_bundle,
|
|
require_context_length,
|
|
scalar_context_length,
|
|
warn_missing_gguf_metadata,
|
|
write_simple_modelfile,
|
|
)
|
|
|
|
log = logging.getLogger("gguf_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 _pythonize(value):
|
|
if isinstance(value, np.generic):
|
|
return value.item()
|
|
if isinstance(value, np.ndarray):
|
|
return value.tolist()
|
|
if isinstance(value, list):
|
|
return [_pythonize(item) for item in value]
|
|
return value
|
|
|
|
|
|
def _text(value) -> str:
|
|
value = _pythonize(value)
|
|
if isinstance(value, bytes):
|
|
return value.decode("utf-8", "replace")
|
|
if isinstance(value, list):
|
|
return "".join(_text(item) for item in value)
|
|
return str(value)
|
|
|
|
|
|
def _logical_shape(tensor) -> tuple[int, ...]:
|
|
return tuple(int(dim) for dim in tensor.shape[::-1])
|
|
|
|
|
|
def _to_float(tensor, gguf_module) -> np.ndarray:
|
|
qtype = tensor.tensor_type
|
|
if qtype == gguf_module.GGMLQuantizationType.F32:
|
|
values = np.asarray(tensor.data, dtype=np.float32)
|
|
elif qtype == gguf_module.GGMLQuantizationType.F16:
|
|
values = np.asarray(tensor.data, dtype=np.float32)
|
|
else:
|
|
values = np.asarray(gguf_module.quants.dequantize(tensor.data, qtype), dtype=np.float32)
|
|
# The reader maps the file read-only. Torch refuses a non-writable view.
|
|
return np.array(values, dtype=np.float32, copy=True)
|
|
|
|
|
|
def _open_reader(gguf, path: Path):
|
|
if not path.is_file():
|
|
raise SystemExit(f"GGUF file does not exist: {path}")
|
|
reader = gguf.GGUFReader(path)
|
|
if reader.byte_order != "I":
|
|
raise SystemExit(f"{path} is not little-endian GGUF")
|
|
return reader
|
|
|
|
|
|
def _architecture(reader) -> str:
|
|
field = reader.fields.get("general.architecture")
|
|
if field is None:
|
|
return "unknown"
|
|
return _text(field.contents())
|
|
|
|
|
|
def _file_type_name(reader, gguf_module) -> str:
|
|
field = reader.fields.get("general.file_type")
|
|
if field is None:
|
|
return ""
|
|
try:
|
|
return gguf_module.LlamaFileType(int(field.contents())).name
|
|
except (TypeError, ValueError):
|
|
return str(field.contents())
|
|
|
|
|
|
def _is_nvfp4(reader, gguf_module) -> bool:
|
|
if _file_type_name(reader, gguf_module) == "MOSTLY_NVFP4":
|
|
return True
|
|
return any(tensor.tensor_type == gguf_module.GGMLQuantizationType.NVFP4 for tensor in reader.tensors)
|
|
|
|
|
|
def _copy_metadata(reader, writer, gguf_module, file_type: int | None, model_name: str | None, context_length: int | None) -> None:
|
|
requested = require_context_length(context_length)
|
|
source_context = None
|
|
wrote_context = False
|
|
for field in reader.fields.values():
|
|
if field.name.startswith("GGUF.") or field.name == "general.architecture" or not field.types:
|
|
continue
|
|
value = _pythonize(field.contents())
|
|
if value is None:
|
|
continue
|
|
value_type = field.types[0]
|
|
if value_type == gguf_module.GGUFValueType.ARRAY and not value:
|
|
continue
|
|
if field.name == "general.file_type" and file_type is not None:
|
|
value = int(file_type)
|
|
if field.name == "general.name" and model_name:
|
|
value = model_name
|
|
if field.name == "general.alignment":
|
|
writer.data_alignment = int(value)
|
|
subtype = field.types[-1] if value_type == gguf_module.GGUFValueType.ARRAY else None
|
|
if field.name == "tokenizer.chat_template" or field.name.startswith("tokenizer.chat_template."):
|
|
value = normalize_chat_template(value)
|
|
if is_model_context_key(field.name):
|
|
source_context = scalar_context_length(value)
|
|
if requested is not None:
|
|
value = requested
|
|
value_type = gguf_module.GGUFValueType.UINT32
|
|
subtype = None
|
|
wrote_context = True
|
|
writer.add_key_value(field.name, value, value_type, sub_type=subtype)
|
|
if file_type is not None and "general.file_type" not in reader.fields:
|
|
writer.add_file_type(int(file_type))
|
|
if model_name and "general.name" not in reader.fields:
|
|
writer.add_name(model_name)
|
|
if requested is not None and not wrote_context:
|
|
writer.add_context_length(requested)
|
|
if requested is not None:
|
|
log.info("context length %d", requested)
|
|
elif source_context is not None:
|
|
log.info("context length %d from the source GGUF", source_context)
|
|
|
|
|
|
def _replaced_scale(name: str, quant_names: set[str], architecture: str) -> bool:
|
|
if not (name.endswith(".scale") or name.endswith(".input_scale")):
|
|
return False
|
|
stem = name.rsplit(".", 1)[0]
|
|
weight = stem + ".weight"
|
|
if weight not in quant_names:
|
|
return False
|
|
return not (name.endswith(".scale") and preserve_scale_tensor(architecture, weight))
|
|
|
|
|
|
def _keep_f32(name: str, rank: int) -> bool:
|
|
return rank < 2 or "norm" in name.lower()
|
|
|
|
|
|
def _add_float(writer, gguf_module, name: str, values: np.ndarray, outtype: str) -> None:
|
|
values = np.ascontiguousarray(values, dtype=np.float32)
|
|
if outtype == "f16":
|
|
writer.add_tensor(name, values.astype(np.float16))
|
|
return
|
|
if outtype == "bf16":
|
|
raw = gguf_module.quants.quantize(values, gguf_module.GGMLQuantizationType.BF16)
|
|
writer.add_tensor(name, raw, raw_dtype=gguf_module.GGMLQuantizationType.BF16)
|
|
return
|
|
writer.add_tensor(name, values)
|
|
|
|
|
|
def _copy_tensor(writer, tensor) -> None:
|
|
writer.add_tensor(tensor.name, tensor.data, raw_dtype=tensor.tensor_type)
|
|
|
|
|
|
def _load_checkpoint(directory: Path, name: str, shape: tuple[int, ...]):
|
|
path = _checkpoint_path(directory, 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(shape):
|
|
return None
|
|
raw = np.array(blob["raw"], copy=True)
|
|
scale = np.array(blob["scale"], copy=True)
|
|
bias = np.array(blob["bias"], copy=True)
|
|
return raw, None if scale.size == 0 else scale, None if bias.size == 0 else bias
|
|
|
|
|
|
def _save_checkpoint(directory: Path, name: str, shape: tuple[int, ...], raw, scale, bias) -> None:
|
|
path = _checkpoint_path(directory, name)
|
|
temporary = path.with_name(path.name + ".tmp")
|
|
with temporary.open("wb") as handle:
|
|
np.savez_compressed(
|
|
handle,
|
|
version=np.int32(QUANT_VERSION),
|
|
shape=np.asarray(shape, dtype=np.int64),
|
|
raw=raw,
|
|
scale=np.zeros((0,), np.float32) if scale is None else np.asarray(scale, dtype=np.float32),
|
|
bias=np.zeros((0,), np.float32) if bias is None else np.asarray(bias, dtype=np.float32),
|
|
)
|
|
temporary.replace(path)
|
|
|
|
|
|
def _assign_output(args) -> None:
|
|
source = Path(args.gguf)
|
|
stem = source.stem
|
|
passthrough = args.no_quant or args.outtype == "nvfp4"
|
|
if args.output is None:
|
|
if passthrough and args.outtype == "auto":
|
|
directory = Path.cwd() / stem
|
|
gguf_name = f"{stem}.gguf"
|
|
elif passthrough:
|
|
label = "NVFP4" if args.outtype == "nvfp4" else args.outtype.upper()
|
|
directory = Path.cwd() / f"{stem}-{label}"
|
|
gguf_name = f"{directory.name}.gguf"
|
|
else:
|
|
directory = Path.cwd() / f"{stem}-NVFP4"
|
|
gguf_name = f"{directory.name}.gguf"
|
|
else:
|
|
output = Path(args.output)
|
|
if output.suffix == ".gguf":
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output = str(output.resolve())
|
|
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_same_file(args) -> None:
|
|
if Path(args.output).resolve() == Path(args.gguf).resolve():
|
|
raise SystemExit("Output GGUF is the same file as the input. Choose a different --output.")
|
|
|
|
|
|
def _target_file_type(gguf_module, outtype: str, source_name: str):
|
|
if outtype == "nvfp4":
|
|
return gguf_module.LlamaFileType.MOSTLY_NVFP4
|
|
if outtype == "f32":
|
|
return gguf_module.LlamaFileType.ALL_F32
|
|
if outtype == "f16":
|
|
return gguf_module.LlamaFileType.MOSTLY_F16
|
|
if outtype == "bf16":
|
|
return gguf_module.LlamaFileType.MOSTLY_BF16
|
|
if source_name == "ALL_F32":
|
|
return gguf_module.LlamaFileType.ALL_F32
|
|
if source_name == "MOSTLY_F16":
|
|
return gguf_module.LlamaFileType.MOSTLY_F16
|
|
if source_name == "MOSTLY_BF16":
|
|
return gguf_module.LlamaFileType.MOSTLY_BF16
|
|
if source_name == "MOSTLY_NVFP4":
|
|
return gguf_module.LlamaFileType.MOSTLY_NVFP4
|
|
return None
|
|
|
|
|
|
def _chat_template(reader) -> str:
|
|
parts = []
|
|
for name, field in reader.fields.items():
|
|
if name == "tokenizer.chat_template" or name.startswith("tokenizer.chat_template."):
|
|
parts.append(flatten_chat_template(field.contents()))
|
|
return "\n".join(part for part in parts if part)
|
|
|
|
|
|
def _kv_int(reader, key: str) -> int:
|
|
field = reader.fields.get(key)
|
|
if field is None:
|
|
return 0
|
|
try:
|
|
return int(field.contents())
|
|
except (TypeError, ValueError):
|
|
return 0
|
|
|
|
|
|
def _mtp_counts(reader, architecture: str) -> tuple[int, int]:
|
|
return _kv_int(reader, f"{architecture}.block_count"), _kv_int(reader, f"{architecture}.nextn_predict_layers")
|
|
|
|
|
|
def _folded_scales(reader, gguf_module, architecture: str, n_blocks: int, n_nextn: int) -> dict[str, float]:
|
|
"""Per-tensor NVFP4 scales the llama.cpp graph never multiplies."""
|
|
by_name = {tensor.name: tensor for tensor in reader.tensors}
|
|
nvfp4 = gguf_module.GGMLQuantizationType.NVFP4
|
|
scales: dict[str, float] = {}
|
|
for tensor in reader.tensors:
|
|
if not tensor.name.endswith(".scale"):
|
|
continue
|
|
weight_name = tensor.name[: -len(".scale")] + ".weight"
|
|
weight = by_name.get(weight_name)
|
|
if weight is None or weight.tensor_type != nvfp4:
|
|
continue
|
|
if preserve_scale_tensor(architecture, weight_name):
|
|
continue
|
|
bias_name = weight_name[: -len(".weight")] + ".bias"
|
|
fold = mtp_layer_tensor(weight_name, n_blocks, n_nextn) or must_bake_nvfp4_scale(
|
|
architecture, weight_name, bias_name in by_name
|
|
)
|
|
if not fold:
|
|
continue
|
|
values = np.asarray(tensor.data, dtype=np.float32).reshape(-1)
|
|
if values.size != 1:
|
|
log.warning(
|
|
"Leaving %s in place. Baking expects one per-tensor scale, got %d values.",
|
|
tensor.name,
|
|
int(values.size),
|
|
)
|
|
continue
|
|
scales[tensor.name] = float(values[0])
|
|
return scales
|
|
|
|
|
|
def export_copy(args) -> None:
|
|
"""Rewrite the GGUF without learned rounding."""
|
|
gguf = import_gguf(args.llama_cpp)
|
|
source = Path(args.gguf).resolve()
|
|
output = Path(args.output).resolve()
|
|
reader = _open_reader(gguf, source)
|
|
if args.outtype == "nvfp4" and not _is_nvfp4(reader, gguf):
|
|
raise SystemExit(
|
|
f"{source} is not an NVFP4 GGUF. "
|
|
"Use --outtype f16, bf16, f32, or auto to write float weights, "
|
|
"or omit --outtype nvfp4 to fit NVFP4 from these weights."
|
|
)
|
|
architecture = _architecture(reader)
|
|
source_type = _file_type_name(reader, gguf)
|
|
file_type = _target_file_type(gguf, args.outtype, source_type)
|
|
if args.outtype == "nvfp4":
|
|
log.info("Copying existing NVFP4 tensors. Learned rounding is skipped.")
|
|
elif args.outtype == "auto":
|
|
log.info("Copying tensors as stored (%s). Learned rounding is skipped.", source_type or "unknown file type")
|
|
else:
|
|
log.info("Writing float GGUF (%s). Learned rounding is skipped.", args.outtype)
|
|
writer = gguf.GGUFWriter(output, architecture, use_temp_file=not args.no_temp_file)
|
|
_copy_metadata(reader, writer, gguf, None if file_type is None else int(file_type), args.model_name, args.context_length)
|
|
n_blocks, n_nextn = _mtp_counts(reader, architecture)
|
|
folded = _folded_scales(reader, gguf, architecture, n_blocks, n_nextn) if args.outtype in ("auto", "nvfp4") else {}
|
|
if folded:
|
|
log.info(
|
|
"Folding %d scale tensors into their NVFP4 blocks. "
|
|
"llama.cpp does not claim those sidecars.",
|
|
len(folded),
|
|
)
|
|
for tensor in reader.tensors:
|
|
if tensor.name in folded:
|
|
continue
|
|
scale_name = tensor.name[: -len(".weight")] + ".scale" if tensor.name.endswith(".weight") else ""
|
|
if scale_name in folded and tensor.tensor_type == gguf.GGMLQuantizationType.NVFP4:
|
|
baked = bake_packed_nvfp4(np.array(tensor.data, copy=True), folded[scale_name])
|
|
writer.add_tensor(tensor.name, baked, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
|
|
continue
|
|
if args.outtype in ("auto", "nvfp4"):
|
|
_copy_tensor(writer, tensor)
|
|
continue
|
|
if _keep_f32(tensor.name, len(_logical_shape(tensor))):
|
|
if tensor.tensor_type == gguf.GGMLQuantizationType.F32:
|
|
_copy_tensor(writer, tensor)
|
|
else:
|
|
_add_float(writer, gguf, tensor.name, _to_float(tensor, gguf), "f32")
|
|
continue
|
|
_add_float(writer, gguf, tensor.name, _to_float(tensor, gguf), args.outtype)
|
|
log.info("Writing %s", output)
|
|
warn_missing_gguf_metadata(writer)
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file(progress=True)
|
|
writer.close()
|
|
write_simple_modelfile(output, args.overwrite, args.mmproj, _chat_template(reader))
|
|
log.info("Wrote %s", output)
|
|
|
|
|
|
def _emit_bias(writer, gguf_module, tensor, correction: np.ndarray | None) -> bool:
|
|
if correction is None:
|
|
_copy_tensor(writer, tensor)
|
|
return False
|
|
values = _to_float(tensor, gguf_module)
|
|
flat = values.reshape(-1)
|
|
if flat.shape != correction.shape:
|
|
log.warning(" bias %s shape %s does not match correction %s", tensor.name, values.shape, correction.shape)
|
|
_copy_tensor(writer, tensor)
|
|
return False
|
|
values = (flat - correction).reshape(values.shape)
|
|
log.info(" bias %s adjusted (mean %+.3e)", tensor.name, float(correction.mean()))
|
|
if tensor.tensor_type == gguf_module.GGMLQuantizationType.F16:
|
|
writer.add_tensor(tensor.name, np.ascontiguousarray(values, dtype=np.float32).astype(np.float16))
|
|
return True
|
|
if tensor.tensor_type == gguf_module.GGMLQuantizationType.BF16:
|
|
raw = gguf_module.quants.quantize(
|
|
np.ascontiguousarray(values, dtype=np.float32),
|
|
gguf_module.GGMLQuantizationType.BF16,
|
|
)
|
|
writer.add_tensor(tensor.name, raw, raw_dtype=gguf_module.GGMLQuantizationType.BF16)
|
|
return True
|
|
writer.add_tensor(tensor.name, np.ascontiguousarray(values, dtype=np.float32))
|
|
return True
|
|
|
|
|
|
def export_quant(args) -> None:
|
|
gguf = import_gguf(args.llama_cpp)
|
|
source = Path(args.gguf).resolve()
|
|
output = Path(args.output).resolve()
|
|
ensure_cuda_interpreter(args.device)
|
|
reader = _open_reader(gguf, source)
|
|
architecture = _architecture(reader)
|
|
bind_exclude(args, architecture)
|
|
notice = scale_bake_notice(architecture)
|
|
if notice:
|
|
log.warning(notice)
|
|
file_type_name = _file_type_name(reader, gguf)
|
|
if file_type_name not in {"ALL_F32", "MOSTLY_F16", "MOSTLY_BF16"}:
|
|
log.warning(
|
|
"Source file type is %s. NVFP4 will be fit to the dequantized weights, "
|
|
"so detail already discarded by the source quant stays discarded.",
|
|
file_type_name or "quantized",
|
|
)
|
|
names = {tensor.name for tensor in reader.tensors}
|
|
n_blocks, n_nextn = _mtp_counts(reader, architecture)
|
|
quant_names = [
|
|
tensor.name
|
|
for tensor in reader.tensors
|
|
if eligible_weight(tensor.name, _logical_shape(tensor), args, gguf)
|
|
]
|
|
quant_name_set = set(quant_names)
|
|
log.info("%d matrices queued for NVFP4, %d tensors copied", len(quant_names), len(reader.tensors) - len(quant_names))
|
|
if not quant_names:
|
|
log.warning("No matrices are eligible for NVFP4. Tensors are copied and the file type stays unchanged.")
|
|
|
|
work = Path(args.checkpoint_dir) if args.checkpoint_dir else Path(str(output) + ".nvfp4-work")
|
|
if args.fresh and work.exists():
|
|
shutil.rmtree(work)
|
|
work.mkdir(parents=True, exist_ok=True)
|
|
require_quant_settings(work, args)
|
|
|
|
_converter_cls, unpack_uint4 = load_learned_converter(args.convert_to_quant)
|
|
del _converter_cls
|
|
quantizer = TensorQuantizer(args, gguf, unpack_uint4)
|
|
writer = gguf.GGUFWriter(output, architecture, use_temp_file=not args.no_temp_file)
|
|
_copy_metadata(
|
|
reader, writer, gguf,
|
|
int(gguf.LlamaFileType.MOSTLY_NVFP4) if quant_names else None,
|
|
args.model_name,
|
|
args.context_length,
|
|
)
|
|
|
|
pending: dict[str, object] = {}
|
|
corrections: dict[str, np.ndarray] = {}
|
|
done = 0
|
|
baked = 0
|
|
corrected = 0
|
|
try:
|
|
for tensor in reader.tensors:
|
|
if _replaced_scale(tensor.name, quant_name_set, architecture):
|
|
continue
|
|
weight_name = tensor.name[: -len(".bias")] + ".weight" if tensor.name.endswith(".bias") else ""
|
|
if weight_name in quant_name_set and not args.no_bias_correction:
|
|
if tensor.name not in corrections:
|
|
pending[tensor.name] = tensor
|
|
continue
|
|
if _emit_bias(writer, gguf, tensor, corrections[tensor.name]):
|
|
corrected += 1
|
|
continue
|
|
if tensor.name not in quant_name_set:
|
|
_copy_tensor(writer, tensor)
|
|
continue
|
|
done += 1
|
|
logical = _logical_shape(tensor)
|
|
bias_name = tensor.name[: -len(".weight")] + ".bias"
|
|
bake = must_bake_nvfp4_scale(architecture, tensor.name, bias_name in names) or mtp_layer_tensor(
|
|
tensor.name, n_blocks, n_nextn
|
|
)
|
|
cached = _load_checkpoint(work, tensor.name, logical)
|
|
if cached is not None and (cached[1] is None) != bake:
|
|
log.info(" checkpoint scale layout does not match %s; recomputing", tensor.name)
|
|
cached = None
|
|
log.info(
|
|
"NVFP4 [%d/%d] %s %s from %s%s",
|
|
done, len(quant_names), tensor.name, logical, tensor.tensor_type.name,
|
|
" (checkpoint)" if cached else "",
|
|
)
|
|
if cached is None:
|
|
weight = _to_float(tensor, gguf)
|
|
raw, scale = quantize_and_pack(quantizer, weight, bake, tensor.name)
|
|
report_reconstruction(quantizer, weight, raw, scale)
|
|
bias = None
|
|
if weight.ndim == 2 and not args.no_bias_correction and bias_name in names:
|
|
per_tensor = None if scale is None else float(np.reshape(scale, -1)[0])
|
|
runtime = quantizer.runtime_dequant(raw if raw.ndim == 2 else raw[0], per_tensor)
|
|
sample = weight if weight.ndim == 2 else weight[0]
|
|
bias = quantizer.bias_correction(sample, runtime)
|
|
_save_checkpoint(work, tensor.name, logical, raw, scale, bias)
|
|
del weight
|
|
else:
|
|
raw, scale, bias = cached
|
|
if bias is not None:
|
|
corrections[bias_name] = bias
|
|
pending_tensor = pending.pop(bias_name, None)
|
|
if pending_tensor is not None and _emit_bias(writer, gguf, pending_tensor, bias):
|
|
corrected += 1
|
|
if bake and scale is None:
|
|
baked += 1
|
|
writer.add_tensor(tensor.name, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
|
|
if scale is not None:
|
|
writer.add_tensor(tensor.name[: -len(".weight")] + ".scale", np.ascontiguousarray(scale, dtype=np.float32))
|
|
del raw
|
|
for name, tensor in list(pending.items()):
|
|
if _emit_bias(writer, gguf, tensor, corrections.get(name)):
|
|
corrected += 1
|
|
pending.clear()
|
|
except KeyboardInterrupt:
|
|
log.warning("Interrupted. Checkpoints are in %s; rerun the same command to resume.", work)
|
|
raise
|
|
|
|
log.info("Writing %s", output)
|
|
warn_missing_gguf_metadata(writer)
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file(progress=True)
|
|
writer.close()
|
|
if not args.keep_checkpoints:
|
|
shutil.rmtree(work, ignore_errors=True)
|
|
write_simple_modelfile(output, args.overwrite, args.mmproj, _chat_template(reader))
|
|
log.info(
|
|
"Wrote %s (%d NVFP4 tensors, %d with scale folded into the blocks, %d biases adjusted)",
|
|
output, done, baked, corrected,
|
|
)
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Quantize a GGUF to NVFP4, or repack one that is already NVFP4, with the same options as hf_to_nvfp4_gguf.py.",
|
|
)
|
|
parser.add_argument("gguf", nargs="?", help="Input GGUF")
|
|
parser.add_argument("-o", "--output", help="Output directory. Default: ./<name>-NVFP4, or ./<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 a source value, or sets one when the source has no single context length.")
|
|
parser.add_argument("--no-quant", action="store_true",
|
|
help="Skip learned rounding. Writes a float GGUF, or copies an existing NVFP4 GGUF with --outtype nvfp4")
|
|
parser.add_argument("--outtype", default="auto", type=_outtype, choices=("auto", "f32", "f16", "bf16", "nvfp4"),
|
|
help="Type used when learned rounding is skipped. nvfp4 copies a GGUF that is already NVFP4. auto leaves tensors as stored (default: auto)")
|
|
parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP,
|
|
help="llama.cpp checkout used for the GGUF reader and 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-min", type=float, default=1e-8, help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).")
|
|
parser.add_argument("--lr-patience", type=int, default=5, help="Patience for learning rate decay (default: 5)")
|
|
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 their source type")
|
|
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="Accepted for the same command line as hf_to_nvfp4_gguf. The GGUF reader maps the file either way.")
|
|
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.gguf:
|
|
parser.error("GGUF file is required")
|
|
_assign_output(args)
|
|
_refuse_same_file(args)
|
|
if not args.modelfile_only:
|
|
refuse_existing_bundle(Path(args.output), args.overwrite, args.mmproj)
|
|
else:
|
|
refuse_existing_bundle(Path(args.output), args.overwrite, args.mmproj, skip_gguf=True)
|
|
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:
|
|
from hf_to_nvfp4_gguf import self_test
|
|
self_test(args)
|
|
return
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
if args.modelfile_only:
|
|
output = Path(args.output)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
log.info("Skipping quantization and GGUF creation; writing Modelfile only")
|
|
reader = _open_reader(import_gguf(args.llama_cpp), Path(args.gguf))
|
|
write_simple_modelfile(output, args.overwrite, args.mmproj, _chat_template(reader))
|
|
return
|
|
if args.no_quant or args.outtype == "nvfp4":
|
|
export_copy(args)
|
|
return
|
|
args.optimizer = resolve_optimizer(args.optimizer)
|
|
export_quant(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|