621 lines
25 KiB
Python
Executable File
621 lines
25 KiB
Python
Executable File
"""Quantize a Hugging Face model directory to NVFP4 Hugging Face format.
|
|
|
|
Reads a Hugging Face model directory (config.json + safetensors) and writes a
|
|
quantized Hugging Face model directory. By default (--format hf), writes standard
|
|
Hugging Face compressed-tensors format with .weight_packed, .weight_scale, and
|
|
.weight_global_scale natively supported by Hugging Face transformers, vLLM, SGLang,
|
|
and llama.cpp.
|
|
|
|
Optionally (--format modelopt / comfy), outputs raw weight tensors (.weight),
|
|
weight_scale, weight_scale_2, and .comfy_quant for ModelOpt and ComfyUI.
|
|
|
|
Weight matrices are quantized using convert_to_quant's learned rounding (or simple
|
|
round-to-grid with --simple). Biases are adjusted using calibration data unless
|
|
--no-bias-correction is specified. Embeddings, layer norms, and excluded layers
|
|
stay in their float format.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
|
|
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_LR_FLOOR,
|
|
DEFAULT_LR_PATIENCE,
|
|
DEFAULT_LR_THRESHOLD,
|
|
DEFAULT_SCALE_OPTIMIZATION,
|
|
DEFAULT_SCALE_REFINEMENT_ROUNDS,
|
|
QUANT_VERSION,
|
|
TensorQuantizer,
|
|
add_quant_arguments,
|
|
bind_exclude,
|
|
coordinate_descent_refine,
|
|
ctq_skips_key,
|
|
ensure_cuda_interpreter,
|
|
load_learned_converter,
|
|
quant_profile,
|
|
report_reconstruction,
|
|
resolve_optimizer,
|
|
sweep_per_tensor_scale,
|
|
ue4m3_bits,
|
|
)
|
|
|
|
INDEX_NAME = "model.safetensors.index.json"
|
|
NON_WEIGHT_PATTERNS = {
|
|
".safetensors",
|
|
".bin",
|
|
".pt",
|
|
".pth",
|
|
".h5",
|
|
".msgpack",
|
|
".ot",
|
|
".tmp",
|
|
}
|
|
|
|
log = logging.getLogger("hf_to_nvfp4_hf")
|
|
|
|
|
|
def _checkpoint_file(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 _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 "version" in blob.files and int(blob["version"]) != QUANT_VERSION:
|
|
return None
|
|
if tuple(int(v) for v in blob["shape"]) != tuple(shape):
|
|
return None
|
|
if need_bias and "bias" not in blob.files:
|
|
return None
|
|
codes = np.array(blob["codes"], copy=True)
|
|
block_scale = np.array(blob["block_scale"], copy=True)
|
|
per_tensor = float(blob["per_tensor"])
|
|
bias = np.array(blob["bias"], copy=True) if "bias" in blob.files else None
|
|
return codes, block_scale, per_tensor, bias
|
|
|
|
|
|
def _save_checkpoint(directory: Path, name: str, shape: tuple[int, ...], codes, block_scale, per_tensor, bias) -> None:
|
|
path = _checkpoint_file(directory, name)
|
|
temporary = path.with_name(path.name + ".tmp")
|
|
with temporary.open("wb") as handle:
|
|
save_dict = {
|
|
"version": np.int32(QUANT_VERSION),
|
|
"shape": np.asarray(shape, dtype=np.int64),
|
|
"codes": np.asarray(codes, dtype=np.uint8),
|
|
"block_scale": np.asarray(block_scale, dtype=np.float32),
|
|
"per_tensor": np.float32(per_tensor),
|
|
}
|
|
if bias is not None:
|
|
save_dict["bias"] = np.asarray(bias, dtype=np.float32)
|
|
np.savez_compressed(handle, **save_dict)
|
|
temporary.replace(path)
|
|
|
|
|
|
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():
|
|
try:
|
|
payload = json.loads(index_path.read_text(encoding="utf-8"))
|
|
weight_map = payload.get("weight_map")
|
|
if isinstance(weight_map, dict) and weight_map:
|
|
names = sorted(set(weight_map.values()))
|
|
missing = [name for name in names if not (model_dir / name).is_file()]
|
|
if not missing:
|
|
return names
|
|
except Exception:
|
|
pass
|
|
|
|
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 found in {model_dir}")
|
|
return shards
|
|
|
|
|
|
def _assign_output(args) -> None:
|
|
source = Path(args.model).resolve()
|
|
stem = source.name
|
|
passthrough = args.no_quant or args.outtype == "nvfp4"
|
|
if not args.output:
|
|
if passthrough and args.outtype == "auto":
|
|
directory = Path.cwd() / stem
|
|
elif passthrough:
|
|
label = "NVFP4" if args.outtype == "nvfp4" else args.outtype.upper()
|
|
directory = Path.cwd() / f"{stem}-{label}"
|
|
else:
|
|
directory = Path.cwd() / f"{stem}-NVFP4"
|
|
else:
|
|
directory = Path(args.output).resolve()
|
|
args.output = str(directory)
|
|
|
|
|
|
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:
|
|
args.exclude = profile.skip
|
|
log.info("%s profile: leaving %s unquantized.", profile.arch, profile.label)
|
|
else:
|
|
log.info("%s profile: --exclude replaces default skip of %s.", profile.arch, profile.label)
|
|
return
|
|
|
|
|
|
def _copy_non_weight_files(source_dir: Path, output_dir: Path) -> None:
|
|
"""Copy configs, tokenizers, chat templates, and documentation into output."""
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
for path in source_dir.iterdir():
|
|
if path.is_dir() or path.name.startswith("."):
|
|
continue
|
|
if path.name == "config.json" or path.name == INDEX_NAME:
|
|
continue
|
|
if any(path.name.endswith(ext) for ext in NON_WEIGHT_PATTERNS):
|
|
continue
|
|
dest = output_dir / path.name
|
|
if not dest.exists() or dest.stat().st_mtime < path.stat().st_mtime:
|
|
shutil.copy2(path, dest)
|
|
|
|
|
|
def _write_config_json(source_dir: Path, output_dir: Path, args, is_quantized: bool) -> None:
|
|
config_path = source_dir / "config.json"
|
|
if not config_path.is_file():
|
|
return
|
|
try:
|
|
config = json.loads(config_path.read_text(encoding="utf-8"))
|
|
except Exception as exc:
|
|
log.warning("Could not parse config.json: %s", exc)
|
|
return
|
|
|
|
fmt = getattr(args, "format", "hf")
|
|
|
|
def _build_ct_quant_config():
|
|
ignore = []
|
|
if not getattr(args, "quantize_output", False):
|
|
ignore.append("lm_head")
|
|
if getattr(args, "exclude", None):
|
|
ignore.append(args.exclude)
|
|
return {
|
|
"quant_method": "compressed-tensors",
|
|
"format": "nvfp4-pack-quantized",
|
|
"quantization_status": "compressed",
|
|
"ignore": ignore,
|
|
"config_groups": {
|
|
"group_0": {
|
|
"weights": {
|
|
"num_bits": 4,
|
|
"type": "float",
|
|
"strategy": "tensor_group",
|
|
"group_size": 16,
|
|
"symmetric": True,
|
|
"dynamic": False,
|
|
"actorder": None,
|
|
},
|
|
"input_activations": None,
|
|
"targets": ["Linear"],
|
|
"format": "nvfp4-pack-quantized",
|
|
}
|
|
},
|
|
}
|
|
|
|
if is_quantized:
|
|
if fmt == "hf":
|
|
config["quantization_config"] = _build_ct_quant_config()
|
|
else:
|
|
config["quantization_config"] = {
|
|
"quant_method": "modelopt",
|
|
"quant_algo": "NVFP4",
|
|
"format": "nvfp4-pack-quantized",
|
|
"group_size": 16,
|
|
"bits": 4,
|
|
}
|
|
config["torch_dtype"] = getattr(args, "output_dtype", "bfloat16")
|
|
elif args.no_quant:
|
|
if args.outtype in ("f32", "f16", "bf16"):
|
|
config.pop("quantization_config", None)
|
|
config["torch_dtype"] = "bfloat16" if args.outtype == "bf16" else ("float16" if args.outtype == "f16" else "float32")
|
|
elif args.outtype == "nvfp4":
|
|
if fmt == "hf":
|
|
config["quantization_config"] = _build_ct_quant_config()
|
|
else:
|
|
config.setdefault("quantization_config", {
|
|
"quant_method": "modelopt",
|
|
"quant_algo": "NVFP4",
|
|
"format": "nvfp4-pack-quantized",
|
|
"group_size": 16,
|
|
"bits": 4,
|
|
})
|
|
|
|
if getattr(args, "context_length", None):
|
|
for key in ("max_position_embeddings", "seq_length", "max_sequence_length"):
|
|
if key in config:
|
|
config[key] = int(args.context_length)
|
|
|
|
out_file = output_dir / "config.json"
|
|
out_file.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
def pack_nvfp4_hf(codes: np.ndarray, block_scales: np.ndarray, per_tensor: float):
|
|
"""Pack codes and scales into NVFP4 tensors."""
|
|
# Packed weights: low nibble is element 0, high nibble is element 1
|
|
packed = (codes[:, 0::2] & 0x0F) | ((codes[:, 1::2] & 0x0F) << 4)
|
|
packed_tensor = torch.from_numpy(np.ascontiguousarray(packed, dtype=np.uint8))
|
|
|
|
scale_bits = ue4m3_bits(block_scales)
|
|
scale_tensor = torch.from_numpy(np.ascontiguousarray(scale_bits, dtype=np.uint8)).view(torch.float8_e4m3fn)
|
|
scale_2_tensor = torch.tensor(float(per_tensor), dtype=torch.float32)
|
|
|
|
return packed_tensor, scale_tensor, scale_2_tensor
|
|
|
|
|
|
def _should_quantize(
|
|
key: str,
|
|
shape: tuple[int, ...],
|
|
args,
|
|
convert_to_quant_dir: Path | None,
|
|
) -> bool:
|
|
if len(shape) != 2:
|
|
return False
|
|
if not (key.endswith(".weight") or key.endswith(".weight_packed")):
|
|
return False
|
|
if shape[-1] % 64 != 0 or shape[-1] < 64:
|
|
return False
|
|
if key in ("lm_head.weight", "lm_head.weight_packed") and not getattr(args, "quantize_output", False):
|
|
return False
|
|
lowered = key.lower()
|
|
if "embed_tokens" in lowered or "wte" in lowered or "word_embeddings" in lowered:
|
|
return False
|
|
if "norm" in lowered or "ln" in lowered:
|
|
return False
|
|
if getattr(args, "exclude_re", None) and args.exclude_re.search(key):
|
|
return False
|
|
if ctq_skips_key(key, convert_to_quant_dir):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _target_dtype(dtype_str: str) -> torch.dtype:
|
|
mapping = {
|
|
"float32": torch.float32,
|
|
"f32": torch.float32,
|
|
"float16": torch.float16,
|
|
"f16": torch.float16,
|
|
"bfloat16": torch.bfloat16,
|
|
"bf16": torch.bfloat16,
|
|
}
|
|
return mapping.get(dtype_str.lower(), torch.bfloat16)
|
|
|
|
|
|
def export_directory(args) -> None:
|
|
source_dir = Path(args.model).resolve()
|
|
if not source_dir.is_dir():
|
|
raise FileNotFoundError(f"Model directory does not exist: {source_dir}")
|
|
|
|
_assign_output(args)
|
|
output_dir = Path(args.output).resolve()
|
|
if output_dir == source_dir:
|
|
raise SystemExit("Output directory cannot be the same as the input directory.")
|
|
if output_dir.exists() and any(output_dir.iterdir()) and not args.overwrite:
|
|
raise SystemExit(f"Output directory {output_dir} already exists and is not empty. Use --overwrite to replace.")
|
|
|
|
shards = discover_shards(source_dir)
|
|
_apply_directory_profile(args, source_dir)
|
|
args.exclude_re = re.compile(args.exclude) if args.exclude else None
|
|
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
log.info("Loading %s", source_dir)
|
|
log.info("Found %d safetensors shard(s)", len(shards))
|
|
log.info("Target format: %s", "Hugging Face (compressed-tensors)" if args.format == "hf" else "ModelOpt / ComfyUI")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
_copy_non_weight_files(source_dir, output_dir)
|
|
|
|
is_passthrough = args.no_quant or args.outtype == "nvfp4"
|
|
float_dtype = _target_dtype(args.output_dtype if not is_passthrough else (args.outtype if args.outtype != "auto" else "bfloat16"))
|
|
|
|
if is_passthrough:
|
|
log.info("Quantization skipped. Converting tensors to %s", args.outtype)
|
|
weight_map: dict[str, str] = {}
|
|
total_size = 0
|
|
for shard_name in shards:
|
|
log.info("Processing %s", shard_name)
|
|
src_shard = source_dir / shard_name
|
|
out_shard = output_dir / shard_name
|
|
out_tensors: dict[str, torch.Tensor] = {}
|
|
with safe_open(src_shard, framework="pt", device="cpu") as handle:
|
|
for key in handle.keys():
|
|
t = handle.get_tensor(key)
|
|
if args.outtype in ("f32", "f16", "bf16") and t.is_floating_point():
|
|
t = t.to(dtype=float_dtype)
|
|
out_tensors[key] = t
|
|
weight_map[key] = shard_name
|
|
total_size += t.numel() * t.element_size()
|
|
save_file(out_tensors, out_shard)
|
|
_write_config_json(source_dir, output_dir, args, is_quantized=False)
|
|
if len(shards) > 1:
|
|
(output_dir / INDEX_NAME).write_text(
|
|
json.dumps({"metadata": {"total_size": total_size}, "weight_map": weight_map}, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
log.info("Wrote %s", output_dir)
|
|
return
|
|
|
|
# Learned NVFP4 Quantization
|
|
ensure_cuda_interpreter(args.device)
|
|
args.optimizer = resolve_optimizer(args.optimizer)
|
|
|
|
# Prepare checkpoint directory for per-tensor resume
|
|
checkpoint_dir = Path(args.checkpoint_dir) if args.checkpoint_dir else output_dir / ".nvfp4-work"
|
|
if args.fresh and checkpoint_dir.exists():
|
|
shutil.rmtree(checkpoint_dir)
|
|
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Initialize quantizer
|
|
import gguf
|
|
quantizer = TensorQuantizer(args, gguf, load_learned_converter(args.convert_to_quant)[1])
|
|
|
|
# First pass: collect all tensor keys and shapes across shards
|
|
all_keys: dict[str, tuple[str, tuple[int, ...]]] = {}
|
|
for shard_name in shards:
|
|
with safe_open(source_dir / shard_name, framework="pt", device="cpu") as handle:
|
|
for key in handle.keys():
|
|
slice_header = handle.get_slice(key)
|
|
all_keys[key] = (shard_name, tuple(int(d) for d in slice_header.get_shape()))
|
|
|
|
quant_count = 0
|
|
passthrough_count = 0
|
|
bias_count = 0
|
|
weight_map: dict[str, str] = {}
|
|
total_size = 0
|
|
all_quant_meta: dict[str, dict] = {}
|
|
|
|
for s_idx, shard_name in enumerate(shards, start=1):
|
|
log.info("\n[%d/%d] Processing shard: %s", s_idx, len(shards), shard_name)
|
|
src_shard = source_dir / shard_name
|
|
out_shard = output_dir / shard_name
|
|
shard_tensors: dict[str, torch.Tensor] = {}
|
|
|
|
with safe_open(src_shard, framework="pt", device="cpu") as handle:
|
|
keys_in_shard = list(handle.keys())
|
|
|
|
for key in keys_in_shard:
|
|
tensor = handle.get_tensor(key)
|
|
shape = tuple(int(d) for d in tensor.shape)
|
|
base_key = key.removesuffix(".weight_packed").removesuffix(".weight")
|
|
|
|
if _should_quantize(key, shape, args, args.convert_to_quant):
|
|
log.info(" Quantizing [%s] %s", "x".join(str(d) for d in shape), key)
|
|
weight_np = tensor.detach().float().cpu().numpy()
|
|
|
|
# Check for sibling bias
|
|
bias_key = f"{base_key}.bias"
|
|
has_bias = bias_key in all_keys and not args.no_bias_correction
|
|
bias_shard = all_keys[bias_key][0] if has_bias else None
|
|
|
|
# Check checkpoint cache
|
|
cached = _load_checkpoint(checkpoint_dir, key, shape, need_bias=has_bias)
|
|
if cached is not None:
|
|
codes, block_scale, per_tensor, bias_corr = cached
|
|
log.info(" Loaded from checkpoint: %s", key)
|
|
else:
|
|
rounds = getattr(args, "scale_refinement_rounds", DEFAULT_SCALE_REFINEMENT_ROUNDS)
|
|
codes, block_scale, per_tensor, dequant = quantizer.quantize_matrix(weight_np, key)
|
|
codes, block_scale = coordinate_descent_refine(
|
|
weight_np, codes, block_scale, per_tensor, bake=False, rounds=rounds, device=quantizer.device,
|
|
)
|
|
codes, block_scale, per_tensor = sweep_per_tensor_scale(
|
|
weight_np, codes, block_scale, per_tensor, bake=False, device=quantizer.device,
|
|
)
|
|
|
|
bias_corr = None
|
|
if has_bias:
|
|
bias_corr = quantizer.bias_correction(weight_np, dequant)
|
|
|
|
_save_checkpoint(checkpoint_dir, key, shape, codes, block_scale, per_tensor, bias_corr)
|
|
|
|
packed_w, scale_1, scale_2 = pack_nvfp4_hf(codes, block_scale, per_tensor)
|
|
|
|
if args.format == "hf":
|
|
# Standard Hugging Face compressed-tensors layout
|
|
shard_tensors[f"{base_key}.weight_packed"] = packed_w
|
|
shard_tensors[f"{base_key}.weight_scale"] = scale_1
|
|
shard_tensors[f"{base_key}.weight_global_scale"] = torch.tensor(1.0 / float(per_tensor), dtype=torch.float32)
|
|
else:
|
|
# ModelOpt / ComfyUI layout
|
|
shard_tensors[f"{base_key}.weight"] = packed_w
|
|
shard_tensors[f"{base_key}.weight_scale"] = scale_1
|
|
shard_tensors[f"{base_key}.weight_scale_2"] = scale_2
|
|
shard_tensors[f"{base_key}.comfy_quant"] = torch.tensor([1], dtype=torch.uint8)
|
|
|
|
quant_count += 1
|
|
all_quant_meta[base_key] = {
|
|
"quant_type": "nvfp4",
|
|
"block_size": 16,
|
|
"scale_type": "float8_e4m3fn",
|
|
}
|
|
|
|
# Bias correction application if bias is in the same shard
|
|
if has_bias and bias_shard == shard_name and bias_corr is not None:
|
|
bias_t = handle.get_tensor(bias_key).float()
|
|
corrected_bias = bias_t - torch.from_numpy(bias_corr).to(bias_t.device)
|
|
shard_tensors[bias_key] = corrected_bias.to(dtype=float_dtype)
|
|
bias_count += 1
|
|
else:
|
|
# Non-quantized tensors (or bias already processed)
|
|
if key in shard_tensors:
|
|
continue
|
|
if tensor.is_floating_point():
|
|
if tensor.ndim <= 1 and key.endswith("norm.weight"):
|
|
tensor = tensor.float()
|
|
else:
|
|
tensor = tensor.to(dtype=float_dtype)
|
|
shard_tensors[key] = tensor
|
|
passthrough_count += 1
|
|
|
|
shard_metadata = {}
|
|
if args.format != "hf":
|
|
shard_metadata["_quantization_metadata"] = json.dumps({"format_version": "1.0", "layers": all_quant_meta}, separators=(",", ":"))
|
|
|
|
for k, t in shard_tensors.items():
|
|
weight_map[k] = shard_name
|
|
total_size += t.numel() * t.element_size()
|
|
|
|
save_file(shard_tensors, out_shard, metadata=shard_metadata if shard_metadata else None)
|
|
|
|
_write_config_json(source_dir, output_dir, args, is_quantized=True)
|
|
|
|
if len(shards) > 1:
|
|
index_payload = {
|
|
"metadata": {"total_size": total_size},
|
|
"weight_map": weight_map,
|
|
}
|
|
(output_dir / INDEX_NAME).write_text(json.dumps(index_payload, indent=2) + "\n", encoding="utf-8")
|
|
|
|
if not args.keep_checkpoints and quant_count:
|
|
shutil.rmtree(checkpoint_dir, ignore_errors=True)
|
|
|
|
log.info(
|
|
"\nWrote %s (%d NVFP4 tensors, %d biases adjusted, %d other tensors)",
|
|
output_dir,
|
|
quant_count,
|
|
bias_count,
|
|
passthrough_count,
|
|
)
|
|
|
|
|
|
def self_test(args) -> None:
|
|
"""Check that packed blocks roundtrip accurately with dequantizer."""
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
from nvfp4gguf import FP4_TABLE
|
|
import gguf
|
|
|
|
log.info("Running NVFP4 roundtrip self-test...")
|
|
_converter_cls, unpack_uint4 = load_learned_converter(args.convert_to_quant)
|
|
quantizer = TensorQuantizer(args, gguf, unpack_uint4)
|
|
|
|
generator = np.random.default_rng(0)
|
|
for shape in ((64, 128), (128, 256)):
|
|
weight = generator.normal(size=shape).astype(np.float32) * 0.05
|
|
codes, block_scale, per_tensor, dequant = quantizer.quantize_matrix(weight)
|
|
packed_w, scale_1, scale_2 = pack_nvfp4_hf(codes, block_scale, per_tensor)
|
|
|
|
# Dequantize packed tensor
|
|
w_u = packed_w.numpy()
|
|
c0 = w_u & 0x0F
|
|
c1 = (w_u >> 4) & 0x0F
|
|
unpacked_codes = np.empty(shape, dtype=np.uint8)
|
|
unpacked_codes[:, 0::2] = c0
|
|
unpacked_codes[:, 1::2] = c1
|
|
assert np.array_equal(codes, unpacked_codes), "Nibble unpacking mismatch!"
|
|
|
|
# Check values
|
|
fp4_vals = FP4_TABLE[unpacked_codes]
|
|
b_scales = block_scale.repeat(16, axis=1)
|
|
recon = fp4_vals * b_scales * per_tensor
|
|
max_diff = np.max(np.abs(recon - dequant))
|
|
assert max_diff < 1e-4, f"Reconstruction difference {max_diff} exceeded tolerance!"
|
|
log.info("Shape %s roundtrip ok (max diff %.3e)", shape, max_diff)
|
|
|
|
log.info("self-test passed successfully.")
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Quantize a Hugging Face model directory to NVFP4 Hugging Face format.",
|
|
)
|
|
parser.add_argument("model", nargs="?", help="Hugging Face model directory (config.json + safetensors)")
|
|
parser.add_argument("-o", "--output", help="Output directory. Default: ./<model-name>-NVFP4")
|
|
parser.add_argument("--format", default="hf", choices=("hf", "modelopt", "comfy"),
|
|
help="Safetensors tensor layout: 'hf' (default) uses Hugging Face compressed-tensors (.weight_packed, .weight_scale, .weight_global_scale). 'modelopt' (or 'comfy') uses raw .weight.")
|
|
parser.add_argument("--overwrite", action="store_true", help="Replace existing files in the output directory")
|
|
parser.add_argument("--no-quant", action="store_true", help="Skip learned rounding. Rewrites or repacks tensors")
|
|
parser.add_argument("--outtype", default="auto", choices=("auto", "f32", "f16", "bf16", "nvfp4"),
|
|
help="Float type for tensors when --no-quant is used (default: auto)")
|
|
parser.add_argument("--output-dtype", default="bfloat16", choices=("bfloat16", "float16"),
|
|
help="Dtype for unquantized 2D weights and embeddings (default: bfloat16)")
|
|
parser.add_argument("--context-length", type=int, default=None,
|
|
help="Override context length stored in config.json")
|
|
parser.add_argument("--model-name", default=None, help="Override model name")
|
|
parser.add_argument("--checkpoint-dir", type=Path, default=None, help="Per-tensor resume directory")
|
|
parser.add_argument("--fresh", action="store_true", help="Delete checkpoint directory before starting")
|
|
parser.add_argument("--keep-checkpoints", action="store_true", help="Keep per-tensor checkpoints upon completion")
|
|
parser.add_argument("--eager", action="store_true", help="Eager loading")
|
|
parser.add_argument("--low-memory", action="store_true", help="Stream each shard instead of loading all at once")
|
|
parser.add_argument("--heur", action="store_true", help="Skip layers heuristic rejects")
|
|
parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
|
|
parser.add_argument("--self-test", action="store_true", help="Run NVFP4 packing self-test and exit")
|
|
parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP,
|
|
help="llama.cpp checkout directory (default: ~/llama.cpp)")
|
|
|
|
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.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:
|
|
args = parse_args(argv)
|
|
if args.self_test:
|
|
self_test(args)
|
|
return
|
|
export_directory(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|