336 lines
14 KiB
Python
336 lines
14 KiB
Python
#!/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("<Q", raw_size)[0]
|
|
header = json.loads(handle.read(header_size))
|
|
if not isinstance(header, dict):
|
|
raise SystemExit(f"{path} has an unreadable safetensors header")
|
|
return header, header_size
|
|
|
|
|
|
def merge_safetensors(paths: list[Path], dest: Path) -> 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("<Q", len(header_bytes)))
|
|
out.write(header_bytes)
|
|
for path, abs_start, length in layout:
|
|
with path.open("rb") as src:
|
|
src.seek(abs_start)
|
|
remaining = length
|
|
while remaining:
|
|
chunk = src.read(min(remaining, 8 * 1024 * 1024))
|
|
if not chunk:
|
|
raise SystemExit(f"Short read from {path}")
|
|
out.write(chunk)
|
|
remaining -= len(chunk)
|
|
temporary.replace(dest)
|
|
|
|
|
|
def quantize_shard(convert_to_nvfp4, args, optimizer: str, source: Path, destination: Path) -> 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: ./<model-name>-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()
|