commit db7f28d087fc4777fdbcc8593553466b8b6511aa Author: zefie Date: Sun Nov 1 16:31:33 2026 -0500 initial commit diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..7e8ac45 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,20 @@ +FROM runpod/pytorch:1.5.0-cu1320-torch2130-ubuntu2404 +WORKDIR /root + +RUN mkdir llama.cpp +COPY llama.cpp ./llama.cpp/ +RUN git clone https://github.com/silveroxides/convert_to_quant && cd convert_to_quant && pip install -r requirements.txt && pip install . + +RUN mkdir nvfp4_gguf_quant +WORKDIR /root/nvfp4_gguf_quant +COPY --exclude=__pycache__ nvfp4gguf ./nvfp4gguf +COPY --chmod=0644 gguf_to_nvfp4_gguf.py hf_to_nvfp4_gguf.py hf_to_nvfp4_safetensors.py requirements.txt safetensors_to_nvfp4_gguf.py . + +RUN pip install -r requirements.txt +COPY --chmod=0644 docker/README.txt /workspace/README.txt +COPY --chmod=0755 docker/hf_to_nvfp4_gguf docker/hf_to_nvfp4_safetensors docker/gguf_to_nvfp4_gguf docker/safetensors_to_nvfp4_gguf /usr/local/bin/ + +RUN apt remove -y python3-cryptography && pip install heretic-llm + +WORKDIR /workspace + diff --git a/__pycache__/gguf_info.cpython-313.pyc b/__pycache__/gguf_info.cpython-313.pyc new file mode 100644 index 0000000..1bdc4c0 Binary files /dev/null and b/__pycache__/gguf_info.cpython-313.pyc differ diff --git a/__pycache__/gguf_to_nvfp4_gguf.cpython-313.pyc b/__pycache__/gguf_to_nvfp4_gguf.cpython-313.pyc new file mode 100644 index 0000000..5837d9a Binary files /dev/null and b/__pycache__/gguf_to_nvfp4_gguf.cpython-313.pyc differ diff --git a/__pycache__/hf_info.cpython-313.pyc b/__pycache__/hf_info.cpython-313.pyc new file mode 100644 index 0000000..30e5c86 Binary files /dev/null and b/__pycache__/hf_info.cpython-313.pyc differ diff --git a/__pycache__/hf_to_nvfp4_gguf.cpython-313.pyc b/__pycache__/hf_to_nvfp4_gguf.cpython-313.pyc new file mode 100644 index 0000000..1914030 Binary files /dev/null and b/__pycache__/hf_to_nvfp4_gguf.cpython-313.pyc differ diff --git a/__pycache__/hf_to_nvfp4_gguf.cpython-314.pyc b/__pycache__/hf_to_nvfp4_gguf.cpython-314.pyc new file mode 100644 index 0000000..e06aad3 Binary files /dev/null and b/__pycache__/hf_to_nvfp4_gguf.cpython-314.pyc differ diff --git a/__pycache__/hf_to_nvfp4_safetensors.cpython-313.pyc b/__pycache__/hf_to_nvfp4_safetensors.cpython-313.pyc new file mode 100644 index 0000000..55285c1 Binary files /dev/null and b/__pycache__/hf_to_nvfp4_safetensors.cpython-313.pyc differ diff --git a/__pycache__/ollama_nvfp4.cpython-313.pyc b/__pycache__/ollama_nvfp4.cpython-313.pyc new file mode 100644 index 0000000..0d6eac5 Binary files /dev/null and b/__pycache__/ollama_nvfp4.cpython-313.pyc differ diff --git a/__pycache__/ollama_nvfp4.cpython-314.pyc b/__pycache__/ollama_nvfp4.cpython-314.pyc new file mode 100644 index 0000000..b859ef7 Binary files /dev/null and b/__pycache__/ollama_nvfp4.cpython-314.pyc differ diff --git a/__pycache__/safetensors_to_nvfp4_gguf.cpython-313.pyc b/__pycache__/safetensors_to_nvfp4_gguf.cpython-313.pyc new file mode 100644 index 0000000..86d4f19 Binary files /dev/null and b/__pycache__/safetensors_to_nvfp4_gguf.cpython-313.pyc differ diff --git a/__pycache__/safetensors_to_nvfp4_gguf.cpython-314.pyc b/__pycache__/safetensors_to_nvfp4_gguf.cpython-314.pyc new file mode 100644 index 0000000..82fac26 Binary files /dev/null and b/__pycache__/safetensors_to_nvfp4_gguf.cpython-314.pyc differ diff --git a/docker/README.txt b/docker/README.txt new file mode 100644 index 0000000..25b8ecd --- /dev/null +++ b/docker/README.txt @@ -0,0 +1,22 @@ +Convert-to-Quant (ctq) is installed for normal quantizations. + +For GGUF quantizations, try: + +hf_to_nvfp4_quant +gguf_to_nvfp4_quant +safetensors_to_nvfp4_quant +hf_to_nvfp4_safetensors + +Depending on the source material. + +hf cli command is also installed, eg you can: + + +hf download coder3101/gemma-4-12B-it-heretic --local-dir gemma-4-12B-it-heretic +hf_to_nvfp4_gguf -o gemma-4-12B-it-heretic_nvfp4 gemma-4-12B-it-heretic + +The same directory can be quantized to one ComfyUI NVFP4 safetensors file: + +hf_to_nvfp4_safetensors -o gemma-4-12B-it-heretic-NVFP4.safetensors gemma-4-12B-it-heretic + +This would create a NVFP4 GGUF for Ollama/llama.cpp from https://huggingface.co/coder3101/gemma-4-12B-it-heretic diff --git a/docker/gguf_to_nvfp4_gguf b/docker/gguf_to_nvfp4_gguf new file mode 100755 index 0000000..d3495a4 --- /dev/null +++ b/docker/gguf_to_nvfp4_gguf @@ -0,0 +1,2 @@ +#!/bin/bash +python /root/nvfp4_gguf_quant/gguf_to_nvfp4_gguf.py --llama-cpp /root/llama.cpp "${@}" diff --git a/docker/hf_to_nvfp4_gguf b/docker/hf_to_nvfp4_gguf new file mode 100755 index 0000000..c6ca43e --- /dev/null +++ b/docker/hf_to_nvfp4_gguf @@ -0,0 +1,2 @@ +#!/bin/bash +python /root/nvfp4_gguf_quant/hf_to_nvfp4_gguf.py --llama-cpp /root/llama.cpp "${@}" diff --git a/docker/hf_to_nvfp4_safetensors b/docker/hf_to_nvfp4_safetensors new file mode 100755 index 0000000..7db788f --- /dev/null +++ b/docker/hf_to_nvfp4_safetensors @@ -0,0 +1,2 @@ +#!/bin/bash +python /root/nvfp4_gguf_quant/hf_to_nvfp4_safetensors.py --convert-to-quant /root/convert_to_quant "${@}" diff --git a/docker/safetensors_to_nvfp4_gguf b/docker/safetensors_to_nvfp4_gguf new file mode 100755 index 0000000..67e4326 --- /dev/null +++ b/docker/safetensors_to_nvfp4_gguf @@ -0,0 +1,2 @@ +#!/bin/bash +python /root/nvfp4_gguf_quant/safetensors_to_nvfp4_gguf.py --llama-cpp /root/llama.cpp "${@}" diff --git a/gguf_info.py b/gguf_info.py new file mode 100644 index 0000000..5a7a8a0 --- /dev/null +++ b/gguf_info.py @@ -0,0 +1,382 @@ +#!/usr/bin/env python3 +"""Print a summary of a GGUF file from its header. + +The weight bytes are not read. Metadata and the tensor directory are enough +for the architecture, quantization, tokenizer, and tensor list. +""" + +from __future__ import annotations + +import argparse +import struct +from collections import Counter +from pathlib import Path + +# GGML type id -> name, and (values per block, bytes per block). +GGML_TYPES = { + 0: ("F32", 1, 4), + 1: ("F16", 1, 2), + 2: ("Q4_0", 32, 18), + 3: ("Q4_1", 32, 20), + 6: ("Q5_0", 32, 22), + 7: ("Q5_1", 32, 24), + 8: ("Q8_0", 32, 34), + 9: ("Q8_1", 32, 36), + 10: ("Q2_K", 256, 84), + 11: ("Q3_K", 256, 110), + 12: ("Q4_K", 256, 144), + 13: ("Q5_K", 256, 176), + 14: ("Q6_K", 256, 210), + 15: ("Q8_K", 256, 292), + 16: ("IQ2_XXS", 256, 66), + 17: ("IQ2_XS", 256, 74), + 18: ("IQ3_XXS", 256, 98), + 19: ("IQ1_S", 256, 50), + 20: ("IQ4_NL", 32, 18), + 21: ("IQ3_S", 256, 110), + 22: ("IQ2_S", 256, 82), + 23: ("IQ4_XS", 256, 136), + 24: ("I8", 1, 1), + 25: ("I16", 1, 2), + 26: ("I32", 1, 4), + 27: ("I64", 1, 8), + 28: ("F64", 1, 8), + 29: ("IQ1_M", 256, 56), + 30: ("BF16", 1, 2), + 34: ("TQ1_0", 256, 54), + 35: ("TQ2_0", 256, 66), + 39: ("MXFP4", 32, 17), + 40: ("NVFP4", 64, 36), + 41: ("Q1_0", 128, 18), + 42: ("Q2_0", 64, 18), +} + +FILE_TYPES = { + 0: "ALL_F32", + 1: "MOSTLY_F16", + 2: "MOSTLY_Q4_0", + 3: "MOSTLY_Q4_1", + 7: "MOSTLY_Q8_0", + 8: "MOSTLY_Q5_0", + 9: "MOSTLY_Q5_1", + 10: "MOSTLY_Q2_K", + 11: "MOSTLY_Q3_K_S", + 12: "MOSTLY_Q3_K_M", + 13: "MOSTLY_Q3_K_L", + 14: "MOSTLY_Q4_K_S", + 15: "MOSTLY_Q4_K_M", + 16: "MOSTLY_Q5_K_S", + 17: "MOSTLY_Q5_K_M", + 18: "MOSTLY_Q6_K", + 19: "MOSTLY_IQ2_XXS", + 20: "MOSTLY_IQ2_XS", + 21: "MOSTLY_Q2_K_S", + 22: "MOSTLY_IQ3_XS", + 23: "MOSTLY_IQ3_XXS", + 24: "MOSTLY_IQ1_S", + 25: "MOSTLY_IQ4_NL", + 26: "MOSTLY_IQ3_S", + 27: "MOSTLY_IQ3_M", + 28: "MOSTLY_IQ2_S", + 29: "MOSTLY_IQ2_M", + 30: "MOSTLY_IQ4_XS", + 31: "MOSTLY_IQ1_M", + 32: "MOSTLY_BF16", + 36: "MOSTLY_TQ1_0", + 37: "MOSTLY_TQ2_0", + 38: "MOSTLY_MXFP4_MOE", + 39: "MOSTLY_NVFP4", + 40: "MOSTLY_Q1_0", + 41: "MOSTLY_Q2_0", + 1024: "GUESSED", +} + +_SCALAR = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8} +_INTEGER = {0: " str: + size = float(num_bytes) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if size < 1024 or unit == "TiB": + return f"{int(size)} B" if unit == "B" else f"{size:.2f} {unit}" + size /= 1024 + return f"{size:.2f} TiB" + + +def _read(handle, n: int) -> bytes: + data = handle.read(n) + if len(data) != n: + raise ValueError("truncated GGUF header") + return data + + +def _read_string(handle) -> bytes: + size = struct.unpack(" 64 * 1024 * 1024: + raise ValueError("GGUF string is too long") + return _read(handle, size) + + +def _read_scalar(handle, value_type: int): + if value_type == 8: + return _read_string(handle).decode("utf-8", "replace") + if value_type not in _SCALAR: + raise ValueError(f"unknown GGUF value type {value_type}") + raw = _read(handle, _SCALAR[value_type]) + if value_type in _INTEGER: + return struct.unpack(_INTEGER[value_type], raw)[0] + if value_type in _FLOAT: + return struct.unpack(_FLOAT[value_type], raw)[0] + return None + + +def _read_array(handle, keep: bool): + subtype = struct.unpack(" limit: + values.append(f"... {count} items") + return values + + +def _type_name(value_type: int) -> str: + names = { + 0: "u8", 1: "i8", 2: "u16", 3: "i16", 4: "u32", 5: "i32", + 6: "f32", 7: "bool", 8: "string", 10: "u64", 11: "i64", 12: "f64", + } + return names.get(value_type, str(value_type)) + + +def load_gguf(path: Path) -> tuple[dict, list[dict], int]: + metadata: dict = {} + tensors: list[dict] = [] + with path.open("rb") as handle: + magic, version = struct.unpack(" str: + for key, value in metadata.items(): + if key.endswith(".context_length") and ".rope." not in key and isinstance(value, int) and value > 0: + return str(value) + return "" + + +def _tensor_bytes(tensor: dict) -> int | None: + spec = GGML_TYPES.get(tensor["type"]) + if spec is None or not tensor["dims"]: + return None + _name, block, type_size = spec + elements = 1 + for dim in tensor["dims"]: + elements *= dim + if elements % block != 0: + return None + return elements // block * type_size + + +def _preview(value, limit: int = 160) -> str: + if isinstance(value, str): + flat = value.replace("\n", "\\n") + if len(flat) > limit: + return f"{flat[:limit]}... ({len(value)} chars)" + return flat + if isinstance(value, list): + text = ", ".join(_preview(item, 40) for item in value[:12]) + if len(value) > 12: + text += f", ... ({len(value)})" + return "[" + text + "]" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, float): + return f"{value:.6g}" + return str(value) + + +def _filetype(value) -> str: + if not isinstance(value, int): + return _preview(value) + name = FILE_TYPES.get(value) + return f"{name} ({value})" if name else str(value) + + +def print_summary(path: Path, metadata: dict, tensors: list[dict], version: int) -> None: + arch = str(metadata.get("general.architecture") or "") + print(f"file {path}") + print(f"size {format_size(path.stat().st_size)}") + print(f"gguf version {version}") + print(f"architecture {arch or '-'}") + for label, key in ( + ("name", "general.name"), + ("basename", "general.basename"), + ("size label", "general.size_label"), + ("finetune", "general.finetune"), + ("type", "general.type"), + ): + if key in metadata and metadata[key] not in ("", None): + print(f"{label:<14}{_preview(metadata[key])}") + if "general.file_type" in metadata: + print(f"filetype {_filetype(metadata['general.file_type'])}") + if "general.quantization_version" in metadata: + print(f"quant version {metadata['general.quantization_version']}") + context = _context(metadata) + if context: + print(f"context {context}") + fields = ( + ("blocks", "block_count"), + ("embedding", "embedding_length"), + ("feed forward", "feed_forward_length"), + ("heads", "attention.head_count"), + ("kv heads", "attention.head_count_kv"), + ("head dim", "attention.key_length"), + ("experts", "expert_count"), + ("vocab", "vocab_size"), + ("sliding window", "attention.sliding_window"), + ) + for label, suffix in fields: + value = _meta(metadata, suffix) + if isinstance(value, int) and value > 0: + print(f"{label:<14}{value}") + swa = _meta(metadata, "attention.sliding_window_pattern") + if isinstance(swa, list) and swa and all(isinstance(item, (bool, int)) for item in swa): + flags = [bool(item) for item in swa] + print(f"swa layers {sum(flags)} sliding, {len(flags) - sum(flags)} full") + tok_model = metadata.get("tokenizer.ggml.model") + if tok_model: + print(f"tokenizer {tok_model}") + for label, key in ( + ("bos", "tokenizer.ggml.bos_token_id"), + ("eos", "tokenizer.ggml.eos_token_id"), + ): + if key in metadata: + print(f"{label:<14}{metadata[key]}") + template = metadata.get("tokenizer.chat_template") + if isinstance(template, str) and template: + print(f"chat template yes ({len(template)} chars)") + elif isinstance(template, list) and template: + print(f"chat template {len(template)} templates") + else: + print("chat template no") + if isinstance(metadata.get("general.description"), str): + print(f"description {_preview(metadata['general.description'], 200)}") + + counts: Counter[str] = Counter() + known_bytes = 0 + unknown = 0 + for tensor in tensors: + spec = GGML_TYPES.get(tensor["type"]) + counts[spec[0] if spec else f"type {tensor['type']}"] += 1 + nbytes = _tensor_bytes(tensor) + if nbytes is None: + unknown += 1 + else: + known_bytes += nbytes + print(f"tensors {len(tensors)}") + for name, count in counts.most_common(): + print(f" {name:<10} {count}") + print(f"tensor bytes {format_size(known_bytes)}") + if unknown: + print(f"unmeasured {unknown} tensors have an unknown type or shape") + + +def print_metadata(metadata: dict) -> None: + print() + print("metadata") + width = max(len(key) for key in metadata) if metadata else 8 + for key in metadata: + print(f" {key:<{width}} {_preview(metadata[key])}") + + +def print_tensors(tensors: list[dict]) -> None: + print() + print(f"{'NAME':<48} {'TYPE':<8} {'SHAPE'}") + for tensor in tensors: + spec = GGML_TYPES.get(tensor["type"]) + kind = spec[0] if spec else str(tensor["type"]) + # File order: the first dimension is the contiguous one. + shape = " x ".join(str(dim) for dim in tensor["dims"]) or "-" + print(f" {tensor['name']:<46} {kind:<8} {shape}") + + +def inspect(path: Path, metadata_dump: bool, tensor_dump: bool) -> None: + if not path.is_file(): + raise SystemExit(f"Not a file: {path}") + try: + metadata, tensors, version = load_gguf(path) + except (OSError, ValueError, struct.error, UnicodeError) as exc: + raise SystemExit(f"{path}: {exc}") from exc + print_summary(path, metadata, tensors, version) + if metadata_dump: + print_metadata(metadata) + if tensor_dump: + print_tensors(tensors) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Print architecture, quantization, and tensor info from a GGUF header.") + parser.add_argument("gguf", nargs="+", type=Path, help="One or more .gguf files") + parser.add_argument("--metadata", action="store_true", help="Print every metadata key. Long strings and the vocabulary are shortened.") + parser.add_argument("--tensors", action="store_true", help="Print every tensor name, type, and shape.") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + for index, path in enumerate(args.gguf): + if index: + print() + inspect(path, args.metadata, args.tensors) + + +if __name__ == "__main__": + main() diff --git a/gguf_to_nvfp4_gguf.py b/gguf_to_nvfp4_gguf.py new file mode 100644 index 0000000..1c61087 --- /dev/null +++ b/gguf_to_nvfp4_gguf.py @@ -0,0 +1,667 @@ +#!/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: ./-NVFP4, or ./ 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-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() diff --git a/hf_info.py b/hf_info.py new file mode 100644 index 0000000..db48505 --- /dev/null +++ b/hf_info.py @@ -0,0 +1,1105 @@ +#!/usr/bin/env python3 +"""Print a summary of a Hugging Face checkpoint. + +Weight bytes are not read. config.json, the tokenizer files, and safetensors +headers are enough for the architecture, quantization, and tensor list. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import struct +import sys +from collections import Counter +from dataclasses import dataclass +from pathlib import Path + +# Same round-trip check as nvfp4gguf.bundle.chat_template_has_tools. +_TOOL_RESULT_MARKERS = ( + "tool_response", + "tool_results", + "role'] == 'tool'", + "role'] == \"tool\"", + "role\"] == 'tool'", + "role\"] == \"tool\"", + "message.role == 'tool'", + "message.role == \"tool\"", + "ipython", +) + +_NESTED_CONFIGS = ( + "text_config", + "llm_config", + "language_config", + "thinker_config", + "decoder", +) + +_INDEX_NAMES = ( + "model.safetensors.index.json", + "diffusion_pytorch_model.safetensors.index.json", +) + +_MAX_HEADER = 64 * 1024 * 1024 + + +@dataclass +class TensorInfo: + name: str + dtype: str + shape: list[int] + nbytes: int | None + file: str + + +@dataclass +class ShardInfo: + path: Path + header_size: int + data_end: int + metadata: dict + + +@dataclass +class Checkpoint: + path: Path + root: Path + config: dict + config_name: str + generation: dict + tokenizer: dict + quant: dict + hf_quant: dict | None + recipe: str | None + base_model: str | None + tensors: list[TensorInfo] + shards: list[ShardInfo] + shard_total: int + partial: bool + index_meta: dict + template: str + template_where: str + template_count: int + nvfp4: bool + comfy: bool + + +def format_size(num_bytes: int) -> str: + size = float(num_bytes) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if size < 1024 or unit == "TiB": + return f"{int(size)} B" if unit == "B" else f"{size:.2f} {unit}" + size /= 1024 + return f"{size:.2f} TiB" + + +def format_params(count: int) -> str: + if count >= 1_000_000_000: + return f"{count} ({count / 1_000_000_000:.2f}B)" + if count >= 1_000_000: + return f"{count} ({count / 1_000_000:.2f}M)" + return str(count) + + +def format_number(value) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, int): + return str(value) + if isinstance(value, float): + if value.is_integer() and abs(value) < 1e15: + return str(int(value)) + return f"{value:.6g}" + return _preview(value) + + +def _preview(value, limit: int = 160) -> str: + if isinstance(value, str): + flat = value.replace("\n", "\\n") + if len(flat) > limit: + return f"{flat[:limit]}... ({len(value)} chars)" + return flat + if isinstance(value, dict): + text = json.dumps(value, ensure_ascii=False, default=str, separators=(", ", ": ")) + if len(text) > limit: + return f"{text[:limit]}... ({len(text)} chars)" + return text + if isinstance(value, list): + text = ", ".join(_preview(item, 40) for item in value[:12]) + if len(value) > 12: + text += f", ... ({len(value)})" + return "[" + text + "]" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, float): + return format_number(value) + if value is None: + return "null" + return str(value) + + +def _line(label: str, value) -> None: + print(f"{label:<16}{value}") + + +def read_json(path: Path): + try: + text = path.read_text(encoding="utf-8") + except (OSError, UnicodeError) as exc: + raise ValueError(f"{path}: {exc}") from exc + try: + return json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}: invalid JSON ({exc.msg})") from exc + + +def read_header(path: Path) -> tuple[dict, int]: + with path.open("rb") as handle: + raw = handle.read(8) + if len(raw) != 8: + raise ValueError(f"{path} is not a safetensors file") + header_size = struct.unpack(" _MAX_HEADER: + raise ValueError(f"{path} is not a safetensors file") + payload = handle.read(header_size) + if len(payload) != header_size: + raise ValueError(f"{path} has a truncated safetensors header") + try: + header = json.loads(payload) + except json.JSONDecodeError as exc: + raise ValueError(f"{path} has an unreadable safetensors header") from exc + if not isinstance(header, dict): + raise ValueError(f"{path} has an unreadable safetensors header") + return header, header_size + + +def _sections(config: dict) -> list[dict]: + found = [config] + for key in _NESTED_CONFIGS: + nested = config.get(key) + if isinstance(nested, dict): + found.append(nested) + return found + + +def _present(config: dict, keys: tuple[str, ...]): + for section in _sections(config): + for key in keys: + if key in section and section[key] not in (None, ""): + return section[key] + return None + + +def _positive_int(config: dict, keys: tuple[str, ...], limit: int = 10**12) -> int | None: + value = _present(config, keys) + if isinstance(value, bool): + return None + if isinstance(value, int) and 0 < value <= limit: + return value + if isinstance(value, float) and value.is_integer() and 0 < value <= limit: + return int(value) + return None + + +def _text(value) -> str: + return value if isinstance(value, str) and value else "" + + +def chat_template_has_tools(template: str) -> bool: + 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 _template_parts(value) -> list[str]: + if isinstance(value, str): + return [value] if value.strip() else [] + if not isinstance(value, list): + return [] + parts = [] + for item in value: + if isinstance(item, str) and item.strip(): + parts.append(item) + elif isinstance(item, dict) and isinstance(item.get("template"), str) and item["template"].strip(): + parts.append(item["template"]) + return parts + + +def load_chat_template(root: Path, tokenizer: dict) -> tuple[str, str, int]: + parts = _template_parts(tokenizer.get("chat_template")) if tokenizer else [] + if parts: + count = len(tokenizer.get("chat_template")) if isinstance(tokenizer.get("chat_template"), list) else 1 + return "\n".join(parts), "tokenizer_config.json", count + path = root / "chat_template.jinja" + if path.is_file(): + try: + return path.read_text(encoding="utf-8"), path.name, 1 + except (OSError, UnicodeError) as exc: + raise ValueError(f"{path}: {exc}") from exc + return "", "", 0 + + +def _token_text(value) -> str: + if isinstance(value, str): + return value + if isinstance(value, dict) and isinstance(value.get("content"), str): + return value["content"] + return "" + + +def token_label(token_id, tokenizer: dict) -> str: + if isinstance(token_id, list): + return ", ".join(token_label(item, tokenizer) for item in token_id) + name = "" + decoder = tokenizer.get("added_tokens_decoder") if tokenizer else None + if isinstance(decoder, dict): + entry = decoder.get(str(token_id)) + if isinstance(entry, dict): + name = _token_text(entry.get("content")) + if name: + return f"{token_id} ({name})" + return str(token_id) + + +def _token_id(source: dict, key: str): + if not source or key not in source: + return None + value = source.get(key) + if isinstance(value, bool) or value is None or value == "": + return None + if isinstance(value, int): + return value + if isinstance(value, list) and value and all(isinstance(item, int) and not isinstance(item, bool) for item in value): + return value + return None + + +def read_base_model(path: Path) -> str | None: + if not path.is_file(): + return None + try: + text = path.read_text(encoding="utf-8", errors="replace")[:65536] + except OSError: + return None + if not text.startswith("---"): + return None + end = text.find("\n---", 3) + if end < 0: + return None + values: list[str] = [] + grabbing = False + for line in text[3:end].splitlines(): + if not grabbing: + if line.startswith("base_model:"): + rest = line.split(":", 1)[1].strip().strip("'\"") + if rest: + values.append(rest) + else: + grabbing = True + continue + stripped = line.strip() + if stripped.startswith("- "): + values.append(stripped[2:].strip().strip("'\"")) + continue + break + cleaned = [value for value in values if value] + return ", ".join(cleaned) if cleaned else None + + +def read_recipe_scheme(path: Path) -> str | None: + if not path.is_file(): + return None + try: + text = path.read_text(encoding="utf-8", errors="replace")[:65536] + except OSError: + return None + for line in text.splitlines(): + stripped = line.strip() + if stripped.startswith("scheme:"): + value = stripped.split(":", 1)[1].strip().strip("'\"") + return value or None + return None + + +def load_quant(config: dict, root: Path) -> tuple[dict, dict | None]: + raw = config.get("quantization_config") + quant = dict(raw) if isinstance(raw, dict) else {} + path = root / "hf_quant_config.json" + if not path.is_file(): + return quant, None + extra = read_json(path) + if not isinstance(extra, dict): + raise ValueError(f"{path} is not a JSON object") + block = extra.get("quantization") if isinstance(extra.get("quantization"), dict) else {} + producer = extra.get("producer") if isinstance(extra.get("producer"), dict) else {} + for key, value in block.items(): + quant.setdefault(key, value) + if "quant_method" not in quant and isinstance(producer.get("name"), str) and producer["name"]: + quant["quant_method"] = producer["name"].lower() + return quant, extra + + +def quant_is_nvfp4(quant: dict) -> bool: + algo = str(quant.get("quant_algo") or "") + fmt = str(quant.get("format") or "") + if algo in {"NVFP4", "W4A16_NVFP4"} or algo.endswith("NVFP4"): + return True + if "nvfp4" in fmt.lower(): + return True + groups = quant.get("config_groups") + if isinstance(groups, dict): + for group in groups.values(): + if isinstance(group, dict) and "nvfp4" in str(group.get("format") or "").lower(): + return True + layers = quant.get("quantized_layers") + if isinstance(layers, dict): + for entry in layers.values(): + if isinstance(entry, dict) and str(entry.get("quant_algo") or "").endswith("NVFP4"): + return True + return False + + +def _weight_schemes(quant: dict) -> list[tuple]: + found = [] + groups = quant.get("config_groups") + if isinstance(groups, dict): + for group in groups.values(): + if not isinstance(group, dict): + continue + weights = group.get("weights") if isinstance(group.get("weights"), dict) else {} + acts = group.get("input_activations") if isinstance(group.get("input_activations"), dict) else {} + found.append((weights.get("num_bits"), weights.get("group_size"), acts.get("num_bits"))) + if found: + return found + bits = quant.get("bits", quant.get("num_bits", quant.get("w_bit"))) + group_size = quant.get("group_size", quant.get("groupsize")) + if bits is not None or group_size is not None: + return [(bits, group_size, None)] + return [] + + +def find_index(root: Path) -> Path | None: + for name in _INDEX_NAMES: + path = root / name + if path.is_file(): + return path + found = sorted(path for path in root.glob("*.safetensors.index.json") if path.is_file()) + return found[0] if found else None + + +def shards_from_index(index_path: Path, payload: dict) -> list[Path]: + weight_map = payload.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"{index_path} has no weight_map") + root = index_path.parent + names: list[Path] = [] + seen: set[str] = set() + for name in weight_map.values(): + if not isinstance(name, str) or not name or name.startswith(("/", "\\")) or ".." in Path(name).parts: + raise ValueError(f"{index_path} has an invalid weight_map entry") + if name not in seen: + seen.add(name) + names.append(root / name) + return names + + +def discover_shards(root: Path) -> list[Path]: + for name in ("model.safetensors", "diffusion_pytorch_model.safetensors"): + path = root / name + if path.is_file(): + return [path] + return sorted(path for path in root.glob("*.safetensors") if path.is_file()) + + +def parse_tensor(name: str, entry, file_name: str) -> tuple[TensorInfo, int]: + if not isinstance(entry, dict): + raise ValueError(f"{file_name}: tensor {name!r} has an invalid header entry") + dtype = entry.get("dtype") + if not isinstance(dtype, str) or not dtype: + dtype = "?" + shape = entry.get("shape") + if shape is None: + dims: list[int] = [] + elif isinstance(shape, list) and all(isinstance(dim, int) and not isinstance(dim, bool) for dim in shape): + dims = [int(dim) for dim in shape] + else: + raise ValueError(f"{file_name}: tensor {name!r} has an invalid shape") + offsets = entry.get("data_offsets") + nbytes = None + end = 0 + if isinstance(offsets, list) and len(offsets) == 2 and all(isinstance(item, int) and not isinstance(item, bool) for item in offsets): + start, end = int(offsets[0]), int(offsets[1]) + if start < 0 or end < start: + raise ValueError(f"{file_name}: tensor {name!r} has invalid data offsets") + nbytes = end - start + return TensorInfo(name, dtype, dims, nbytes, file_name), end + + +def load_shards(paths: list[Path]) -> tuple[list[TensorInfo], list[ShardInfo]]: + tensors: list[TensorInfo] = [] + shards: list[ShardInfo] = [] + owners: dict[str, str] = {} + for path in paths: + header, header_size = read_header(path) + meta = header.pop("__metadata__", None) + if not isinstance(meta, dict): + meta = {} + data_end = 0 + for name, entry in header.items(): + if name.startswith("__"): + continue + if name in owners: + raise ValueError(f"Tensor {name!r} appears in both {owners[name]} and {path.name}") + owners[name] = path.name + tensor, end = parse_tensor(name, entry, path.name) + data_end = max(data_end, end) + tensors.append(tensor) + shards.append(ShardInfo(path, header_size, data_end, meta)) + return tensors, shards + + +def _natural_key(name: str) -> tuple: + return tuple((0, int(part)) if part.isdigit() else (1, part) for part in re.split(r"(\d+)", name)) + + +def tensor_type(tensor: TensorInfo, quant_method: str, nvfp4: bool, comfy: bool) -> str: + leaf = tensor.name.rsplit(".", 1)[-1] + if nvfp4 and leaf == "weight_packed": + return "NVFP4" + if comfy and leaf == "weight" and tensor.dtype == "U8": + return "NVFP4" + if quant_method.lower() in {"gptq", "awq"} and leaf == "qweight": + return quant_method.upper() + return tensor.dtype + + +def _numel(shape: list[int]) -> int: + count = 1 + for dim in shape: + count *= dim + return count + + +def parameter_count(tensors: list[TensorInfo], type_of) -> int | None: + """Logical weights and biases. + + Packed NVFP4 stores two FP4 values in each byte, so those element counts + are doubled. Scale tensors are not parameters. Other packed layouts are + left uncounted rather than reported as a partial total. + """ + total = 0 + seen = False + for tensor in tensors: + leaf = tensor.name.rsplit(".", 1)[-1] + packed = leaf in {"weight_packed", "qweight"} + if packed and type_of(tensor) != "NVFP4": + return None + if leaf not in {"weight", "bias", "weight_packed"}: + continue + seen = True + count = _numel(tensor.shape) + if type_of(tensor) == "NVFP4": + count *= 2 + total += count + return total if seen else None + + +def _is_model_dir(path: Path) -> bool: + return path.name in {"config.json", "adapter_config.json"} or path.name.endswith(".safetensors.index.json") + + +def _load(path: Path) -> Checkpoint: + if not path.exists(): + raise SystemExit(f"Not found: {path}") + if path.is_file() and path.suffix == ".gguf": + raise SystemExit(f"{path} is a GGUF file. Use gguf_info.py for GGUF files.") + if path.is_file() and _is_model_dir(path): + root = path.parent + selected: list[Path] | None = None + elif path.is_file(): + if path.suffix != ".safetensors": + raise SystemExit(f"{path} is not a model directory or a safetensors file") + root = path.parent + selected = [path] + elif path.is_dir(): + root = path + selected = None + else: + raise SystemExit(f"Not a file or directory: {path}") + + config_path = root / "config.json" + config_name = "config.json" + if not config_path.is_file() and (root / "adapter_config.json").is_file(): + config_path = root / "adapter_config.json" + config_name = config_path.name + config: dict = {} + if config_path.is_file(): + loaded = read_json(config_path) + if not isinstance(loaded, dict): + raise ValueError(f"{config_path} is not a JSON object") + config = loaded + + generation_path = root / "generation_config.json" + generation: dict = {} + if generation_path.is_file(): + loaded = read_json(generation_path) + if isinstance(loaded, dict): + generation = loaded + + tokenizer_path = root / "tokenizer_config.json" + tokenizer: dict = {} + if tokenizer_path.is_file(): + loaded = read_json(tokenizer_path) + if isinstance(loaded, dict): + tokenizer = loaded + + quant, hf_quant = load_quant(config, root) + recipe = read_recipe_scheme(root / "recipe.yaml") + base_model = read_base_model(root / "README.md") + repo = config.get("_name_or_path") + if not base_model and isinstance(repo, str) and "/" in repo and not repo.startswith(("/", ".")): + base_model = repo + + index_path = find_index(root) + index_payload: dict = {} + weight_map: dict = {} + indexed: list[Path] = [] + if index_path is not None: + loaded = read_json(index_path) + if not isinstance(loaded, dict): + raise ValueError(f"{index_path} is not a JSON object") + index_payload = loaded + raw_map = loaded.get("weight_map") + weight_map = raw_map if isinstance(raw_map, dict) else {} + indexed = shards_from_index(index_path, loaded) + if selected is None: + shard_paths = indexed if indexed else discover_shards(root) + missing = [path.name for path in shard_paths if not path.is_file()] + if missing: + raise ValueError(f"Index entries are missing from {root}: {', '.join(missing)}") + else: + shard_paths = selected + + if not config and not shard_paths: + raise SystemExit(f"No config.json or safetensors files in {root}") + + tensors, shards = load_shards(shard_paths) if shard_paths else ([], []) + opened = {shard.path.name for shard in shards} + index_names = {path.name for path in indexed} + partial = bool(index_names) and opened != index_names and opened.issubset(index_names) + if index_names and opened == index_names: + missing = [name for name in weight_map if isinstance(name, str) and name not in {tensor.name for tensor in tensors}] + if missing: + raise ValueError(f"Index lists {missing[0]!r}, which is not in the safetensors headers") + + comfy = any(tensor.name.endswith(".comfy_quant") for tensor in tensors) or any( + "_quantization_metadata" in shard.metadata for shard in shards + ) + packed = any(tensor.name.endswith("weight_packed") for tensor in tensors) + recipe_nvfp4 = bool(recipe) and "NVFP4" in recipe.upper() + nvfp4 = quant_is_nvfp4(quant) or comfy or (recipe_nvfp4 and packed) + template, template_where, template_count = load_chat_template(root, tokenizer) + index_meta = index_payload.get("metadata") if isinstance(index_payload.get("metadata"), dict) else {} + + return Checkpoint( + path=path, + root=root, + config=config, + config_name=config_name, + generation=generation, + tokenizer=tokenizer, + quant=quant, + hf_quant=hf_quant, + recipe=recipe, + base_model=base_model, + tensors=tensors, + shards=shards, + shard_total=len(indexed) if indexed else len(shards), + partial=partial, + index_meta=index_meta, + template=template, + template_where=template_where, + template_count=template_count, + nvfp4=nvfp4, + comfy=comfy, + ) + + +def load_checkpoint(path: Path) -> Checkpoint: + try: + return _load(path) + except (OSError, ValueError, struct.error, UnicodeError) as exc: + raise SystemExit(str(exc)) from exc + + +def _type_of(checkpoint: Checkpoint): + method = str(checkpoint.quant.get("quant_method") or "") + + def kind(tensor: TensorInfo) -> str: + return tensor_type(tensor, method, checkpoint.nvfp4, checkpoint.comfy) + + return kind + + +def _quant_title(checkpoint: Checkpoint) -> str: + quant = checkpoint.quant + algo = quant.get("quant_algo") + if isinstance(algo, str) and algo: + return algo + if checkpoint.nvfp4: + return "NVFP4" + fmt = quant.get("format") + if isinstance(fmt, str) and fmt: + return fmt + method = quant.get("quant_method") + if isinstance(method, str) and method: + return method + if checkpoint.recipe: + return checkpoint.recipe + return "" + + +def print_identity(checkpoint: Checkpoint) -> None: + _line("file", checkpoint.path) + if checkpoint.shards: + size = sum(shard.path.stat().st_size for shard in checkpoint.shards) + _line("size", format_size(size)) + model_type = _text(_present(checkpoint.config, ("model_type",))) + architectures = checkpoint.config.get("architectures") + if isinstance(architectures, list): + arch_text = ", ".join(str(item) for item in architectures if item not in (None, "")) + else: + arch_text = _text(architectures) + if model_type: + _line("architecture", model_type) + if arch_text and arch_text != model_type: + _line("class", arch_text) + elif arch_text: + _line("architecture", arch_text) + if checkpoint.path.is_dir() or checkpoint.config: + _line("name", checkpoint.root.name) + else: + _line("name", checkpoint.path.stem) + if checkpoint.base_model and checkpoint.base_model != checkpoint.root.name: + _line("base model", checkpoint.base_model) + dtype = _text(_present(checkpoint.config, ("dtype", "torch_dtype"))) + if dtype: + _line("dtype", dtype) + + +def print_quantization(checkpoint: Checkpoint) -> None: + title = _quant_title(checkpoint) + if checkpoint.nvfp4 or checkpoint.quant: + if title: + _line("quantization", title) + elif checkpoint.recipe: + _line("recipe", checkpoint.recipe) + return + else: + return + method = _text(checkpoint.quant.get("quant_method")) + fmt = _text(checkpoint.quant.get("format")) + if method and method != title: + _line("quant method", method) + if fmt and fmt != title: + _line("quant format", fmt) + elif checkpoint.comfy and not fmt: + _line("quant format", "comfy") + schemes = _weight_schemes(checkpoint.quant) + bits = [] + groups = [] + act_bits = [] + for weight_bits, group_size, activation_bits in schemes: + if isinstance(weight_bits, int) and not isinstance(weight_bits, bool) and weight_bits not in bits: + bits.append(weight_bits) + if isinstance(group_size, int) and not isinstance(group_size, bool) and group_size > 0 and group_size not in groups: + groups.append(group_size) + if isinstance(activation_bits, int) and not isinstance(activation_bits, bool) and activation_bits not in act_bits: + act_bits.append(activation_bits) + if bits: + _line("bits", ", ".join(str(item) for item in bits)) + if groups: + _line("group size", ", ".join(str(item) for item in groups)) + if act_bits: + _line("act bits", ", ".join(str(item) for item in act_bits)) + status = _text(checkpoint.quant.get("quantization_status")) + if status: + _line("quant status", status) + version = checkpoint.quant.get("version") + if isinstance(version, (str, int, float)) and not isinstance(version, bool) and version != "": + _line("quant version", format_number(version) if isinstance(version, (int, float)) else version) + ignored = checkpoint.quant.get("ignore") + if isinstance(ignored, list) and ignored: + if all(isinstance(item, str) for item in ignored): + text = ", ".join(ignored) + if len(text) > 160: + text = text[:160] + "..." + _line("quant ignore", text) + else: + _line("quant ignore", _preview(ignored)) + elif isinstance(ignored, str) and ignored: + _line("quant ignore", ignored) + layers = _comfy_layers(checkpoint) + if layers: + _line("quant layers", layers) + + +def _comfy_layers(checkpoint: Checkpoint) -> int: + total = 0 + found = False + for shard in checkpoint.shards: + raw = shard.metadata.get("_quantization_metadata") + parsed = None + if isinstance(raw, str): + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + parsed = None + elif isinstance(raw, dict): + parsed = raw + if not isinstance(parsed, dict): + continue + layers = parsed.get("layers") + if isinstance(layers, dict): + total += len(layers) + found = True + return total if found else 0 + + +def _submodule_text(block: dict) -> str: + parts = [] + for label, keys in ( + ("layers", ("num_hidden_layers", "num_layers", "depth")), + ("width", ("hidden_size", "d_model", "embed_dim")), + ("heads", ("num_attention_heads", "num_heads")), + ): + for key in keys: + value = block.get(key) + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + parts.append(f"{value} {label}") + break + return ", ".join(parts) + + +def print_architecture(checkpoint: Checkpoint) -> None: + config = checkpoint.config + context = _positive_int(config, ( + "max_position_embeddings", + "max_seq_len", + "n_ctx", + "n_positions", + "max_sequence_length", + "seq_length", + "model_max_length", + )) + if context: + _line("context", context) + tok_length = _positive_int(checkpoint.tokenizer, ("model_max_length",)) + if tok_length and tok_length != context: + _line("tok length", tok_length) + fields = ( + ("blocks", ("num_hidden_layers", "n_layer", "num_layers", "n_layers")), + ("embedding", ("hidden_size", "n_embd", "d_model", "hidden_dim")), + ("feed forward", ("intermediate_size", "n_inner", "ffn_dim", "intermediate_dim")), + ("heads", ("num_attention_heads", "n_head", "num_heads")), + ("kv heads", ("num_key_value_heads", "num_kv_heads", "n_head_kv")), + ("experts", ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts")), + ("experts used", ("num_experts_per_tok", "num_experts_per_token")), + ("expert ffn", ("moe_intermediate_size", "expert_intermediate_size")), + ("vocab", ("vocab_size", "padded_vocab_size")), + ) + values = {label: _positive_int(config, keys) for label, keys in fields} + head_dim = _positive_int(config, ("head_dim", "attention_head_dim", "head_size")) + if head_dim is None and values["embedding"] and values["heads"] and values["embedding"] % values["heads"] == 0: + if _present(config, ("qk_nope_head_dim", "qk_rope_head_dim")) is None: + head_dim = values["embedding"] // values["heads"] + for label, _keys in fields: + if values[label]: + _line(label, values[label]) + if label == "kv heads" and head_dim: + _line("head dim", head_dim) + activation = _text(_present(config, ("hidden_act", "activation_function", "hidden_activation"))) + if activation: + _line("activation", activation) + rope = _present(config, ("rope_theta",)) + if isinstance(rope, (int, float)) and not isinstance(rope, bool) and rope != 0: + _line("rope theta", format_number(rope)) + window = _positive_int(config, ("sliding_window", "sliding_window_size")) + if window and _present(config, ("use_sliding_window",)) is not False: + _line("sliding window", window) + layer_types = _present(config, ("layer_types", "sliding_window_pattern")) + if isinstance(layer_types, list) and layer_types: + if all(isinstance(item, bool) for item in layer_types): + sliding = sum(bool(item) for item in layer_types) + _line("swa layers", f"{sliding} sliding, {len(layer_types) - sliding} full") + elif all(isinstance(item, str) for item in layer_types): + counts = Counter(layer_types) + if not (len(counts) == 1 and next(iter(counts)) in {"full_attention", "full"}): + text = ", ".join(f"{count} {name}" for name, count in counts.most_common()) + _line("layers", text) + for key, label in (("vision_config", "vision"), ("audio_config", "audio"), ("vision_tower", "vision")): + block = config.get(key) + if isinstance(block, dict): + text = _submodule_text(block) + if text: + _line(label, text) + elif isinstance(block, str) and block: + _line(label, block) + if _present(config, ("tie_word_embeddings", "tie_embeddings")) is True: + _line("tied embed", "yes") + + +def print_tokenizer(checkpoint: Checkpoint) -> None: + tokenizer = checkpoint.tokenizer + kind = _text(tokenizer.get("tokenizer_class")) + if not kind and (checkpoint.root / "tokenizer.model").is_file(): + kind = "sentencepiece" + elif not kind and (checkpoint.root / "tokenizer.json").is_file(): + kind = "tokenizer.json" + if kind: + _line("tokenizer", kind) + bos = _token_id(checkpoint.config, "bos_token_id") + if bos is None: + bos = _token_id(checkpoint.generation, "bos_token_id") + eos = _token_id(checkpoint.generation, "eos_token_id") + if eos is None: + eos = _token_id(checkpoint.config, "eos_token_id") + pad = _token_id(checkpoint.generation, "pad_token_id") + if pad is None: + pad = _token_id(checkpoint.config, "pad_token_id") + if bos is not None: + _line("bos", token_label(bos, tokenizer)) + elif _token_text(tokenizer.get("bos_token")): + _line("bos", _token_text(tokenizer.get("bos_token"))) + if eos is not None: + _line("eos", token_label(eos, tokenizer)) + elif _token_text(tokenizer.get("eos_token")): + _line("eos", _token_text(tokenizer.get("eos_token"))) + if pad is not None: + _line("pad", token_label(pad, tokenizer)) + elif _token_text(tokenizer.get("pad_token")): + _line("pad", _token_text(tokenizer.get("pad_token"))) + if checkpoint.template: + detail = f"{len(checkpoint.template)} chars" + if checkpoint.template_count > 1: + detail = f"{checkpoint.template_count} templates, {detail}" + if checkpoint.template_where: + detail = f"{detail}, {checkpoint.template_where}" + if chat_template_has_tools(checkpoint.template): + detail = f"{detail}, tools" + _line("chat template", f"yes ({detail})") + else: + _line("chat template", "no") + + +def print_tensor_summary(checkpoint: Checkpoint) -> None: + if checkpoint.partial: + _line("shards", f"{len(checkpoint.shards)} of {checkpoint.shard_total}") + elif len(checkpoint.shards) > 1: + _line("shards", str(len(checkpoint.shards))) + truncated = [ + shard.path.name + for shard in checkpoint.shards + if shard.path.stat().st_size < 8 + shard.header_size + shard.data_end + ] + if truncated: + shown = ", ".join(truncated[:4]) + if len(truncated) > 4: + shown += f", ... ({len(truncated)})" + _line("truncated", shown) + kind = _type_of(checkpoint) + params = None if checkpoint.partial else parameter_count(checkpoint.tensors, kind) + if params: + _line("parameters", format_params(params)) + _line("tensors", str(len(checkpoint.tensors))) + if not checkpoint.tensors: + return + counts: Counter[str] = Counter() + stored: dict[str, Counter[str]] = {} + known = 0 + unknown = 0 + for tensor in checkpoint.tensors: + name = kind(tensor) + counts[name] += 1 + stored.setdefault(name, Counter())[tensor.dtype] += 1 + if tensor.nbytes is None: + unknown += 1 + else: + known += tensor.nbytes + width = max([10, *(len(name) for name in counts)]) + for name, count in counts.most_common(): + extra = "" + dtypes = stored[name] + if list(dtypes) != [name]: + if len(dtypes) == 1: + extra = f" (stored {next(iter(dtypes))})" + else: + detail = ", ".join(f"{dtype} x {number}" for dtype, number in dtypes.most_common()) + extra = f" (stored {detail})" + print(f" {name:<{width}} {count}{extra}") + _line("tensor bytes", format_size(known)) + if unknown: + _line("unmeasured", f"{unknown} tensors have no data length") + + +def print_summary(checkpoint: Checkpoint) -> None: + print_identity(checkpoint) + print_quantization(checkpoint) + print_architecture(checkpoint) + print_tokenizer(checkpoint) + print_tensor_summary(checkpoint) + + +def _metadata_tokenizer(tokenizer: dict) -> dict: + view = {} + for key, value in tokenizer.items(): + if key == "added_tokens_decoder" and isinstance(value, dict): + view[key] = f"{len(value)} added tokens" + else: + view[key] = value + return view + + +def _metadata_safetensors(metadata: dict) -> dict: + view = {} + for key, value in metadata.items(): + if key == "format" and value == "pt": + continue + if key == "_quantization_metadata": + parsed = None + if isinstance(value, str): + try: + parsed = json.loads(value) + except json.JSONDecodeError: + parsed = None + elif isinstance(value, dict): + parsed = value + layers = parsed.get("layers") if isinstance(parsed, dict) else None + if isinstance(parsed, dict) and isinstance(layers, dict): + version = parsed.get("format_version", "?") + view[key] = f"format {version}, {len(layers)} layers" + continue + view[key] = value + return view + + +def _print_mapping(title: str, mapping: dict) -> None: + if not mapping: + return + print() + print(title) + width = max(len(str(key)) for key in mapping) + for key, value in mapping.items(): + print(f" {key:<{width}} {_preview(value)}") + + +def print_metadata(checkpoint: Checkpoint) -> None: + if checkpoint.config: + _print_mapping(checkpoint.config_name, checkpoint.config) + if checkpoint.hf_quant: + _print_mapping("hf_quant_config.json", checkpoint.hf_quant) + if checkpoint.generation: + _print_mapping("generation_config.json", checkpoint.generation) + if checkpoint.tokenizer: + _print_mapping("tokenizer_config.json", _metadata_tokenizer(checkpoint.tokenizer)) + if checkpoint.template and not _template_parts(checkpoint.tokenizer.get("chat_template") if checkpoint.tokenizer else None): + _print_mapping(checkpoint.template_where or "chat_template", {"chat_template": checkpoint.template}) + if checkpoint.index_meta: + _print_mapping("safetensors index", checkpoint.index_meta) + metas = [] + for shard in checkpoint.shards: + cleaned = _metadata_safetensors(shard.metadata) + if cleaned and cleaned not in metas: + metas.append(cleaned) + title = "safetensors metadata" if len(checkpoint.shards) == 1 else f"safetensors metadata ({shard.path.name})" + if len(checkpoint.shards) > 1 and all(_metadata_safetensors(item.metadata) == cleaned for item in checkpoint.shards): + title = "safetensors metadata" + _print_mapping(title, cleaned) + if len(checkpoint.shards) > 1 and all(_metadata_safetensors(item.metadata) == cleaned for item in checkpoint.shards): + break + + +def _shape_text(tensor: TensorInfo) -> str: + return " x ".join(str(dim) for dim in tensor.shape) or "-" + + +def print_tensors(checkpoint: Checkpoint) -> None: + print() + if not checkpoint.tensors: + print("tensors") + print(" (none)") + return + kind = _type_of(checkpoint) + tensors = sorted(checkpoint.tensors, key=lambda tensor: _natural_key(tensor.name)) + if any(kind(tensor) == "NVFP4" for tensor in tensors): + print("NVFP4 shapes are packed storage, 2 FP4 values per byte.") + name_width = max(len(tensor.name) for tensor in tensors) + type_width = max(8, max(len(kind(tensor)) for tensor in tensors)) + shapes = [_shape_text(tensor) for tensor in tensors] + show_shard = len({tensor.file for tensor in tensors}) > 1 + shape_width = max(len("SHAPE"), *(len(shape) for shape in shapes)) if show_shard else 0 + shape_header = f"{'SHAPE':<{shape_width}}" if show_shard else "SHAPE" + header = f"{'NAME':<{name_width + 2}} {'TYPE':<{type_width}} {shape_header}" + if show_shard: + header += " SHARD" + print(header) + for tensor, shape in zip(tensors, shapes): + shown = f"{shape:<{shape_width}}" if show_shard else shape + row = f" {tensor.name:<{name_width}} {kind(tensor):<{type_width}} {shown}" + if show_shard: + row += f" {tensor.file}" + print(row) + + +def inspect(path: Path, metadata_dump: bool, tensor_dump: bool) -> None: + checkpoint = load_checkpoint(path) + print_summary(checkpoint) + if metadata_dump: + print_metadata(checkpoint) + if tensor_dump: + print_tensors(checkpoint) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Print architecture, quantization, and tensor info from a Hugging Face checkpoint." + ) + parser.add_argument("model", nargs="+", type=Path, help="A model directory or a .safetensors file") + parser.add_argument( + "--metadata", + action="store_true", + help="Print config, tokenizer, and safetensors header metadata. Long values are shortened.", + ) + parser.add_argument("--tensors", action="store_true", help="Print every tensor name, type, and shape.") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + for index, path in enumerate(args.model): + if index: + print() + inspect(path, args.metadata, args.tensors) + + +if __name__ == "__main__": + try: + main() + except BrokenPipeError: + # head and less close the pipe early. Don't turn that into a traceback. + try: + sys.stdout.close() + except BrokenPipeError: + pass + try: + sys.stderr.close() + except BrokenPipeError: + pass + os._exit(0) diff --git a/hf_to_nvfp4_gguf.py b/hf_to_nvfp4_gguf.py new file mode 100644 index 0000000..e2c0011 --- /dev/null +++ b/hf_to_nvfp4_gguf.py @@ -0,0 +1,1118 @@ +#!/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 "" in known: + stops = [""] + [token for token in stops if token != ""] + 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: ./-NVFP4, or ./ 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-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() diff --git a/hf_to_nvfp4_safetensors.py b/hf_to_nvfp4_safetensors.py new file mode 100644 index 0000000..3e72042 --- /dev/null +++ b/hf_to_nvfp4_safetensors.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 +"""Quantize a Hugging Face directory to one ComfyUI NVFP4 safetensors file. + +convert_to_quant reads one safetensors file and writes one file with packed +NVFP4 weights, ``weight_scale`` / ``weight_scale_2``, ``.comfy_quant`` +tensors, and a ``_quantization_metadata`` header. A checkpoint is usually +sharded. This script runs that same conversion on every shard, then merges +the results into a single ``.safetensors`` file. + +A bias is corrected only when it lives in the same shard as its weight, +which is how convert_to_quant sees it. Norms, embeddings, ``lm_head``, and +the other names on convert_to_quant's exclusion list stay in their original +dtype. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import struct +from pathlib import Path + +from nvfp4gguf import ( + CTQ_NUM_ITER, + DEFAULT_CONVERT_TO_QUANT, + DEFAULT_EARLY_STOP_STALL, + DEFAULT_SCALE_OPTIMIZATION, + DEFAULT_SCALE_REFINEMENT_ROUNDS, + load_learned_converter, + quant_profile, + resolve_optimizer, +) + +INDEX_NAME = "model.safetensors.index.json" +# safetensors aligns each tensor to its dtype width, not to 8. +_DTYPE_ALIGN = { + "BOOL": 1, "U8": 1, "I8": 1, + "F8_E4M3": 1, "F8_E5M2": 1, "F8_E8M0": 1, + "F16": 2, "BF16": 2, "I16": 2, "U16": 2, + "F32": 4, "I32": 4, "U32": 4, + "F64": 8, "I64": 8, "U64": 8, +} + + +def discover_shards(model_dir: Path) -> list[str]: + """Return safetensors shard filenames in the checkpoint directory.""" + index_path = model_dir / INDEX_NAME + if index_path.is_file(): + payload = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = payload.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise SystemExit(f"{index_path} has no weight_map") + names = sorted(set(weight_map.values())) + missing = [name for name in names if not (model_dir / name).is_file()] + if missing: + raise SystemExit(f"Index entries are missing from {model_dir}: {', '.join(missing)}") + return names + + single = model_dir / "model.safetensors" + if single.is_file(): + return ["model.safetensors"] + + shards = sorted(path.name for path in model_dir.glob("*.safetensors") if path.is_file()) + if not shards: + raise SystemExit(f"No safetensors shards in {model_dir}") + return shards + + +def output_file(args, source: Path) -> Path: + if not args.output: + return (Path.cwd() / f"{source.name}-NVFP4.safetensors").resolve() + output = Path(args.output) + if output.exists() and output.is_dir(): + return (output / f"{source.name}-NVFP4.safetensors").resolve() + if output.suffix != ".safetensors": + output = output.with_name(output.name + ".safetensors") + return output.resolve() + + +def read_header(path: Path) -> tuple[dict, int]: + with path.open("rb") as handle: + raw_size = handle.read(8) + if len(raw_size) != 8: + raise SystemExit(f"{path} is not a safetensors file") + header_size = struct.unpack(" None: + """Concatenate quantized shards into one safetensors file. + + Tensor bytes are copied, not loaded. ``_quantization_metadata`` layers + from each shard are combined so ComfyUI still sees every quantized weight. + """ + metadata: dict[str, str] = {} + layers: dict = {} + pieces: list[tuple[str, dict, Path, int, int]] = [] + owners: dict[str, str] = {} + for path in paths: + header, header_size = read_header(path) + file_meta = dict(header.pop("__metadata__", {}) or {}) + raw_quant = file_meta.pop("_quantization_metadata", None) + if raw_quant: + parsed = json.loads(raw_quant) + file_layers = parsed.get("layers", {}) + overlap = sorted(set(file_layers) & set(layers)) + if overlap: + raise SystemExit(f"Quant metadata for {overlap[0]!r} appears in more than one shard") + layers.update(file_layers) + for key, value in file_meta.items(): + metadata.setdefault(key, value) + data_base = 8 + header_size + for key, entry in header.items(): + if key in owners: + raise SystemExit(f"Tensor {key!r} appears in both {owners[key]} and {path.name}") + owners[key] = path.name + start, end = entry["data_offsets"] + pieces.append((key, entry, path, data_base + int(start), int(end) - int(start))) + + if layers: + metadata["_quantization_metadata"] = json.dumps( + {"format_version": "1.0", "layers": layers}, + separators=(",", ":"), + ) + + # safetensors rejects gaps. Grouping by descending dtype width keeps every + # tensor aligned without padding, which is how save_file lays a file out. + ordered = sorted( + enumerate(pieces), + key=lambda item: (-_DTYPE_ALIGN.get(item[1][1].get("dtype", ""), 8), item[0]), + ) + cursor = 0 + new_header: dict = {} + if metadata: + new_header["__metadata__"] = metadata + layout: list[tuple[Path, int, int]] = [] + for _index, (key, entry, path, abs_start, length) in ordered: + align = _DTYPE_ALIGN.get(entry.get("dtype", ""), 8) + if cursor % align != 0: + raise SystemExit(f"Cannot align {key} ({entry.get('dtype')}) at offset {cursor}") + info = {name: value for name, value in entry.items() if name != "data_offsets"} + info["data_offsets"] = [cursor, cursor + length] + new_header[key] = info + layout.append((path, abs_start, length)) + cursor += length + + header_bytes = json.dumps(new_header, separators=(",", ":")).encode("utf-8") + header_bytes += b" " * ((8 - (len(header_bytes) % 8)) % 8) + dest.parent.mkdir(parents=True, exist_ok=True) + temporary = dest.with_name(dest.name + ".tmp") + with temporary.open("wb") as out: + out.write(struct.pack(" None: + if destination.exists(): + destination.unlink() + convert_to_nvfp4( + str(source), + str(destination), + exclude_layers=args.exclude_layers, + output_dtype=args.output_dtype, + simple=args.simple, + num_iter=args.num_iter, + heur=args.heur, + calib_samples=args.calib_samples, + seed=args.seed, + optimizer=optimizer, + lr=args.lr, + lr_schedule=args.lr_schedule, + top_p=args.top_p, + min_k=args.min_k, + max_k=args.max_k, + full_matrix=args.full_matrix, + early_stop_loss=args.early_stop_loss, + early_stop_stall=args.early_stop_stall, + scale_refinement_rounds=args.scale_refinement_rounds, + scale_optimization=args.scale_optimization, + low_memory=args.low_memory, + use_speed=args.use_speed, + verbose_output=args.verbose, + ) + if not destination.is_file() or destination.stat().st_size == 0: + raise SystemExit(f"convert_to_quant did not write {destination}") + + +def _apply_directory_profile(args, source: Path) -> None: + config_path = source / "config.json" + if not config_path.is_file(): + return + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return + names: list[str] = [] + architectures = config.get("architectures") + if isinstance(architectures, list): + names.extend(str(item) for item in architectures) + elif isinstance(architectures, str): + names.append(architectures) + model_type = config.get("model_type") + if isinstance(model_type, str) and model_type: + names.append(model_type) + for name in names: + profile = quant_profile(name) + if profile is None: + continue + if not args.exclude_layers: + args.exclude_layers = profile.skip + print(f"{profile.arch} profile: leaving {profile.label} unquantized.") + else: + print(f"{profile.arch} profile: --exclude-layers replaces the default skip of {profile.label}.") + return + + +def convert_directory(args) -> None: + source = Path(args.model).resolve() + _apply_directory_profile(args, source) + if not source.is_dir(): + raise SystemExit(f"Not a directory: {source}") + output = output_file(args, source) + if output.is_relative_to(source): + raise SystemExit("Write the safetensors file outside the input directory.") + if output.exists() and not args.overwrite: + raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.") + shards = discover_shards(source) + + load_learned_converter(args.convert_to_quant) + from convert_to_quant.formats.nvfp4_conversion import convert_to_nvfp4 + from convert_to_quant.utils.logging import setup_logging + + # Without this, convert_to_quant drops the "(13/256) Processing tensor" lines + # and only the inner 4000-step bar is visible. + setup_logging("VERBOSE" if args.verbose else "NORMAL") + + optimizer = resolve_optimizer(args.optimizer) + print(f"Quantizing {len(shards)} shard(s) from {source}") + print(f"Output file: {output}") + if len(shards) == 1: + quantize_shard(convert_to_nvfp4, args, optimizer, source / shards[0], output) + print(f"Wrote {output}") + return + + work = output.with_name(output.name + ".shards") + if work.exists(): + shutil.rmtree(work) + work.mkdir(parents=True) + try: + written: list[Path] = [] + for index, name in enumerate(shards, start=1): + destination = work / name + print(f"\n[{index}/{len(shards)}] {name}") + quantize_shard(convert_to_nvfp4, args, optimizer, source / name, destination) + written.append(destination) + print(f"\nMerging {len(written)} shards into {output.name}") + merge_safetensors(written, output) + finally: + shutil.rmtree(work, ignore_errors=True) + print(f"Wrote {output}") + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Quantize a Hugging Face directory to one ComfyUI NVFP4 safetensors file.", + ) + parser.add_argument("model", nargs="?", help="Hugging Face model directory (config.json + safetensors)") + parser.add_argument("-o", "--output", help="Output .safetensors file. Default: ./-NVFP4.safetensors") + parser.add_argument("--overwrite", action="store_true", help="Replace an existing output file") + parser.add_argument("--convert-to-quant", type=Path, default=None, + help="convert_to_quant checkout, if it is not installed") + 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-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 progress. At the learning-rate floor, progress means a 1% loss drop (default: 2000)", + ) + parser.add_argument("--calib-samples", type=int, default=3072) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--use-speed", action="store_true", help="Enable Prodigy's use_speed") + parser.add_argument("--low-memory", action="store_true", help="Stream each shard instead of loading it all at once") + parser.add_argument("--heur", action="store_true", help="Skip layers convert_to_quant's size heuristic rejects") + parser.add_argument("--exclude-layers", default=None, help="Regex of tensor names to leave unquantized") + parser.add_argument("--output-dtype", default="bfloat16", choices=("bfloat16", "float16"), + help="Dtype recorded for quantized layers and used for unquantized 2D weights") + parser.add_argument("--verbose", action="store_true") + args = parser.parse_args(argv) + if not args.model: + parser.error("model directory is required") + 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: + convert_directory(parse_args(argv)) + + +if __name__ == "__main__": + main() diff --git a/llama.cpp b/llama.cpp new file mode 160000 index 0000000..63b4b30 --- /dev/null +++ b/llama.cpp @@ -0,0 +1 @@ +Subproject commit 63b4b303bd87456adc2f09d56c7e1f74e9e26673 diff --git a/nvfp4gguf/__init__.py b/nvfp4gguf/__init__.py new file mode 100644 index 0000000..a6cae77 --- /dev/null +++ b/nvfp4gguf/__init__.py @@ -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", +] diff --git a/nvfp4gguf/__pycache__/__init__.cpython-313.pyc b/nvfp4gguf/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..0eaec0e Binary files /dev/null and b/nvfp4gguf/__pycache__/__init__.cpython-313.pyc differ diff --git a/nvfp4gguf/__pycache__/bundle.cpython-313.pyc b/nvfp4gguf/__pycache__/bundle.cpython-313.pyc new file mode 100644 index 0000000..7aca80c Binary files /dev/null and b/nvfp4gguf/__pycache__/bundle.cpython-313.pyc differ diff --git a/nvfp4gguf/__pycache__/quant.cpython-313.pyc b/nvfp4gguf/__pycache__/quant.cpython-313.pyc new file mode 100644 index 0000000..e65e60f Binary files /dev/null and b/nvfp4gguf/__pycache__/quant.cpython-313.pyc differ diff --git a/nvfp4gguf/bundle.py b/nvfp4gguf/bundle.py new file mode 100644 index 0000000..7d53e46 --- /dev/null +++ b/nvfp4gguf/bundle.py @@ -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 diff --git a/nvfp4gguf/quant.py b/nvfp4gguf/quant.py new file mode 100644 index 0000000..1e048f2 --- /dev/null +++ b/nvfp4gguf/quant.py @@ -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") diff --git a/ollama_nvfp4.py b/ollama_nvfp4.py new file mode 100644 index 0000000..b496eb7 --- /dev/null +++ b/ollama_nvfp4.py @@ -0,0 +1,1838 @@ +#!/usr/bin/env python3 +"""List Ollama's sha256 model store in readable form, and quantize a model to NVFP4. + +Ollama keeps weights in blobs/sha256- and names them through JSON manifests. +`list` and `show` resolve those hashes. `list` prints a short capability +code for each model. `list --full-scan` also reads GGUF headers, which is +slower, and uses them to find encoders and capabilities that the manifest +does not name. `list --ctx-scan` reads each parameter blob and names the +models whose Modelfile sets num_ctx. `--info MODEL` runs gguf_info.py on +that model's sha256 weight blob. Arguments after the model name are passed +through, so `--info MODEL --metadata` is gguf_info.py on the blob with +`--metadata`. `export` writes the original GGUF and a Modelfile into a +directory. `quant` does the same after fitting NVFP4, and can +register the result as a new Ollama model. The Modelfile keeps the source +template, stop list, system prompt, renderer, and parser, which is what +Ollama uses for thinking and for stopping generation. + +Most Ollama blobs are already quantized (Q4_K_M, Q8_0, and so on). NVFP4 is then fit +to those dequantized weights, so it cannot recover detail the source quant discarded. +A float16 or BF16 blob is the better source. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import logging +import os +import re +import shutil +import struct +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path + +log = logging.getLogger("ollama_nvfp4") + +MODEL_MEDIA = "application/vnd.ollama.image.model" +CONFIG_MEDIA = "application/vnd.docker.container.image.v1+json" +TEMPLATE_MEDIA = "application/vnd.ollama.image.template" +MESSAGES_MEDIA = "application/vnd.ollama.image.messages" +ADAPTER_MEDIA = "application/vnd.ollama.image.adapter" +PROJECTOR_MEDIA = "application/vnd.ollama.image.projector" +EMBED_MEDIA = "application/vnd.ollama.image.embed" +PARAMS_MEDIA = "application/vnd.ollama.image.params" +LAYER_KIND = { + TEMPLATE_MEDIA: "TEMPLATE", + "application/vnd.ollama.image.prompt": "TEMPLATE", + "application/vnd.ollama.image.system": "SYSTEM", + "application/vnd.ollama.image.license": "LICENSE", + PARAMS_MEDIA: "PARAMETER", + MESSAGES_MEDIA: "MESSAGE", + PROJECTOR_MEDIA: "PROJECTOR", + ADAPTER_MEDIA: "ADAPTER", +} +FLOAT_TYPES = {"F32", "F16", "BF16"} + + +def default_models_dir() -> Path: + env = os.environ.get("OLLAMA_MODELS") + if env: + return Path(env) + system = Path("/models/ollama") + if (system / "blobs").is_dir(): + return system + return Path.home() / ".ollama" / "models" + + +def digest_to_name(digest: str) -> str: + return digest.replace(":", "-") if digest.startswith("sha256:") else f"sha256-{digest}" + + +def blob_path(root: Path, digest: str) -> Path: + return root / "blobs" / digest_to_name(digest) + + +def format_size(num_bytes: int) -> str: + size = float(num_bytes) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if size < 1024 or unit == "TiB": + return f"{int(size)} {unit}" if unit == "B" else f"{size:.1f} {unit}" + size /= 1024 + return f"{size:.1f} TiB" + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def read_json(path: Path) -> dict: + with path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + if not isinstance(data, dict): + raise ValueError(f"{path} is not a JSON object") + return data + + +@dataclass +class OllamaModel: + name: str + manifest_path: Path + store: str + config_digest: str + config: dict + layers: list[dict] + model_digest: str + model_size: int + gguf_path: Path + metadata: dict = field(default_factory=dict) + runner: str = "" + manifest_format: str = "" + + @property + def family(self) -> str: + families = self.config.get("model_families") or [] + if families: + return str(families[0]) + return str(self.config.get("model_family") or "") + + @property + def params(self) -> str: + return str(self.config.get("model_type") or self.metadata.get("general.size_label") or "") + + @property + def quant(self) -> str: + if self.config.get("file_type"): + return str(self.config["file_type"]) + file_type = self.metadata.get("general.file_type") + return "" if file_type is None else str(file_type) + + @property + def context(self) -> str: + found = _context_from_metadata(self.metadata) + if found: + return found + # Ollama's metadata cache is often mode 600 and owned by root. The GGUF + # header is readable and is what /api/show uses. + return _context_from_gguf(self.gguf_path) + + +def _is_model_context_key(name: str) -> bool: + return name.endswith(".context_length") and ".rope." not in name + + +def _scalar_context(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 _context_from_metadata(metadata: dict) -> str: + for key, value in metadata.items(): + if not _is_model_context_key(key): + continue + length = _scalar_context(value) + if length is not None: + return str(length) + return "" + + +def _context_from_gguf(path: Path) -> str: + """Read only the GGUF header. GGUFReader maps the whole weight file.""" + if not path.is_file(): + return "" + scalar = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8} + integer = {4: " bytes: + data = handle.read(n) + if len(data) != n: + raise ValueError("truncated GGUF header") + return data + + def read_string() -> bytes: + size = struct.unpack(" 16 * 1024 * 1024: + raise ValueError("GGUF string is too long") + return read(size) + + def skip(value_type: int) -> None: + if value_type in scalar: + read(scalar[value_type]) + return + if value_type == 8: + read_string() + return + if value_type != 9: + raise ValueError("unknown GGUF value type") + subtype = struct.unpack(" tuple[str, str]: + for store in ("manifests", "manifests-v2"): + base = root / store + try: + parts = manifest.relative_to(base).parts + except ValueError: + continue + if len(parts) < 3: + break + tag = parts[-1] + model = parts[-2] + namespace = parts[-3] + name = f"{model}:{tag}" if namespace == "library" else f"{namespace}/{model}:{tag}" + return name, store + return manifest.name, "loose" + + +def iter_manifests(root: Path): + for store in ("manifests", "manifests-v2"): + base = root / store + if not base.is_dir(): + continue + for path in sorted(base.rglob("*")): + if path.is_file(): + yield path + + +def load_models(root: Path) -> list[OllamaModel]: + models: list[OllamaModel] = [] + for manifest_path in iter_manifests(root): + try: + manifest = read_json(manifest_path) + except (OSError, json.JSONDecodeError, ValueError): + continue + layers = manifest.get("layers") or [] + model_layer = next((layer for layer in layers if layer.get("mediaType") == MODEL_MEDIA), None) + config_info = manifest.get("config") or {} + if model_layer is None or "digest" not in config_info: + continue + config_path = blob_path(root, config_info["digest"]) + try: + config = read_json(config_path) if config_path.is_file() else {} + except (OSError, json.JSONDecodeError, ValueError): + config = {} + meta_path = root / "metadata" / digest_to_name(model_layer["digest"]) + meta_file = meta_path.with_suffix(".json") + metadata = {} + if meta_file.is_file(): + try: + metadata = read_json(meta_file).get("kv") or {} + except (OSError, json.JSONDecodeError, ValueError): + metadata = {} + name, store = manifest_name(root, manifest_path) + models.append(OllamaModel( + name=name, + manifest_path=manifest_path, + store=store, + config_digest=config_info["digest"], + config=config, + layers=layers, + model_digest=model_layer["digest"], + model_size=int(model_layer.get("size") or 0), + gguf_path=blob_path(root, model_layer["digest"]), + metadata=metadata, + runner=str(manifest.get("runner") or ""), + manifest_format=str(manifest.get("format") or ""), + )) + models.sort(key=lambda item: item.name.lower()) + return models + + +def referenced_digests(root: Path) -> set[str]: + found: set[str] = set() + for manifest_path in iter_manifests(root): + try: + manifest = read_json(manifest_path) + except (OSError, json.JSONDecodeError, ValueError): + continue + config = manifest.get("config") or {} + if isinstance(config.get("digest"), str): + found.add(config["digest"]) + for layer in manifest.get("layers") or []: + if isinstance(layer.get("digest"), str): + found.add(layer["digest"]) + return found + + +def resolve_model(models: list[OllamaModel], query: str) -> OllamaModel: + exact = [model for model in models if model.name == query] + if len(exact) == 1: + return exact[0] + if len(exact) > 1: + raise SystemExit("More than one model is named " + query + ":\n" + "\n".join(model.manifest_path.as_posix() for model in exact)) + folded = query.lower() + partial = [model for model in models if folded in model.name.lower() or folded in model.model_digest.lower()] + if len(partial) == 1: + return partial[0] + if not partial: + raise SystemExit(f"No Ollama model matches {query!r}. Run list to see names.") + preview = "\n".join(f" {model.name}" for model in partial[:20]) + raise SystemExit(f"{query!r} matches {len(partial)} models:\n{preview}") + + +# Letters match Ollama's capability set. X is text because T is tool calls. +CAPABILITY_LETTERS = ( + ("completion", "X", "text"), + ("tools", "T", "tool calls"), + ("vision", "V", "vision (images and video)"), + ("audio", "A", "audio input"), + ("image", "I", "image generation"), + ("embedding", "E", "embeddings"), + ("thinking", "R", "thinking"), + ("insert", "M", "infill (fill-in-the-middle)"), + ("decision", "D", "decision (classification, not chat)"), + ("speech", "S", "speech output"), +) +# Families and architectures that always need a vision encoder. Gemma 3 is +# not here: its 270M and 1B builds are text, and the larger ones are not. +_VISION_IDS = { + "qwen2vl", "qwen25vl", "qwen2.5vl", "qwen3vl", "qwen3vlmoe", + "llama4", "mllama", "llava", "bakllava", "moondream", + "minicpm-v", "minicpmv", "mistral3", "pixtral", "internvl", + "cogvlm", "deepseek2-ocr", "glmocr", "glm-ocr", "gemma3n", + "hunyuan_vl", "paddleocr", "sensenova_u1", "llama3.2-vision", +} +_EMBED_IDS = { + "nomic-bert", "nomic-bert-moe", "modern-bert", "neo-bert", + "jina-bert-v2", "jina-bert-v3", "eurobert", + "gemma-embedding", "gemma-embedding2", "llama-embed", "pangu-embedded", +} +_SPEECH_ARCH = {"qwen3tts", "pockettts", "talkie", "wavtokenizer-dec"} +_THINKING_PARSERS = { + "qwen3-thinking", "qwen3-vl-thinking", "deepseek3", "gemma4", + "lfm2-thinking", "olmo3-think", "harmony", +} +_TOOL_PARSERS = {"qwen3-coder", "functiongemma", "harmony"} +_AUDIO_PROJECTORS = { + "qwen2a", "qwen3a", "glma", "voxtral", "voxtral_rt", "meralion", + "lfm2a", "musicflamingo", "granite_speech", "mimo_audio", "parakeet", + "dots3note_a", "d1omni_a", "muse-spark-audio", "ultravox", + "qwen2.5o", "nemotron_omni", +} +_SPEECH_PROJECTORS = { + "qwen3tts_spkenc", "qwen3tts_gen", "pockettts_spkenc", "pockettts_gen", "soprano", +} +_VISION_PROJECTORS = { + "gemma3", "gemma3nv", "gemma3na", "gemma4v", "gemma4a", "gemma4uv", "gemma4ua", + "phi4", "idefics3", "pixtral", "llama4", "qwen2vl_merger", "qwen2.5vl_merger", + "exaone4_5", "qwen3vl_merger", "ling3vl", "step3vl", "internvl", "lfm2", + "lfm3vl", "kimivl", "paddleocr", "kimik25", "lightonocr", "cogvlm", + "janus_pro", "dots_ocr", "dots3note_v", "deepseekocr", "deepseekocr2", + "deepseek4v", "deepseek41v", "glm4v", "glm5v", "youtuvl", "nemotron_v2_vl", + "nemotron_omni", "qwen2.5o", "hunyuanvl", "minimax_m3", "minicpmv4_6", + "minicpmv4_7", "minicpm_ocr_1_0", "mimovl", "granite4_vision", "muse-glimmer", + "cohere2v", "d1omni_v", +} +_TEMPLATE_MEDIAS = { + TEMPLATE_MEDIA, + "application/vnd.ollama.image.prompt", +} +_GEMMA3_VISION_SIZE = re.compile(r"(^|[^0-9.])(4(?:\.3)?|12|27)\s*b\b") +# Shared chat templates often contain an image or audio branch even for a text +# checkpoint. A real requirement comes from the name, the architecture, or a +# projector the manifest actually lists. +_NAME_VISION = re.compile(r"(?:^|[^a-z0-9])(?:vl|vision|ocr)(?:[^a-z0-9]|$)") +_NAME_AUDIO = re.compile(r"(?:^|[^a-z0-9])audio(?:[^a-z0-9]|$)") +_GGUF_CACHE: dict[tuple[str, int, int], "GGUFProfile"] = {} + + +@dataclass +class GGUFProfile: + context: str = "" + architecture: str = "" + general_name: str = "" + general_type: str = "" + pooling: bool = False + decision: bool = False + vision_blocks: bool = False + audio_blocks: bool = False + has_vision_encoder: bool | None = None + has_audio_encoder: bool | None = None + projector_type: str = "" + image_tokens: bool = False + video_tokens: bool = False + audio_tokens: bool = False + templates: list[str] = field(default_factory=list) + vision_tensors: bool = False + audio_tensors: bool = False + speech_tensors: bool = False + readable: bool = False + + +@dataclass +class CapabilityReport: + marks: list[tuple[str, str, bool]] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + projector: str = "-" + context: str = "" + + def letters(self) -> str: + return " ".join(letter if ok else f"{letter}!" for letter, _label, ok in self.marks) + + def column(self) -> str: + text = self.letters() + if self.notes: + return f"{text} {'; '.join(self.notes)}" if text else "; ".join(self.notes) + return text + + def labels(self) -> str: + rendered = [] + for _letter, label, ok in self.marks: + rendered.append(label if ok else f"{label}, weights missing") + return ", ".join(rendered) + + +def _gguf_cache_key(path: Path) -> tuple[str, int, int] | None: + try: + stat = path.stat() + except OSError: + return None + return (str(path), stat.st_mtime_ns, stat.st_size) + + +def _note_gguf_key(profile: GGUFProfile, key: str, value) -> None: + low = key.lower() + if low == "general.architecture" and isinstance(value, str): + profile.architecture = value + elif low == "general.name" and isinstance(value, str): + profile.general_name = value + elif low == "general.type" and isinstance(value, str): + profile.general_type = value + elif low.endswith("pooling_type"): + profile.pooling = True + elif low == "decision.type" or low.endswith(".decision.type"): + profile.decision = isinstance(value, str) and bool(value.strip()) + elif low.endswith("vision.block_count"): + profile.vision_blocks = True + elif low.endswith("audio.block_count"): + profile.audio_blocks = True + elif low.endswith("has_vision_encoder") and isinstance(value, int): + profile.has_vision_encoder = value != 0 + elif low.endswith("has_audio_encoder") and isinstance(value, int): + profile.has_audio_encoder = value != 0 + elif low.endswith("projector_type") and isinstance(value, str): + profile.projector_type = value + elif low.endswith("image_token_id"): + profile.image_tokens = True + elif low.endswith("video_token_id"): + profile.video_tokens = True + elif low.endswith("audio_token_id"): + profile.audio_tokens = True + elif "chat_template" in low and isinstance(value, str) and value: + profile.templates.append(value[:500_000]) + elif _is_model_context_key(key): + length = _scalar_context(value) + if length is not None and not profile.context: + profile.context = str(length) + + +def _note_tensor_name(profile: GGUFProfile, name: str) -> None: + low = name.lower() + head, _, rest = low.partition(".") + if head in {"v", "vis", "vision", "mm", "mmproj"}: + profile.vision_tensors = True + elif head in {"audio", "gen_audio"} or low.startswith("a.blk"): + profile.audio_tensors = True + elif head in {"tts", "spk", "speaker", "codec"} or rest.startswith("spk"): + profile.speech_tensors = True + + +def _scan_gguf(path: Path) -> GGUFProfile: + """Read a GGUF header for capability metadata. Does not map the weights.""" + key = _gguf_cache_key(path) + if key is None: + return GGUFProfile() + cached = _GGUF_CACHE.get(key) + if cached is not None: + return cached + profile = _scan_gguf_uncached(path) + _GGUF_CACHE[key] = profile + return profile + + +def _scan_gguf_uncached(path: Path) -> GGUFProfile: + profile = GGUFProfile() + scalar = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8} + integer = {0: " bytes: + data = handle.read(n) + if len(data) != n: + raise ValueError("truncated GGUF header") + return data + + def read_string() -> bytes: + size = struct.unpack(" 16 * 1024 * 1024: + raise ValueError("GGUF string is too long") + return read(size) + + def read_scalar(value_type: int): + if value_type == 8: + return read_string().decode("utf-8", "replace") + if value_type not in scalar: + raise ValueError("unknown GGUF value type") + raw = read(scalar[value_type]) + if value_type in integer: + return struct.unpack(integer[value_type], raw)[0] + return None + + def skip_array(subtype: int, count: int, keep_strings: bool) -> None: + if subtype in scalar: + handle.seek(scalar[subtype] * count, os.SEEK_CUR) + return + if subtype != 8: + raise ValueError("unsupported GGUF array") + kept = 0 + for _ in range(count): + text = read_string() + if keep_strings and kept < 4: + profile.templates.append(text.decode("utf-8", "replace")[:500_000]) + kept += 1 + + try: + with path.open("rb") as handle: + magic, version = struct.unpack(" set[str]: + raw = config.get("capabilities") or [] + if isinstance(raw, str): + raw = [raw] + if not isinstance(raw, list): + return set() + return {str(item).strip().lower() for item in raw if str(item).strip()} + + +def _model_ids(model: OllamaModel, main: GGUFProfile) -> set[str]: + found = set() + if model.family: + found.add(model.family.lower()) + for item in model.config.get("model_families") or []: + found.add(str(item).lower()) + if main.architecture: + found.add(main.architecture.lower()) + return found + + +def _template_capabilities(text: str) -> set[str]: + if not text: + return set() + low = text.lower() + found: set[str] = set() + if "tools" in low or "tool_call" in low or ".tools" in low: + found.add("tools") + if "" in low or "" in low or "enable_thinking" in low or "<|think|>" in low: + found.add("thinking") + if ".suffix" in low: + found.add("insert") + return found + + +def _name_modalities(text: str) -> set[str]: + """Vision or audio claimed by the model name, not by a shared chat template.""" + low = text.lower() + found: set[str] = set() + if _NAME_VISION.search(low): + found.add("vision") + if _NAME_AUDIO.search(low): + found.add("audio") + return found + + +def _blob_text(path: Path, limit: int = 2_000_000) -> str: + if not path.is_file(): + return "" + try: + with path.open("r", encoding="utf-8", errors="replace") as handle: + return handle.read(limit) + except OSError: + return "" + + +def _looks_like_embedding(model: OllamaModel) -> bool: + if _model_ids(model, GGUFProfile()) & _EMBED_IDS: + return True + return "embed" in f"{model.name} {model.family}".lower() + + +def _gemma3_expects_vision(model: OllamaModel, main: GGUFProfile, ids: set[str]) -> bool: + if "gemma3n" in ids: + return False + if "gemma3" not in ids and "gemma3" not in model.name.lower(): + return False + text = " ".join((model.params, model.name, main.general_name)) + return _GEMMA3_VISION_SIZE.search(text.lower()) is not None + + +def _projector_modes(profile: GGUFProfile) -> tuple[bool, bool, bool]: + """Return vision, audio, speech for one projector header.""" + ptype = (profile.projector_type or "").lower() + speech = ptype in _SPEECH_PROJECTORS or profile.speech_tensors + audio = bool( + profile.has_audio_encoder + or profile.audio_blocks + or profile.audio_tensors + or profile.audio_tokens + or ptype in _AUDIO_PROJECTORS + ) + if ptype == "gemma3nv": + audio = False + if profile.has_vision_encoder is True: + vision = True + elif profile.has_vision_encoder is False: + vision = False + elif profile.vision_blocks or profile.vision_tensors or profile.image_tokens or profile.video_tokens: + vision = True + elif ptype in _VISION_PROJECTORS: + vision = True + elif audio or speech: + vision = False + else: + vision = True + return vision, audio, speech + + +def assess_capabilities( + model: OllamaModel, + *, + full: bool, + weights_present: bool, + main: GGUFProfile, + projectors: list[tuple[bool, GGUFProfile | None]], + templates: list[str], +) -> CapabilityReport: + """Decide which capabilities a model has, and which are missing their weights. + + `full` is the GGUF header scan. Without it, only the manifest, config, + parser, and small template blobs are used, so a packed-in encoder stays + invisible and a missing sidecar is reported only when the manifest names it. + """ + explicit = _config_capabilities(model.config) + ids = _model_ids(model, main if full else GGUFProfile()) + scanned = full and main.readable + present_projectors = [(ok, profile) for ok, profile in projectors if ok and profile is not None] + missing_projector = any(not ok for ok, _profile in projectors) + readable_modes = [_projector_modes(profile) for _ok, profile in present_projectors if profile.readable] + # A projector file that is on disk but unreadable still counts as present. + opaque_projector = any(ok and (profile is None or not profile.readable) for ok, profile in projectors) + + wanted: set[str] = set(explicit) + template_found: set[str] = set() + for text in templates: + template_found |= _template_capabilities(text) + if scanned: + for text in main.templates: + template_found |= _template_capabilities(text) + wanted |= template_found + wanted |= _name_modalities(model.name) + if scanned: + wanted |= _name_modalities(main.general_name) + + parser = str(model.config.get("parser") or "").lower() + renderer = str(model.config.get("renderer") or "").lower() + if parser in _TOOL_PARSERS or renderer in _TOOL_PARSERS: + wanted.add("tools") + if parser in _THINKING_PARSERS or renderer in _THINKING_PARSERS or ids & {"gptoss", "gpt-oss"}: + wanted.add("thinking") + + embed = "embedding" in explicit or (scanned and main.pooling) or (not scanned and _looks_like_embedding(model)) + decision = "decision" in explicit or (scanned and main.decision) + if embed: + wanted.add("embedding") + if decision: + wanted.add("decision") + if "completion" in explicit or (not embed and not decision): + wanted.add("completion") + + vision_inline = scanned and main.vision_tensors + audio_inline = scanned and main.audio_tensors + speech_inline = scanned and main.speech_tensors + projector_vision = any(vision for vision, _audio, _speech in readable_modes) or opaque_projector + projector_audio = any(audio for _vision, audio, _speech in readable_modes) + projector_speech = any(speech for _vision, _audio, speech in readable_modes) + + if projector_vision or missing_projector: + wanted.add("vision") + if scanned and (main.vision_blocks or main.image_tokens or main.video_tokens or vision_inline): + wanted.add("vision") + if scanned and (ids & _VISION_IDS or _gemma3_expects_vision(model, main, ids)): + wanted.add("vision") + if projector_audio or (scanned and (main.audio_blocks or main.audio_tokens or audio_inline)): + wanted.add("audio") + if "image" in explicit or (scanned and main.general_type.lower() in {"image", "flux", "sdxl"}): + wanted.add("image") + if projector_speech or speech_inline or (scanned and main.architecture.lower() in _SPEECH_ARCH): + wanted.add("speech") + if "speech" in explicit: + wanted.add("speech") + + if decision: + wanted.difference_update({"completion", "tools", "thinking", "insert"}) + + # Audio sidecars are separate files. A name that says audio, with no projector + # on disk, is missing its encoder even before the GGUF header is read. + # Vision is often packed into the model file, so a vision name stays + # unmarked until --full-scan looks for those tensors. + name_claims_audio = "audio" in _name_modalities(model.name) or ( + scanned and "audio" in _name_modalities(main.general_name) + ) + + have = set(wanted) + notes: list[str] = [] + if not weights_present: + have.clear() + notes.append("model blob missing") + else: + # Without a header scan, declared vision can be packed into the model + # file, so it is left unmarked. A projector blob the manifest names and + # cannot find is broken either way. A model whose name says audio is + # broken when that sidecar is absent, because the audio encoder is not + # packed into the language file. + if "vision" in wanted and not (vision_inline or projector_vision): + if scanned or missing_projector: + have.discard("vision") + if missing_projector: + notes.append("mmproj blob missing") + elif main.vision_blocks and not main.vision_tensors: + notes.append("vision encoder missing from the model file") + else: + notes.append("no mmproj") + if "audio" in wanted and not (audio_inline or projector_audio): + if scanned or missing_projector or name_claims_audio: + have.discard("audio") + if missing_projector and "mmproj blob missing" not in notes: + notes.append("mmproj blob missing") + elif scanned and main.audio_blocks and not main.audio_tensors: + notes.append("audio encoder missing from the model file") + elif "no mmproj" not in notes and "mmproj blob missing" not in notes: + notes.append("no audio mmproj") + if "speech" in wanted and not (speech_inline or projector_speech or opaque_projector): + needs_sidecar = scanned and main.architecture.lower() in _SPEECH_ARCH and not speech_inline + if needs_sidecar or missing_projector: + have.discard("speech") + if missing_projector and "mmproj blob missing" not in notes: + notes.append("mmproj blob missing") + elif "no mmproj" not in notes and "mmproj blob missing" not in notes: + notes.append("no speech mmproj") + + broken = [name for name in ("vision", "audio", "speech") if name in wanted and name not in have] + if missing_projector and not broken and "mmproj blob missing" not in notes: + notes.append("mmproj blob missing") + + present_count = sum(1 for ok, _profile in projectors if ok) + missing_count = len(projectors) - present_count + inline = vision_inline or audio_inline or speech_inline + if missing_count and present_count: + projector = f"yes, {missing_count} MISSING" + elif missing_count: + projector = "MISSING" + elif present_count > 1: + projector = f"yes x{present_count}" + elif present_count == 1: + projector = "yes+inline" if inline else "yes" + elif inline: + projector = "inline" + else: + projector = "-" + + marks = [] + for name, letter, label in CAPABILITY_LETTERS: + if name not in wanted: + continue + marks.append((letter, label, name in have and weights_present)) + return CapabilityReport(marks=marks, notes=notes, projector=projector, context=main.context) + + +def capability_report(root: Path, model: OllamaModel, full: bool) -> CapabilityReport: + weights_present = model.gguf_path.is_file() + projectors: list[tuple[bool, GGUFProfile | None]] = [] + templates: list[str] = [] + try: + for layer in model.layers: + media = str(layer.get("mediaType") or "") + digest = str(layer.get("digest") or "") + if not digest: + continue + path = blob_path(root, digest) + if media == PROJECTOR_MEDIA: + profile = _scan_gguf(path) if full and path.is_file() else None + projectors.append((path.is_file(), profile)) + elif media in _TEMPLATE_MEDIAS: + templates.append(_blob_text(path)) + main = _scan_gguf(model.gguf_path) if full and weights_present else GGUFProfile() + return assess_capabilities( + model, + full=full, + weights_present=weights_present, + main=main, + projectors=projectors, + templates=templates, + ) + except (OSError, ValueError): + return CapabilityReport(notes=["could not read capabilities"]) + + +def _report_context(model: OllamaModel, report: CapabilityReport, full: bool) -> str: + found = _context_from_metadata(model.metadata) + if found: + return found + if full and report.context: + return report.context + return _context_from_gguf(model.gguf_path) + + +def print_capability_legend(full: bool) -> None: + print() + print("CAPS " + " ".join(f"{letter} {label}" for _name, letter, label in CAPABILITY_LETTERS)) + print("A letter with ! has that capability, but the file that implements it is absent.") + print("MMPROJ yes = projector blob, inline = encoder is inside the model file, MISSING = the manifest names a blob that is absent, - = no projector file.") + if full: + print("Each GGUF header was read. An encoder packed into the model file shows up as inline. A capability whose encoder is absent is marked !.") + else: + print("Caps here come from the manifest, config, template blobs, and the model name.") + print("A model with audio in its name is marked A! when it has no projector. Vision from the name is left unmarked until --full-scan, because that encoder is sometimes inside the model file.") + print("--full-scan reads each GGUF header, which walks the vocabulary, and can add tools, thinking, and encoders stored only in that file.") + + +def modelfile_num_ctx(root: Path, model: OllamaModel) -> tuple[str, str]: + """Read PARAMETER num_ctx from the model's parameter blob. + + The first item is the value to print: the number, '-' when it is unset, + MISSING when the blob is gone, or '?' when the blob is not JSON. + The second item is '' unless this model needs the num_ctx removed. + """ + saw_params = False + value = None + for layer in model.layers: + if layer.get("mediaType") != PARAMS_MEDIA: + continue + digest = str(layer.get("digest") or "") + if not digest: + continue + saw_params = True + path = blob_path(root, digest) + if not path.is_file(): + return "MISSING", "" + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError): + return "?", "" + if isinstance(data, dict) and "num_ctx" in data: + value = data["num_ctx"] + if not saw_params or value is None: + return "-", "" + rendered = value if isinstance(value, str) else json.dumps(value, separators=(",", ":")) + return rendered, rendered + + +def cmd_list(root: Path, full_scan: bool = False, ctx_scan: bool = False) -> None: + models = load_models(root) + if not models: + print(f"No manifests in {root}") + return + rows = [] + needs_fix: list[tuple[str, str]] = [] + for model in models: + if full_scan: + print(f"Reading {model.name}...", file=sys.stderr) + report = capability_report(root, model, full_scan) + present = "MISSING" if not model.gguf_path.is_file() else format_size(model.model_size) + num_ctx = "-" + if ctx_scan: + num_ctx, fix = modelfile_num_ctx(root, model) + if fix: + needs_fix.append((model.name, fix)) + rows.append((model, present, _report_context(model, report, full_scan), report, num_ctx)) + name_width = max(len(model.name) for model, _present, _context, _report, _num_ctx in rows) + ctx_width = max([len("CTX")] + [len(context) for _model, _present, context, _report, _num_ctx in rows]) + proj_width = max([len("MMPROJ")] + [len(report.projector) for _model, _present, _context, report, _num_ctx in rows]) + num_width = max([len("NUM_CTX")] + [len(num_ctx) for *_rest, num_ctx in rows]) if ctx_scan else 0 + header = ( + f"{'NAME':<{name_width}} {'PARAMS':<8} {'QUANT':<12} {'FAMILY':<14} " + f"{'SIZE':>10} {'CTX':<{ctx_width}} {'MMPROJ':<{proj_width}} " + ) + if ctx_scan: + header += f"{'NUM_CTX':<{num_width}} " + header += "CAPS" + print(header) + print("-" * len(header)) + for model, present, context, report, num_ctx in rows: + line = ( + f"{model.name:<{name_width}} {model.params:<8} {model.quant:<12} " + f"{model.family:<14} {present:>10} {context:<{ctx_width}} " + f"{report.projector:<{proj_width}} " + ) + if ctx_scan: + line += f"{num_ctx:<{num_width}} " + line += report.column() + print(line) + blobs = list((root / "blobs").glob("sha256-*")) if (root / "blobs").is_dir() else [] + referenced = referenced_digests(root) + orphan = [path for path in blobs if "sha256:" + path.name.removeprefix("sha256-") not in referenced] + orphan_bytes = sum(path.stat().st_size for path in orphan) + print() + print(f"{len(models)} models, {len(blobs)} blobs, {len(orphan)} blobs not referenced by a manifest ({format_size(orphan_bytes)})") + if ctx_scan: + print() + if needs_fix: + noun = "model" if len(needs_fix) == 1 else "models" + print(f"{len(needs_fix)} {noun} set num_ctx in the Modelfile and need reinstalling without it:") + for name, value in needs_fix: + print(f" {name} {value}") + else: + print("No Modelfile sets num_ctx.") + print_capability_legend(full_scan) + + +def cmd_show(root: Path, query: str) -> None: + model = resolve_model(load_models(root), query) + report = capability_report(root, model, True) + print(model.name) + print(f" store {model.store}") + print(f" manifest {model.manifest_path}") + print(f" family {model.family}") + print(f" params {model.params}") + print(f" quant {model.quant}") + print(f" context {_report_context(model, report, True) or '-'}") + print(f" caps {report.letters() or '-'}") + if report.labels(): + print(f" cap detail {report.labels()}") + if report.notes: + print(f" broken {'; '.join(report.notes)}") + print(f" projector {report.projector}") + if model.config.get("renderer"): + print(f" renderer {model.config['renderer']}") + if model.config.get("parser"): + print(f" parser {model.config['parser']}") + if model.config.get("requires"): + print(f" requires {model.config['requires']}") + architecture = model.metadata.get("general.architecture", "") + if architecture: + print(f" arch {architecture}") + for suffix in ("block_count", "embedding_length", "expert_count", "vocab_size"): + key = f"{architecture}.{suffix}" + if key in model.metadata and not isinstance(model.metadata[key], list): + print(f" {suffix:<18} {model.metadata[key]}") + print(f" weights {model.model_digest}") + print(f" file {model.gguf_path} ({format_size(model.model_size)})") + print(" layers") + for layer in model.layers: + kind = str(layer.get("mediaType", "")).removeprefix("application/vnd.ollama.image.") + digest = str(layer.get("digest", "")) + short = digest[:19] + "…" if len(digest) > 19 else digest + path = blob_path(root, digest) + state = "ok" if path.is_file() else "MISSING" + print(f" {kind:<16} {format_size(int(layer.get('size') or 0)):>10} {short} {state}") + + +def _pythonize(value): + import numpy as np + if isinstance(value, np.generic): + return value.item() + if isinstance(value, list): + return [_pythonize(item) for item in value] + return value + + +def _logical_shape(tensor) -> tuple[int, ...]: + return tuple(int(dim) for dim in tensor.shape[::-1]) + + +def _to_float(tensor, gguf_module): + import numpy as np + qtype = tensor.tensor_type + if qtype == gguf_module.GGMLQuantizationType.F32: + return np.asarray(tensor.data, dtype=np.float32) + if qtype == gguf_module.GGMLQuantizationType.F16: + return np.asarray(tensor.data, dtype=np.float32) + return np.asarray(gguf_module.quants.dequantize(tensor.data, qtype), dtype=np.float32) + + +def _checkpoint_file(directory: Path, name: str) -> Path: + digest = hashlib.sha1(name.encode("utf-8")).hexdigest()[:16] + safe = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in name)[:80] + return directory / f"{safe}.{digest}.npz" + + +def _load_checkpoint(directory: Path, name: str, shape: tuple[int, ...]): + import numpy as np + path = _checkpoint_file(directory, name) + if not path.is_file(): + return None + with np.load(path) as blob: + 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) + return raw, None if scale.size == 0 else scale + + +def _save_checkpoint(directory: Path, name: str, shape: tuple[int, ...], raw, scale) -> None: + import numpy as np + path = _checkpoint_file(directory, name) + temporary = path.with_name(path.name + ".tmp") + with temporary.open("wb") as handle: + np.savez_compressed( + handle, + 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), + ) + temporary.replace(path) + + +def _copy_metadata(reader, writer, gguf_module, context_length: int | None = None) -> None: + from nvfp4gguf.bundle import is_model_context_key, normalize_chat_template, require_context_length, scalar_context_length + 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": + value = int(gguf_module.LlamaFileType.MOSTLY_NVFP4) + 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 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 _source_description(tensor) -> str: + return tensor.tensor_type.name + + +def quantize_gguf(args, source: Path, output: Path, quant_args) -> None: + from nvfp4gguf import ( + TensorQuantizer, + bind_exclude, + eligible_weight, + ensure_cuda_interpreter, + import_gguf, + load_learned_converter, + must_bake_nvfp4_scale, + mtp_layer_tensor, + preserve_scale_tensor, + quantize_and_pack, + report_reconstruction, + require_quant_settings, + scale_bake_notice, + ) + gguf = import_gguf(args.llama_cpp) + + ensure_cuda_interpreter(quant_args.device) + if output.exists() and not _wants_overwrite(args): + raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.") + output.parent.mkdir(parents=True, exist_ok=True) + reader = gguf.GGUFReader(source) + if reader.byte_order != "I": + raise SystemExit(f"{source} is not little-endian GGUF") + architecture = reader.fields["general.architecture"].contents() + if isinstance(architecture, bytes): + architecture = architecture.decode("utf-8", "replace") + else: + architecture = str(architecture) + if not hasattr(quant_args, "exclude"): + quant_args.exclude = None + bind_exclude(quant_args, architecture) + notice = scale_bake_notice(architecture) + if notice: + log.warning(notice) + file_type = reader.fields.get("general.file_type") + file_type_name = "" + if file_type is not None: + try: + file_type_name = gguf.LlamaFileType(int(file_type.contents())).name + except ValueError: + file_type_name = str(file_type.contents()) + float_source = file_type_name in {"ALL_F32", "MOSTLY_F16", "MOSTLY_BF16"} + if not float_source: + message = ( + f"Source file type is {file_type_name or 'quantized'}. " + "NVFP4 will be fit to the dequantized weights, so detail already discarded by the source quant stays discarded." + ) + if args.float_only: + raise SystemExit(message) + log.warning(message) + + names = {tensor.name for tensor in reader.tensors} + + def field_int(key: str) -> int: + field = reader.fields.get(key) + if field is None: + return 0 + try: + return int(field.contents()) + except (TypeError, ValueError): + return 0 + + n_blocks = field_int(f"{architecture}.block_count") + n_nextn = field_int(f"{architecture}.nextn_predict_layers") + quant_names = [] + for tensor in reader.tensors: + logical = _logical_shape(tensor) + if eligible_weight(tensor.name, logical, quant_args, gguf): + quant_names.append(tensor.name) + 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)) + + work = output.with_name(output.name + ".nvfp4-work") + if args.fresh and work.exists(): + shutil.rmtree(work) + work.mkdir(parents=True, exist_ok=True) + require_quant_settings(work, quant_args) + _converter_cls, unpack_uint4 = load_learned_converter(quant_args.convert_to_quant) + del _converter_cls + quantizer = TensorQuantizer(quant_args, gguf, unpack_uint4) + writer = gguf.GGUFWriter(output, architecture, use_temp_file=not args.no_temp_file) + _copy_metadata(reader, writer, gguf, getattr(args, "context_length", None)) + + done = 0 + for tensor in reader.tensors: + scale_for = "" + if tensor.name.endswith(".scale") or tensor.name.endswith(".input_scale"): + stem = tensor.name.rsplit(".", 1)[0] + scale_for = stem + ".weight" + if scale_for in quant_name_set and not ( + tensor.name.endswith(".scale") and preserve_scale_tensor(architecture, scale_for) + ): + continue + if tensor.name not in quant_name_set: + writer.add_tensor(tensor.name, tensor.data, raw_dtype=tensor.tensor_type) + continue + done += 1 + logical = _logical_shape(tensor) + bias_name = tensor.name.replace(".weight", ".bias") + bake = must_bake_nvfp4_scale( + architecture, + tensor.name, + (".ffn_" in tensor.name or tensor.name.startswith("ffn_")) and 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, _source_description(tensor), + " (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) + _save_checkpoint(work, tensor.name, logical, raw, scale) + del weight + else: + raw, scale = cached + writer.add_tensor(tensor.name, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4) + if scale is not None: + writer.add_tensor(tensor.name[: -len(".weight")] + ".scale", scale) + del raw + + log.info("Writing %s", output) + from nvfp4gguf.bundle import warn_missing_gguf_metadata + 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) + log.info("Wrote %s", output) + + +def _split_name(name: str) -> tuple[str, str, str]: + if ":" not in name: + name = name + ":latest" + model, tag = name.rsplit(":", 1) + if "/" in model: + namespace, model = model.split("/", 1) + else: + namespace = "library" + return namespace, model, tag + + +def _read_blob(root: Path, digest: str) -> str: + path = blob_path(root, digest) + if not path.is_file(): + raise FileNotFoundError(f"Missing blob {path}") + return path.read_text(encoding="utf-8") + + +def _quote_modelfile_value(text: str) -> str: + """Quote a Modelfile argument the way Ollama's parser reads it back.""" + if text == "" or "\n" in text or text[:1].isspace() or text[-1:].isspace() 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 _format_parameter(value) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return str(value) + if isinstance(value, (list, dict)): + return _quote_modelfile_value(json.dumps(value, separators=(",", ":"))) + return _quote_modelfile_value(str(value)) + + +def _config_lines(config: dict) -> list[str]: + """Renderer, parser, requires, and explicit capabilities. + + Ollama uses the renderer and parser for thinking and tool calls. Gemma 4, + Qwen3.5, and LFM2 thinking models stop and split thinking only when these + are present. A copied template alone does not turn thinking on. + """ + lines: list[str] = [] + if config.get("requires"): + lines.append(f"REQUIRES {_quote_modelfile_value(str(config['requires']))}") + if config.get("renderer"): + lines.append(f"RENDERER {_quote_modelfile_value(str(config['renderer']))}") + if config.get("parser"): + lines.append(f"PARSER {_quote_modelfile_value(str(config['parser']))}") + capabilities = config.get("capabilities") or [] + if isinstance(capabilities, str): + capabilities = [capabilities] + for capability in capabilities: + lines.append(f"CAPABILITY {_quote_modelfile_value(str(capability))}") + return lines + + +def _quote_block(command: str, text: str) -> str: + body = text.replace("\r\n", "\n") + return f"{command} {_quote_modelfile_value(body)}\n" + + +def _message_lines(text: str) -> list[str]: + try: + messages = json.loads(text) + except json.JSONDecodeError: + return ["# Messages blob was not JSON and was left out."] + if not isinstance(messages, list): + return ["# Messages blob was not a list and was left out."] + lines: list[str] = [] + for item in messages: + if not isinstance(item, dict): + continue + role = str(item.get("role") or "user") + content = item.get("content") if isinstance(item.get("content"), str) else "" + if not content and item.get("tool_calls"): + content = json.dumps(item["tool_calls"], separators=(",", ":")) + lines.append(f"MESSAGE {role} {_quote_modelfile_value(content)}") + return lines + + +# This fine-tune drafts inside a block and often never emits . +# temperature 0.3, repeat_penalty 1.2, and repeat_last_n 128 got that story prompt +# to close the thought and stop. The chat template stays the source template. +_GEMMA_THINK_SAMPLER = ( + ("temperature", 0.3), + ("repeat_penalty", 1.2), + ("repeat_last_n", 128), +) + + +def _gemma_thinking_sampler(source: OllamaModel | None) -> bool: + if source is None or source.family.lower() not in {"gemma3", "gemma3n"}: + return False + if source.config.get("renderer") or source.config.get("parser"): + return False + text = " ".join(( + source.name, + str(source.metadata.get("general.name") or ""), + str(source.metadata.get("general.finetune") or ""), + )).lower() + return "thinking" in text + + +def _extra_sampler(source: OllamaModel | None, existing: dict) -> list[tuple[str, object]]: + if not _gemma_thinking_sampler(source): + return [] + missing = [(key, value) for key, value in _GEMMA_THINK_SAMPLER if key not in existing] + if missing: + log.warning( + "This Gemma thinking fine-tune drafts without writing . " + "Adding temperature 0.3, repeat_penalty 1.2, and repeat_last_n 128 so the reply reaches the stop token." + ) + return missing + + +def render_modelfile(root: Path, source: OllamaModel | None, from_path: str) -> str: + """Rebuild the source Modelfile without changing its prompt or sampling. + + Renderer and parser are what turn thinking on for Gemma 4, Qwen, and LFM2. + The chat template and stop list are copied as stored, because rewriting them + is what made Gemma keep generating. + """ + lines = [f"FROM {from_path}", ""] + if source is None: + return "\n".join(lines) + if any(layer.get("mediaType") == PROJECTOR_MEDIA for layer in source.layers): + lines.append("FROM ./mmproj.gguf") + lines.append("") + config_lines = _config_lines(source.config) + if config_lines: + lines.extend(config_lines) + lines.append("") + notes: list[str] = [] + for layer in source.layers: + media = str(layer.get("mediaType") or "") + kind = LAYER_KIND.get(media) + if media in (PROJECTOR_MEDIA, ADAPTER_MEDIA) or kind in (None, "PROJECTOR", "ADAPTER"): + if media == ADAPTER_MEDIA: + notes.append("# Ollama 0.40 no longer loads LoRA adapters, so the adapter layer was not written as ADAPTER.") + elif kind is None and media and media != MODEL_MEDIA: + notes.append(f"# Layer {media} is kept by --install and has no Modelfile command.") + continue + try: + text = _read_blob(root, layer["digest"]) + except (OSError, KeyError) as exc: + notes.append(f"# Could not read {kind.lower()} blob: {exc}") + continue + if kind == "PARAMETER": + try: + params = json.loads(text) + except json.JSONDecodeError: + notes.append("# Parameters blob was not JSON and was left out.") + continue + if isinstance(params, dict): + if "num_ctx" in params: + log.info("Leaving num_ctx out of the Modelfile so OLLAMA_CONTEXT_LENGTH applies.") + for key, value in params.items(): + if key == "num_ctx": + continue + values = value if isinstance(value, list) else [value] + for item in values: + lines.append(f"PARAMETER {key} {_format_parameter(item)}") + for key, value in _extra_sampler(source, params): + lines.append(f"PARAMETER {key} {_format_parameter(value)}") + lines.append("") + continue + if kind == "MESSAGE": + lines.extend(_message_lines(text)) + lines.append("") + continue + lines.append(_quote_block(kind, text)) + if notes: + lines.extend(notes) + lines.append("") + return "\n".join(lines).rstrip() + "\n" + + +def export_bundle(root: Path, export_dir: Path, gguf_file: Path, source: OllamaModel | None, overwrite: bool, copy_weights: bool = True) -> Path: + export_dir.mkdir(parents=True, exist_ok=True) + modelfile_path = export_dir / "Modelfile" + if modelfile_path.exists() and not overwrite: + raise SystemExit(f"{modelfile_path} already exists. Pass --overwrite to replace it.") + try: + relative = Path(os.path.relpath(gguf_file, export_dir)) + from_path = relative.as_posix() if relative.is_absolute() else relative.as_posix() + except ValueError: + from_path = gguf_file.as_posix() + if not from_path.startswith(("/", "./", "../")): + from_path = "./" + from_path + modelfile_path.write_text(render_modelfile(root, source, from_path), encoding="utf-8") + if copy_weights and source is not None: + for layer in source.layers: + if layer.get("mediaType") != PROJECTOR_MEDIA: + continue + dest = export_dir / "mmproj.gguf" + src = blob_path(root, layer["digest"]) + if not src.is_file(): + log.warning("Projector blob is missing: %s", src) + continue + if dest.exists() and not overwrite: + raise SystemExit(f"{dest} already exists. Pass --overwrite to replace it.") + if dest.exists(): + dest.unlink() + try: + os.link(src, dest) + except OSError: + shutil.copy2(src, dest) + log.info("Copied vision projector to %s", dest) + log.info("Wrote %s", modelfile_path) + return modelfile_path + + +def install_model(root: Path, gguf_file: Path, source: OllamaModel, new_name: str) -> Path: + digest = sha256_file(gguf_file) + blob = root / "blobs" / f"sha256-{digest}" + blob.parent.mkdir(parents=True, exist_ok=True) + if not blob.exists(): + try: + os.link(gguf_file, blob) + except OSError: + shutil.copy2(gguf_file, blob) + config = {key: value for key, value in source.config.items() if key != "rootfs"} + config["model_format"] = source.config.get("model_format") or "gguf" + config["model_family"] = source.family + config["model_families"] = source.config.get("model_families") or ([source.family] if source.family else []) + config["model_type"] = source.params + config["file_type"] = "NVFP4" + config.setdefault("architecture", "amd64") + config.setdefault("os", "linux") + config["rootfs"] = {"type": "layers", "diff_ids": []} + layers = [] + diff_ids = [] + for layer in source.layers: + copied = dict(layer) + copied.pop("from", None) + if layer.get("mediaType") == MODEL_MEDIA: + copied = { + "mediaType": MODEL_MEDIA, + "digest": f"sha256:{digest}", + "size": gguf_file.stat().st_size, + } + elif layer.get("mediaType") == "application/vnd.ollama.image.params": + try: + current = json.loads(_read_blob(root, layer["digest"])) + except (OSError, KeyError, json.JSONDecodeError): + current = None + extra = _extra_sampler(source, current if isinstance(current, dict) else {}) + if isinstance(current, dict) and (extra or "num_ctx" in current): + current.pop("num_ctx", None) + for key, value in extra: + current[key] = value + raw = json.dumps(current, separators=(",", ":")).encode("utf-8") + params_digest = hashlib.sha256(raw).hexdigest() + params_blob = root / "blobs" / f"sha256-{params_digest}" + params_blob.parent.mkdir(parents=True, exist_ok=True) + if not params_blob.exists(): + params_blob.write_bytes(raw) + copied = { + "mediaType": layer["mediaType"], + "digest": f"sha256:{params_digest}", + "size": len(raw), + } + layers.append(copied) + diff_ids.append(copied["digest"]) + config["rootfs"]["diff_ids"] = diff_ids + config_bytes = json.dumps(config, separators=(",", ":")).encode("utf-8") + config_digest = hashlib.sha256(config_bytes).hexdigest() + config_blob = root / "blobs" / f"sha256-{config_digest}" + config_blob.write_bytes(config_bytes) + namespace, model, tag = _split_name(new_name) + store = "manifests-v2" if source.store == "manifests-v2" else "manifests" + host = "ollama.com" if store == "manifests-v2" else "registry.ollama.ai" + manifest_path = root / store / host / namespace / model / tag + manifest_path.parent.mkdir(parents=True, exist_ok=True) + manifest = { + "schemaVersion": 2, + "mediaType": "application/vnd.docker.distribution.manifest.v2+json", + "config": {"mediaType": CONFIG_MEDIA, "digest": f"sha256:{config_digest}", "size": len(config_bytes)}, + "layers": layers, + } + if source.runner: + manifest["runner"] = source.runner + elif store == "manifests-v2": + manifest["runner"] = "llamacpp" + if source.manifest_format: + manifest["format"] = source.manifest_format + elif store == "manifests-v2": + manifest["format"] = "gguf" + manifest_path.write_text(json.dumps(manifest, separators=(",", ":")), encoding="utf-8") + return manifest_path + + +def _resolve_source(args) -> tuple[Path, OllamaModel | None, str]: + query = args.model + if Path(query).is_file(): + source = Path(query).resolve() + return source, None, source.stem + root = Path(args.models).resolve() + source_model = resolve_model(load_models(root), query) + if not source_model.gguf_path.is_file(): + raise SystemExit(f"Weight blob is missing: {source_model.gguf_path}") + stem = source_model.name.replace("/", "-").replace(":", "-") + return source_model.gguf_path, source_model, stem + + +def _output_dir(args, stem: str) -> Path: + if args.output: + return Path(args.output).resolve() + return Path.cwd() / stem + + +def _wants_overwrite(args) -> bool: + return bool(getattr(args, "overwrite", False) or getattr(args, "force", False)) + + +def _place_gguf(src: Path, dest: Path, overwrite: bool) -> None: + dest.parent.mkdir(parents=True, exist_ok=True) + if dest.exists() and not overwrite: + raise SystemExit(f"{dest} already exists. Pass --overwrite to replace it.") + if dest.exists(): + dest.unlink() + try: + os.link(src, dest) + except OSError: + shutil.copy2(src, dest) + + +def cmd_export(args) -> None: + """Write the original GGUF and a Modelfile into a directory. No quantization.""" + logging.basicConfig(level=logging.INFO, format="%(message)s") + source, source_model, stem = _resolve_source(args) + out_dir = _output_dir(args, stem) + gguf_file = out_dir / f"{out_dir.name}.gguf" + overwrite = _wants_overwrite(args) + if not args.modelfile_only: + _place_gguf(source, gguf_file, overwrite) + else: + out_dir.mkdir(parents=True, exist_ok=True) + log.info("Skipping GGUF copy; writing Modelfile only") + export_bundle(Path(args.models).resolve(), out_dir, gguf_file, source_model, overwrite, copy_weights=not args.modelfile_only) + log.info("Exported %s", out_dir) + + +def cmd_quant(args) -> None: + logging.basicConfig(level=logging.INFO, format="%(message)s") + source, source_model, stem = _resolve_source(args) + out_dir = _output_dir(args, f"{stem}-NVFP4") + output = out_dir / f"{out_dir.name}.gguf" + modelfile_path = out_dir / "Modelfile" + overwrite = _wants_overwrite(args) + if modelfile_path.exists() and not overwrite: + raise SystemExit(f"{modelfile_path} already exists. Pass --overwrite to replace it.") + quant_args = argparse.Namespace( + device=args.device, + optimizer=args.optimizer, + num_iter=args.num_iter, + lr=args.lr, + lr_min=args.lr_min, + top_p=args.top_p, + min_k=args.min_k, + max_k=args.max_k, + full_matrix=args.full_matrix, + simple=args.simple, + lr_schedule=args.lr_schedule, + early_stop_loss=args.early_stop_loss, + early_stop_stall=args.early_stop_stall, + early_stop_floor_stall=args.early_stop_floor_stall, + target_loss=args.target_loss, + scale_refinement_rounds=args.scale_refinement_rounds, + scale_optimization=args.scale_optimization, + convert_to_quant=args.convert_to_quant, + quantize_output=args.quantize_output, + exclude=args.exclude, + exclude_re=None, + calib_samples=args.calib_samples, + seed=args.seed, + no_bias_correction=args.no_bias_correction, + ) + if args.modelfile_only: + out_dir.mkdir(parents=True, exist_ok=True) + log.info("Skipping quantization and GGUF creation; writing Modelfile only") + else: + quantize_gguf(args, source, output, quant_args) + export_bundle(Path(args.models).resolve(), out_dir, output, source_model, overwrite, copy_weights=not args.modelfile_only) + log.info("Quantized bundle %s" if not args.modelfile_only else "Wrote Modelfile %s", out_dir) + if args.install: + if source_model is None: + raise SystemExit("--install needs an Ollama model name, not a bare GGUF path") + if not output.is_file(): + raise SystemExit(f"{output} is not there, so --install has no GGUF. Run quant without --modelfile-only first.") + root = Path(args.models).resolve() + new_name = args.name or _nvfp4_name(source_model.name) + manifest = install_model(root, output, source_model, new_name) + log.info("Registered %s", new_name) + log.info("Manifest %s", manifest) + + +_INFO_USAGE = "usage: ollama_nvfp4.py [--models DIR] --info MODEL [arguments for gguf_info.py...]" + + +def _without_models_dir(argv: list[str]) -> tuple[list[str], Path]: + """Pull --models out of an --info command so it is not forwarded to gguf_info.py.""" + kept: list[str] = [] + models: Path | None = None + index = 0 + while index < len(argv): + arg = argv[index] + if arg == "--models" or arg.startswith("--models="): + if arg == "--models": + if index + 1 >= len(argv): + raise SystemExit(f"{_INFO_USAGE}\n--models needs a directory") + value = argv[index + 1] + index += 2 + else: + value = arg.split("=", 1)[1] + index += 1 + models = Path(value) + continue + kept.append(arg) + index += 1 + return kept, models if models is not None else default_models_dir() + + +def cmd_info(argv: list[str]) -> None: + """Resolve MODEL to its sha256 weight blob and run gguf_info.py on that file. + + Arguments after the model name are forwarded unchanged. --models is kept + here and may appear before or after --info. + """ + kept, models_dir = _without_models_dir(argv) + try: + info_at = kept.index("--info") + except ValueError: + raise SystemExit(_INFO_USAGE) from None + head, tail = kept[:info_at], kept[info_at + 1:] + if head or not tail or tail[0].startswith("-"): + extra = f"\nunrecognized: {' '.join(head)}" if head else "" + raise SystemExit(_INFO_USAGE + extra) + model_name, forwarded = tail[0], tail[1:] + root = Path(models_dir).resolve() + model = resolve_model(load_models(root), model_name) + if not model.gguf_path.is_file(): + raise SystemExit(f"Weight blob is missing: {model.gguf_path}") + script = Path(__file__).resolve().parent / "gguf_info.py" + if not script.is_file(): + raise SystemExit(f"gguf_info.py is not next to this script: {script}") + command = [sys.executable, str(script), str(model.gguf_path), *forwarded] + completed = subprocess.run(command) + raise SystemExit(completed.returncode) + + +def _nvfp4_name(name: str) -> str: + if ":" not in name: + return name + ":nvfp4" + model, tag = name.rsplit(":", 1) + if tag == "latest": + return model + ":nvfp4" + return f"{model}:{tag}-nvfp4" + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + from nvfp4gguf import ( + DEFAULT_EARLY_STOP_FLOOR_STALL, + DEFAULT_TARGET_LOSS, + DEFAULT_EARLY_STOP_STALL, + DEFAULT_SCALE_OPTIMIZATION, + DEFAULT_SCALE_REFINEMENT_ROUNDS, + ) + parser = argparse.ArgumentParser( + description=( + "Read an Ollama sha256 model store and quantize a model to NVFP4 GGUF. " + "--info MODEL runs gguf_info.py on that model's sha256 weight blob. " + "Arguments after MODEL are passed through, for example --metadata or --tensors." + ), + ) + parser.add_argument("--models", type=Path, default=default_models_dir(), help="Ollama models directory (blobs/ and manifests/)") + sub = parser.add_subparsers(dest="command", required=True) + + listed = sub.add_parser("list", help="Show every manifest as a readable model") + listed.add_argument( + "--full-scan", + action="store_true", + help="Read each GGUF header for capabilities and missing encoders. Slower on large models.", + ) + listed.add_argument( + "--ctx-scan", + action="store_true", + help="Read parameter blobs and list models whose Modelfile sets num_ctx.", + ) + + show = sub.add_parser("show", help="Show one model, its blobs, and its GGUF metadata") + show.add_argument("model", help="Name from list, or a unique substring") + + export = sub.add_parser("export", help="Write the original GGUF and a Modelfile into a directory") + export.add_argument("model", help="Ollama model name, or a path to a GGUF file") + export.add_argument("-o", "--output", type=Path, help="Output directory. Default: ./ in the current directory") + export.add_argument("--modelfile-only", action="store_true", help="Write the Modelfile and skip copying the GGUF and projector") + export.add_argument("--overwrite", action="store_true", help="Replace any output file that already exists") + export.add_argument("--force", action="store_true", help="Same as --overwrite") + + quant = sub.add_parser("quant", help="Quantize into a directory containing the NVFP4 GGUF and a Modelfile") + quant.add_argument("model", help="Ollama model name, or a path to a GGUF file") + quant.add_argument("-o", "--output", type=Path, help="Output directory. Default: ./-NVFP4 in the current directory") + quant.add_argument("--name", help="Ollama name to register with --install. Default: :nvfp4 or -nvfp4") + quant.add_argument("--install", action="store_true", help="Register the NVFP4 GGUF in the Ollama store") + quant.add_argument("--modelfile-only", action="store_true", help="Write the Modelfile and skip quantization and GGUF creation") + quant.add_argument("--overwrite", action="store_true", help="Replace any output file that already exists") + quant.add_argument("--force", action="store_true", help="Same as --overwrite") + quant.add_argument("--float-only", action="store_true", help="Refuse to requantize a model that is already quantized") + quant.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.") + quant.add_argument("--llama-cpp", type=Path, default=Path.home() / "llama.cpp", + help="llama.cpp checkout used for the GGUF reader and writer (default: ~/llama.cpp)") + quant.add_argument("--convert-to-quant", type=Path, default=None) + quant.add_argument("--device", default="auto") + quant.add_argument("--simple", action="store_true", help="Round onto the E2M1 grid without learned rounding") + quant.add_argument("--optimizer", default="auto", choices=("auto", "prodigy", "adamw", "radam", "original")) + quant.add_argument("--num-iter", type=int, default=4000) + quant.add_argument("--lr", type=float, default=1.0, help="Initial learning rate") + quant.add_argument("--lr-min", type=float, default=1e-8, help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).") + quant.add_argument("--lr-schedule", default="plateau", choices=("plateau", "exponential", "adaptive")) + quant.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)", + ) + quant.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.", + ) + quant.add_argument("--full-matrix", action="store_true") + quant.add_argument("--top-p", type=float, default=0.2) + quant.add_argument("--min-k", type=int, default=128) + quant.add_argument("--max-k", type=int, default=1280) + quant.add_argument("--early-stop-loss", type=float, default=5e-9) + quant.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)", + ) + quant.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.", + ) + quant.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)", + ) + quant.add_argument("--calib-samples", type=int, default=3072) + quant.add_argument("--seed", type=int, default=42) + quant.add_argument("--no-bias-correction", action="store_true") + quant.add_argument("--quantize-output", action="store_true") + quant.add_argument("--exclude", default=None, help="Regex of GGUF tensor names to leave unchanged") + quant.add_argument("--fresh", action="store_true") + quant.add_argument("--keep-checkpoints", action="store_true") + quant.add_argument("--no-temp-file", action="store_true") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> None: + if argv is None: + argv = sys.argv[1:] + if "--info" in argv: + cmd_info(argv) + return + args = parse_args(argv) + root = Path(args.models).resolve() + if args.command == "list": + cmd_list(root, full_scan=args.full_scan, ctx_scan=args.ctx_scan) + return + if args.command == "show": + cmd_show(root, args.model) + return + if args.command == "export": + cmd_export(args) + return + if args.command == "quant": + cmd_quant(args) + return + raise SystemExit(f"Unknown command {args.command}") + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..086b8c6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,19 @@ +# Dependencies for hf_to_nvfp4_gguf.py, ollama_nvfp4.py, +# safetensors_to_nvfp4_gguf.py, and nvfp4gguf. +# +# llama.cpp (GGUF writer and Hugging Face conversion) and convert_to_quant +# (LearnedNVFP4Converter) are local checkouts, not pip packages. Point the +# scripts at them with --llama-cpp and --convert-to-quant. +# +# Install a CUDA build of torch that matches your GPU. A CPU wheel works +# with --device cpu. + +numpy>=2.2.6 +torch>=2.6.0 +safetensors>=0.4.2 +tqdm>=4.27 +PyYAML>=5.1 +transformers>=4.46 +sentencepiece>=0.1.98,<0.3.0 +protobuf>=4.21.0,<5.0.0 +prodigy-plus-schedule-free diff --git a/safetensors_to_nvfp4_gguf.py b/safetensors_to_nvfp4_gguf.py new file mode 100644 index 0000000..7af827e --- /dev/null +++ b/safetensors_to_nvfp4_gguf.py @@ -0,0 +1,532 @@ +#!/usr/bin/env python3 +"""Quantize a safetensors file to NVFP4 GGUF without a Hugging Face model directory. + +The learned rounding is the same convert_to_quant pass the other scripts use, +so a file you are already quantizing with convert_to_quant can be packed into +GGUF and compared tensor by tensor. Tensor names are kept. There is no +tokenizer and no architecture config; the GGUF is a weight container. + +By default, tensors that convert_to_quant's base exclusion list would skip +(norms, biases, embeddings, lm_head, and the other names in that list) stay +in their original dtype. Pass --all-weights to quantize every 2D matrix whose +last dimension is a multiple of 64. A sibling bias is corrected the same way +as hf_to_nvfp4_gguf.py. Repeat --mmproj to copy projector GGUFs beside the +output and add a FROM line for each one. + +--no-quant skips that rounding and writes the safetensors into a GGUF as +float. --outtype chooses f32, f16, bf16, or auto, which keeps each tensor's +stored dtype. --outtype nvfp4 repacks a ComfyUI NVFP4 safetensors file +(packed weight, weight_scale, and weight_scale_2) without fitting it again. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import logging +import shutil +from pathlib import Path + +import numpy as np +import torch +from safetensors import safe_open + +from nvfp4gguf import ( + DEFAULT_LLAMA_CPP, + TensorQuantizer, + add_quant_arguments, + bind_exclude, + ctq_skips_key, + ensure_cuda_interpreter, + import_gguf, + load_learned_converter, + looks_like_spark, + quantize_and_pack, + report_reconstruction, + require_quant_settings, +) +from nvfp4gguf.bundle import ( + outtype as _outtype, + refuse_existing_bundle, + require_context_length, + warn_missing_gguf_metadata, + write_simple_modelfile, +) + +log = logging.getLogger("safetensors_to_nvfp4_gguf") + + +def _checkpoint_file(directory: Path, name: str) -> Path: + digest = hashlib.sha1(name.encode("utf-8")).hexdigest()[:16] + safe = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in name)[:80] + return directory / f"{safe}.{digest}.npz" + + +def _load_checkpoint(directory: Path, name: str, shape: tuple[int, ...], need_bias: bool): + path = _checkpoint_file(directory, name) + if not path.is_file(): + return None + with np.load(path) as blob: + if tuple(int(v) for v in blob["shape"]) != tuple(shape): + return None + if need_bias and "bias" not in blob.files: + return None + raw = np.array(blob["raw"], copy=True) + scale = np.array(blob["scale"], copy=True) + bias = np.array(blob["bias"], copy=True) if "bias" in blob.files else np.zeros((0,), np.float32) + 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_file(directory, name) + temporary = path.with_name(path.name + ".tmp") + with temporary.open("wb") as handle: + np.savez_compressed( + handle, + 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 _should_quantize( + key: str, + shape: tuple[int, ...], + all_weights: bool, + convert_to_quant_dir: Path | None, + exclude_re, +) -> bool: + if len(shape) != 2 or not key.endswith(".weight"): + return False + if shape[-1] % 64 != 0 or shape[-1] < 64: + return False + if exclude_re is not None and exclude_re.search(key): + return False + if not all_weights and ctq_skips_key(key, convert_to_quant_dir): + return False + return True + + +def _apply_bias(name: str, tensor: torch.Tensor, correction: np.ndarray | None) -> torch.Tensor: + if correction is None: + return tensor + values = tensor.detach().cpu().float() + flat = values.reshape(-1) + if tuple(flat.shape) != tuple(correction.shape): + log.warning(" bias %s shape %s does not match correction %s", name, tuple(flat.shape), tuple(correction.shape)) + return tensor + flat = flat - torch.from_numpy(np.ascontiguousarray(correction, dtype=np.float32)) + values = flat.reshape(values.shape).to(dtype=tensor.dtype) + log.info(" bias %s adjusted (mean %+.3e)", name, float(np.mean(correction))) + return values + + +def _float_file_type(gguf, outtype: str): + return { + "f32": gguf.LlamaFileType.ALL_F32, + "f16": gguf.LlamaFileType.MOSTLY_F16, + "bf16": gguf.LlamaFileType.MOSTLY_BF16, + }[outtype] + + +def _auto_file_type(gguf, dtypes: list[str]): + names = {dtype.upper() for dtype in dtypes} + if "BF16" in names: + return gguf.LlamaFileType.MOSTLY_BF16 + if "F16" in names: + return gguf.LlamaFileType.MOSTLY_F16 + if "F32" in names: + return gguf.LlamaFileType.ALL_F32 + return gguf.LlamaFileType.GUESSED + + +def _cast_float(tensor: torch.Tensor, outtype: str) -> torch.Tensor: + if outtype == "auto" or not tensor.is_floating_point(): + return tensor + dtype = {"f32": torch.float32, "f16": torch.float16, "bf16": torch.bfloat16}[outtype] + if tensor.dtype == dtype: + return tensor + return tensor.to(dtype) + + +def _default_output(source: Path, skip_quant: bool, outtype: str) -> Path: + if not skip_quant: + return source.with_name(source.stem + "-NVFP4.gguf") + if outtype == "auto": + return source.with_name(source.stem + ".gguf") + label = "NVFP4" if outtype == "nvfp4" else outtype.upper() + return source.with_name(f"{source.stem}-{label}.gguf") + + +def _add_original(writer, gguf, name: str, tensor: torch.Tensor) -> None: + tensor = tensor.detach().cpu().contiguous() + if tensor.dtype == torch.bfloat16: + raw = tensor.view(torch.uint16).numpy().view(np.uint8) + writer.add_tensor(name, np.ascontiguousarray(raw), raw_dtype=gguf.GGMLQuantizationType.BF16) + return + if tensor.dtype == torch.float16: + writer.add_tensor(name, tensor.numpy()) + return + if tensor.dtype == torch.float32: + writer.add_tensor(name, tensor.numpy()) + return + writer.add_tensor(name, tensor.float().numpy()) + + +def _prepare_output(args) -> tuple[Path, Path, bool]: + source = Path(args.safetensors).resolve() + if not source.is_file(): + raise SystemExit(f"Not a file: {source}") + skip_quant = bool(args.no_quant or args.outtype == "nvfp4") + output = Path(args.output).resolve() if args.output else _default_output(source, skip_quant, args.outtype) + return source, output, skip_quant + + +def _comfy_formats(handle) -> dict[str, str]: + """Map a weight prefix to the format named in its .comfy_quant sidecar.""" + found: dict[str, str] = {} + suffix = ".comfy_quant" + for key in handle.keys(): + if not key.endswith(suffix): + continue + raw = handle.get_tensor(key).detach().cpu().numpy().tobytes() + try: + payload = json.loads(raw.decode("utf-8")) + except (UnicodeError, json.JSONDecodeError): + continue + fmt = payload.get("format") if isinstance(payload, dict) else None + if isinstance(fmt, str): + found[key[: -len(suffix)]] = fmt + return found + + +def _pack_comfy_nvfp4(weight: torch.Tensor, scale: torch.Tensor) -> np.ndarray: + """Repack ComfyUI / ModelOpt NVFP4 into ggml super-blocks. + + The code bytes are already nibble-packed the way ggml stores them. Each + group of four 16-value blocks becomes 4 UE4M3 scale bytes followed by + those 32 code bytes. + """ + if scale.ndim != 2 or weight.ndim != 2: + raise ValueError(f"NVFP4 weight {tuple(weight.shape)} and scale {tuple(scale.shape)} must be matrices") + rows, n_blocks = scale.shape + if weight.shape != (rows, n_blocks * 8): + raise ValueError(f"NVFP4 weight {tuple(weight.shape)} does not match scale {tuple(scale.shape)}") + if n_blocks % 4 != 0: + raise ValueError(f"NVFP4 scale count {n_blocks} is not a multiple of 4") + n_super = n_blocks // 4 + codes = np.ascontiguousarray(weight.detach().cpu().contiguous().view(torch.uint8).numpy()).reshape(rows, n_super, 32) + scales = np.ascontiguousarray(scale.detach().cpu().contiguous().view(torch.uint8).numpy()).reshape(rows, n_super, 4) + scales = scales & np.uint8(0x7F) + return np.ascontiguousarray(np.concatenate([scales, codes], axis=-1).reshape(rows, n_super * 36)) + + +def _scale_is_trivial(scale: torch.Tensor) -> bool: + return scale.numel() <= 1 and abs(float(scale.detach().float().sum()) - 1.0) < 1e-6 + + +def _float8_to_bf16(tensor: torch.Tensor, scale: torch.Tensor | None) -> torch.Tensor: + values = tensor.detach().float() + if scale is not None: + factor = scale.detach().float() + if factor.numel() == 1: + values = values * factor.reshape(()) + elif tuple(factor.shape) == tuple(values.shape): + values = values * factor + else: + raise ValueError(f"float8 scale {tuple(factor.shape)} does not match weight {tuple(values.shape)}") + return values.to(torch.bfloat16) + + +def convert_packed_nvfp4(args) -> None: + """Copy ComfyUI NVFP4 blocks into a GGUF. Float tensors are stored as stored.""" + logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s") + source, output, _skip = _prepare_output(args) + overwrite = bool(args.force or args.overwrite) + if output.exists() and not overwrite: + raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.") + output.parent.mkdir(parents=True, exist_ok=True) + if args.mmproj: + modelfile = output.parent / "Modelfile" + if modelfile.exists() and not overwrite: + raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.") + refuse_existing_bundle(output, overwrite, args.mmproj, skip_gguf=True) + + gguf = import_gguf(args.llama_cpp) + with safe_open(source, framework="pt", device="cpu") as handle: + keys = list(handle.keys()) + formats = _comfy_formats(handle) + packed = [name for name, fmt in formats.items() if fmt == "nvfp4"] + if not packed: + raise SystemExit( + f"{source} has no ComfyUI NVFP4 weights to repack. " + "Use --no-quant with --outtype auto, f16, bf16, or f32 to write float weights without quantizing them." + ) + present = set(keys) + for name in packed: + if name + ".weight" not in present or name + ".weight_scale" not in present: + raise SystemExit(f"{name} is marked nvfp4 but is missing weight or weight_scale") + log.info("Repacking %d NVFP4 tensors from %s. Learned rounding is skipped.", len(packed), source.name) + writer = gguf.GGUFWriter(output, "safetensors", use_temp_file=not args.no_temp_file) + writer.add_file_type(int(gguf.LlamaFileType.MOSTLY_NVFP4)) + writer.add_string("general.name", source.stem) + writer.add_string("general.description", "NVFP4 weights repacked from safetensors without refitting") + writer.add_uint32("general.quantization_version", 2) + context_length = require_context_length(args.context_length) + if context_length is not None: + writer.add_context_length(context_length) + log.info("context length %d", context_length) + + skip = {name + ".comfy_quant" for name in formats} + for name, fmt in formats.items(): + if fmt in ("nvfp4", "float8_e4m3fn"): + skip.add(name + ".weight_scale") + skip.add(name + ".weight_scale_2") + promoted = 0 + with safe_open(source, framework="pt", device="cpu") as handle: + for key in keys: + if key in skip: + continue + if key.endswith(".weight") and formats.get(key[: -len(".weight")]) == "nvfp4": + base = key[: -len(".weight")] + raw = _pack_comfy_nvfp4(handle.get_tensor(key), handle.get_tensor(base + ".weight_scale")) + writer.add_tensor(key, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4) + scale2_name = base + ".weight_scale_2" + if scale2_name in present: + scale2 = handle.get_tensor(scale2_name) + if not _scale_is_trivial(scale2): + writer.add_tensor(base + ".scale", scale2.detach().float().numpy().reshape(-1).astype(np.float32)) + continue + if key.endswith(".weight") and formats.get(key[: -len(".weight")]) == "float8_e4m3fn": + base = key[: -len(".weight")] + scale_name = base + ".weight_scale" + scale = handle.get_tensor(scale_name) if scale_name in present else None + _add_original(writer, gguf, key, _float8_to_bf16(handle.get_tensor(key), scale)) + promoted += 1 + log.info("Stored %s as BF16 (source is float8_e4m3fn)", key) + continue + _add_original(writer, gguf, key, handle.get_tensor(key)) + if promoted: + log.info("%d float8 tensors were stored as BF16. GGUF has no float8 weight type.", promoted) + 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 args.mmproj: + write_simple_modelfile(output, overwrite, args.mmproj) + log.info("Wrote %s", output) + + +def convert_float(args) -> None: + """Write the safetensors into a GGUF without learned rounding.""" + logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s") + if args.outtype == "nvfp4": + convert_packed_nvfp4(args) + return + source, output, _skip = _prepare_output(args) + overwrite = bool(args.force or args.overwrite) + if output.exists() and not overwrite: + raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.") + output.parent.mkdir(parents=True, exist_ok=True) + if args.mmproj: + modelfile = output.parent / "Modelfile" + if modelfile.exists() and not overwrite: + raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.") + refuse_existing_bundle(output, overwrite, args.mmproj, skip_gguf=True) + + gguf = import_gguf(args.llama_cpp) + with safe_open(source, framework="pt", device="cpu") as handle: + keys = list(handle.keys()) + dtypes = [handle.get_slice(key).get_dtype() for key in keys] + file_type = _float_file_type(gguf, args.outtype) if args.outtype != "auto" else _auto_file_type(gguf, dtypes) + log.info("Writing float GGUF (%s). NVFP4 quantization is skipped.", args.outtype) + writer = gguf.GGUFWriter(output, "safetensors", use_temp_file=not args.no_temp_file) + writer.add_file_type(int(file_type)) + writer.add_string("general.name", source.stem) + writer.add_string("general.description", "Weights converted from safetensors without quantization") + context_length = require_context_length(args.context_length) + if context_length is not None: + writer.add_context_length(context_length) + log.info("context length %d", context_length) + with safe_open(source, framework="pt", device="cpu") as handle: + for key in keys: + _add_original(writer, gguf, key, _cast_float(handle.get_tensor(key), 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() + if args.mmproj: + write_simple_modelfile(output, overwrite, args.mmproj) + log.info("Wrote %s", output) + + +def convert(args) -> None: + logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s") + if args.no_quant or args.outtype == "nvfp4": + convert_float(args) + return + ensure_cuda_interpreter(args.device) + source, output, _skip = _prepare_output(args) + overwrite = bool(args.force or args.overwrite) + if output.exists() and not overwrite: + raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.") + output.parent.mkdir(parents=True, exist_ok=True) + if args.mmproj: + modelfile = output.parent / "Modelfile" + if modelfile.exists() and not overwrite: + raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.") + refuse_existing_bundle(output, overwrite, args.mmproj, skip_gguf=True) + + gguf = import_gguf(args.llama_cpp) + with safe_open(source, framework="pt", device="cpu") as handle: + keys = list(handle.keys()) + shapes = {key: tuple(handle.get_slice(key).get_shape()) for key in keys} + bind_exclude(args, "spark2_5" if looks_like_spark(keys) else None) + chosen = [ + key for key in keys + if _should_quantize(key, shapes[key], args.all_weights, args.convert_to_quant, args.exclude_re) + ] + log.info("%d of %d tensors queued for NVFP4", len(chosen), len(keys)) + if not chosen: + raise SystemExit("No 2D .weight tensor has a last dimension divisible by 64.") + + work = Path(args.checkpoint_dir) if args.checkpoint_dir else output.with_name(output.name + ".nvfp4-work") + if args.fresh and work.exists(): + shutil.rmtree(work) + work.mkdir(parents=True, exist_ok=True) + require_quant_settings(work, args) + quantizer = TensorQuantizer(args, gguf, load_learned_converter(args.convert_to_quant)[1]) + writer = gguf.GGUFWriter(output, "safetensors", use_temp_file=not args.no_temp_file) + writer.add_file_type(int(gguf.LlamaFileType.MOSTLY_NVFP4)) + writer.add_string("general.name", source.stem) + writer.add_string("general.description", "NVFP4 weights converted from safetensors for comparison with convert_to_quant") + writer.add_uint32("general.quantization_version", 2) + context_length = require_context_length(args.context_length) + if context_length is not None: + writer.add_context_length(context_length) + log.info("context length %d", context_length) + + done = 0 + chosen_set = set(chosen) + pending_bias: dict[str, torch.Tensor] = {} + corrections: dict[str, np.ndarray] = {} + + def _write_bias(key: str, tensor: torch.Tensor) -> None: + weight_name = key[: -len(".bias")] + ".weight" + correction = corrections.get(key) if weight_name in chosen_set and not args.no_bias_correction else None + _add_original(writer, gguf, key, _apply_bias(key, tensor, correction)) + + with safe_open(source, framework="pt", device="cpu") as handle: + for key in keys: + tensor = handle.get_tensor(key) + weight_name = key[: -len(".bias")] + ".weight" if key.endswith(".bias") else "" + if weight_name in chosen_set and not args.no_bias_correction and key not in corrections: + pending_bias[key] = tensor + continue + if key not in chosen_set: + if len(shapes[key]) == 2 and key.endswith(".weight") and shapes[key][-1] % 64 != 0: + log.info("Keeping %s %s (last dimension is not a multiple of 64)", key, shapes[key]) + if key.endswith(".bias") and weight_name in chosen_set: + _write_bias(key, tensor) + else: + _add_original(writer, gguf, key, tensor) + continue + done += 1 + bias_name = key[: -len(".weight")] + ".bias" + need_bias = not args.no_bias_correction and bias_name in shapes and len(shapes[key]) == 2 + cached = _load_checkpoint(work, key, shapes[key], need_bias) + log.info( + "NVFP4 [%d/%d] %s %s%s", + done, len(chosen), key, shapes[key], " (checkpoint)" if cached else "", + ) + if cached is None: + weight = tensor.float().numpy() + raw, scale = quantize_and_pack(quantizer, weight, False, key) + report_reconstruction(quantizer, weight, raw, scale) + bias = None + if need_bias: + per_tensor = None if scale is None else float(np.reshape(scale, -1)[0]) + runtime = quantizer.runtime_dequant(raw, per_tensor) + bias = quantizer.bias_correction(weight, runtime) + _save_checkpoint(work, key, shapes[key], raw, scale, bias) + del weight + else: + raw, scale, bias = cached + if bias is not None: + corrections[bias_name] = bias + pending = pending_bias.pop(bias_name, None) + if pending is not None: + _write_bias(bias_name, pending) + writer.add_tensor(key, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4) + if scale is not None: + writer.add_tensor(key[: -len(".weight")] + ".scale", scale) + del tensor, raw + for key, tensor in pending_bias.items(): + _write_bias(key, tensor) + + 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) + if args.mmproj: + write_simple_modelfile(output, overwrite, args.mmproj) + log.info("Wrote %s", output) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Quantize a safetensors file to an NVFP4 GGUF weight container.") + parser.add_argument("safetensors", nargs="?", help="Input .safetensors file") + parser.add_argument("-o", "--output", help="Output GGUF path. Default: -NVFP4.gguf, or .gguf with --no-quant and auto") + parser.add_argument("--no-quant", action="store_true", + help="Skip learned rounding. Writes a float GGUF, or repacks ComfyUI NVFP4 weights 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 repacks ComfyUI NVFP4 weights. auto keeps each tensor's stored dtype (default: auto)") + parser.add_argument("--all-weights", action="store_true", + help="Quantize every eligible 2D weight. By default, convert_to_quant's exclusion list is skipped too.") + parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP, + help="llama.cpp checkout used for the GGUF writer (default: ~/llama.cpp)") + parser.add_argument("--overwrite", action="store_true", help="Replace an existing GGUF, Modelfile, or copied projector") + parser.add_argument("--force", action="store_true", help="Same as --overwrite") + parser.add_argument("--mmproj", action="append", type=Path, default=None, + help="Vision projector GGUF. Copied next to the output 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. A safetensors file has no context length to copy.") + parser.add_argument("--checkpoint-dir", type=Path, default=None, help="Per-tensor resume directory") + parser.add_argument("--fresh", action="store_true") + parser.add_argument("--keep-checkpoints", action="store_true") + parser.add_argument("--no-temp-file", action="store_true") + 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") + add_quant_arguments(parser) + args = parser.parse_args(argv) + if args.self_test: + if args.device == "auto": + args.device = "cpu" + return args + if not args.safetensors: + parser.error("safetensors file is required") + 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 + convert(args) + + +if __name__ == "__main__": + main()