762 lines
33 KiB
Python
Executable File
762 lines
33 KiB
Python
Executable File
"""Convert and quantize a GGUF model to NVFP4 Hugging Face format.
|
|
|
|
Reads a GGUF file and writes a Hugging Face format model directory (config.json +
|
|
safetensors). 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.
|
|
|
|
If the source GGUF contains float or non-NVFP4 weights, weight matrices are quantized
|
|
to NVFP4 using convert_to_quant's learned rounding (or simple round-to-grid with --simple).
|
|
If the source GGUF is already NVFP4, weights are repacked into safetensors without loss
|
|
of precision.
|
|
|
|
If --no-quant is specified, weights are dequantized and written to safetensors in the
|
|
target float format (--outtype f32, f16, bf16, or auto).
|
|
"""
|
|
|
|
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.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"
|
|
MAX_SHARD_BYTES = 5 * 1024 * 1024 * 1024 # 5 GB max shard size
|
|
|
|
log = logging.getLogger("gguf_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 _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 _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)
|
|
elif qtype == gguf_module.GGMLQuantizationType.BF16:
|
|
values = np.asarray(tensor.data, dtype=np.float32)
|
|
else:
|
|
values = np.asarray(gguf_module.quants.dequantize(tensor.data, qtype), dtype=np.float32)
|
|
return np.array(values, dtype=np.float32, copy=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 _assign_output(args) -> None:
|
|
source = Path(args.gguf)
|
|
stem = source.stem
|
|
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 pack_nvfp4_hf(codes: np.ndarray, block_scales: np.ndarray, per_tensor: float):
|
|
"""Pack codes and scales into NVFP4 tensors."""
|
|
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 unpack_gguf_nvfp4_to_hf(raw: np.ndarray, per_tensor: float):
|
|
"""Convert GGUF NVFP4 super-blocks (36 bytes / 64 values) into ModelOpt tensors."""
|
|
rows = raw.shape[0]
|
|
n_super = raw.shape[1] // 36
|
|
n_blocks = n_super * 4
|
|
cols = n_blocks * 16
|
|
|
|
raw_t = torch.from_numpy(raw).reshape(rows, n_super, 36)
|
|
d_ue = raw_t[:, :, :4].reshape(rows, n_blocks)
|
|
qs = raw_t[:, :, 4:].reshape(rows, n_blocks, 8)
|
|
|
|
vals = torch.empty((rows, n_blocks, 16), dtype=torch.uint8)
|
|
vals[:, :, :8] = qs & 0x0F
|
|
vals[:, :, 8:] = qs >> 4
|
|
|
|
w = ((vals[:, :, 0::2] & 0x0F) | ((vals[:, :, 1::2] & 0x0F) << 4)).reshape(rows, cols // 2)
|
|
|
|
scale_tensor = d_ue.contiguous().view(torch.float8_e4m3fn)
|
|
scale_2_tensor = torch.tensor(float(per_tensor), dtype=torch.float32)
|
|
return w, scale_tensor, scale_2_tensor
|
|
|
|
|
|
def unpermute_rope(weights: np.ndarray, n_head: int, n_kv_head: int | None = None) -> np.ndarray:
|
|
"""Invert RoPE permutation applied during HF to GGUF conversion."""
|
|
head_count = n_kv_head if n_kv_head is not None else n_head
|
|
if head_count <= 0 or weights.shape[0] % head_count != 0:
|
|
return weights
|
|
head_dim = weights.shape[0] // head_count
|
|
if head_dim % 2 != 0:
|
|
return weights
|
|
w = weights.reshape(head_count, head_dim // 2, 2, *weights.shape[1:])
|
|
w = np.swapaxes(w, 1, 2)
|
|
return w.reshape(weights.shape)
|
|
|
|
|
|
def build_tensor_name_map(arch_name: str, block_count: int, gguf_module) -> dict[str, str]:
|
|
"""Build a mapping from GGUF tensor names to standard Hugging Face names."""
|
|
arch_enum = getattr(gguf_module.MODEL_ARCH, arch_name.upper(), None)
|
|
mapping: dict[str, str] = {}
|
|
|
|
if arch_enum is not None:
|
|
try:
|
|
tm = gguf_module.get_tensor_name_map(arch_enum, block_count)
|
|
rev: dict[str, list[str]] = {}
|
|
for hf_key, (_tensor_enum, gguf_name) in tm.mapping.items():
|
|
if hf_key == gguf_name:
|
|
continue
|
|
if gguf_name not in rev:
|
|
rev[gguf_name] = []
|
|
rev[gguf_name].append(hf_key)
|
|
|
|
for g_name, hf_candidates in rev.items():
|
|
preferred = [c for c in hf_candidates if c.startswith("model.") or c.startswith("transformer.")]
|
|
mapping[g_name] = preferred[0] if preferred else hf_candidates[0]
|
|
except Exception:
|
|
pass
|
|
|
|
defaults = {
|
|
"token_embd": "model.embed_tokens",
|
|
"output_norm": "model.norm",
|
|
"output": "lm_head",
|
|
}
|
|
for g, h in defaults.items():
|
|
mapping.setdefault(g, h)
|
|
|
|
for bid in range(block_count):
|
|
bid_defaults = {
|
|
f"blk.{bid}.attn_q": f"model.layers.{bid}.self_attn.q_proj",
|
|
f"blk.{bid}.attn_k": f"model.layers.{bid}.self_attn.k_proj",
|
|
f"blk.{bid}.attn_v": f"model.layers.{bid}.self_attn.v_proj",
|
|
f"blk.{bid}.attn_output": f"model.layers.{bid}.self_attn.o_proj",
|
|
f"blk.{bid}.attn_norm": f"model.layers.{bid}.input_layernorm",
|
|
f"blk.{bid}.ffn_gate": f"model.layers.{bid}.mlp.gate_proj",
|
|
f"blk.{bid}.ffn_up": f"model.layers.{bid}.mlp.up_proj",
|
|
f"blk.{bid}.ffn_down": f"model.layers.{bid}.mlp.down_proj",
|
|
f"blk.{bid}.ffn_norm": f"model.layers.{bid}.post_attention_layernorm",
|
|
f"blk.{bid}.attn_q_norm": f"model.layers.{bid}.self_attn.q_norm",
|
|
f"blk.{bid}.attn_k_norm": f"model.layers.{bid}.self_attn.k_norm",
|
|
f"blk.{bid}.post_ffw_norm": f"model.layers.{bid}.post_feedforward_layernorm",
|
|
f"blk.{bid}.pre_ffw_norm": f"model.layers.{bid}.pre_feedforward_layernorm",
|
|
}
|
|
for g, h in bid_defaults.items():
|
|
mapping.setdefault(g, h)
|
|
|
|
return mapping
|
|
|
|
|
|
def _map_gguf_name(gguf_name: str, name_map: dict[str, str]) -> str:
|
|
base, _, suffix = gguf_name.rpartition(".")
|
|
if base in name_map:
|
|
return f"{name_map[base]}.{suffix}"
|
|
return gguf_name
|
|
|
|
|
|
def _synthesize_hf_configs(reader, output_dir: Path, arch: str, args, is_quantized: bool) -> None:
|
|
"""Generate config.json and tokenizer files from GGUF metadata."""
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
fields = reader.fields
|
|
|
|
def get_field(key: str, default=None):
|
|
f = fields.get(key)
|
|
return _pythonize(f.contents()) if f is not None else default
|
|
|
|
arch_map = {
|
|
"llama": ("LlamaForCausalLM", "llama"),
|
|
"qwen2": ("Qwen2ForCausalLM", "qwen2"),
|
|
"gemma": ("GemmaForCausalLM", "gemma"),
|
|
"gemma2": ("Gemma2ForCausalLM", "gemma2"),
|
|
"mistral": ("MistralForCausalLM", "mistral"),
|
|
"phi3": ("Phi3ForCausalLM", "phi3"),
|
|
"deepseek2": ("DeepseekV2ForCausalLM", "deepseek_v2"),
|
|
"gptneox": ("GPTNeoXForCausalLM", "gpt_neox"),
|
|
}
|
|
class_name, model_type = arch_map.get(arch.lower(), (f"{arch.capitalize()}ForCausalLM", arch.lower()))
|
|
|
|
tokens = get_field("tokenizer.ggml.tokens", [])
|
|
vocab_size = len(tokens) if tokens else get_field(f"{arch}.vocab_size", 32000)
|
|
|
|
config = {
|
|
"architectures": [class_name],
|
|
"model_type": model_type,
|
|
"vocab_size": vocab_size,
|
|
"hidden_size": get_field(f"{arch}.embedding_length", 4096),
|
|
"intermediate_size": get_field(f"{arch}.feed_forward_length", 11008),
|
|
"num_hidden_layers": get_field(f"{arch}.block_count", 32),
|
|
"num_attention_heads": get_field(f"{arch}.attention.head_count", 32),
|
|
"num_key_value_heads": get_field(f"{arch}.attention.head_count_kv", get_field(f"{arch}.attention.head_count", 32)),
|
|
"max_position_embeddings": getattr(args, "context_length", None) or get_field(f"{arch}.context_length", 4096),
|
|
"rms_norm_eps": get_field(f"{arch}.attention.layer_norm_rms_epsilon", 1e-5),
|
|
"rope_theta": get_field(f"{arch}.rope.freq_base", 10000.0),
|
|
"bos_token_id": get_field("tokenizer.ggml.bos_token_id", 1),
|
|
"eos_token_id": get_field("tokenizer.ggml.eos_token_id", 2),
|
|
"torch_dtype": getattr(args, "output_dtype", "bfloat16"),
|
|
}
|
|
|
|
fmt = getattr(args, "format", "hf")
|
|
ignore = []
|
|
if not getattr(args, "quantize_output", False):
|
|
ignore.append("lm_head")
|
|
if getattr(args, "exclude", None):
|
|
ignore.append(args.exclude)
|
|
|
|
if is_quantized:
|
|
if fmt == "hf":
|
|
config["quantization_config"] = {
|
|
"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",
|
|
}
|
|
},
|
|
}
|
|
else:
|
|
config["quantization_config"] = {
|
|
"quant_method": "modelopt",
|
|
"quant_algo": "NVFP4",
|
|
"format": "nvfp4-pack-quantized",
|
|
"group_size": 16,
|
|
"bits": 4,
|
|
}
|
|
|
|
(output_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
|
|
|
|
# Tokenizer configs
|
|
chat_template = get_field("tokenizer.chat_template")
|
|
tok_config = {
|
|
"bos_token_id": config["bos_token_id"],
|
|
"eos_token_id": config["eos_token_id"],
|
|
"model_max_length": config["max_position_embeddings"],
|
|
}
|
|
if chat_template:
|
|
tok_config["chat_template"] = _text(chat_template)
|
|
(output_dir / "tokenizer_config.json").write_text(json.dumps(tok_config, indent=2) + "\n", encoding="utf-8")
|
|
|
|
special_tokens = {}
|
|
if tokens:
|
|
bos_id = config["bos_token_id"]
|
|
eos_id = config["eos_token_id"]
|
|
if 0 <= bos_id < len(tokens):
|
|
special_tokens["bos_token"] = _text(tokens[bos_id])
|
|
if 0 <= eos_id < len(tokens):
|
|
special_tokens["eos_token"] = _text(tokens[eos_id])
|
|
vocab = {_text(t): i for i, t in enumerate(tokens)}
|
|
(output_dir / "vocab.json").write_text(json.dumps(vocab, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
(output_dir / "special_tokens_map.json").write_text(json.dumps(special_tokens, indent=2) + "\n", encoding="utf-8")
|
|
|
|
merges = get_field("tokenizer.ggml.merges")
|
|
if merges:
|
|
merge_lines = [_text(m) for m in merges]
|
|
(output_dir / "merges.txt").write_text("#version: 0.2\n" + "\n".join(merge_lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
def _copy_template_files(template_dir: Path, output_dir: Path, args, is_quantized: bool) -> None:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
fmt = getattr(args, "format", "hf")
|
|
|
|
for path in template_dir.iterdir():
|
|
if path.is_dir() or path.name.startswith("."):
|
|
continue
|
|
if any(path.name.endswith(ext) for ext in (".safetensors", ".bin", ".pt", ".pth", ".h5")):
|
|
continue
|
|
if path.name == INDEX_NAME:
|
|
continue
|
|
dest = output_dir / path.name
|
|
if path.name == "config.json":
|
|
try:
|
|
cfg = json.loads(path.read_text(encoding="utf-8"))
|
|
if is_quantized:
|
|
if fmt == "hf":
|
|
cfg["quantization_config"] = {
|
|
"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",
|
|
}
|
|
},
|
|
}
|
|
else:
|
|
cfg["quantization_config"] = {
|
|
"quant_method": "modelopt",
|
|
"quant_algo": "NVFP4",
|
|
"format": "nvfp4-pack-quantized",
|
|
"group_size": 16,
|
|
"bits": 4,
|
|
}
|
|
cfg["torch_dtype"] = getattr(args, "output_dtype", "bfloat16")
|
|
if getattr(args, "context_length", None):
|
|
cfg["max_position_embeddings"] = int(args.context_length)
|
|
dest.write_text(json.dumps(cfg, indent=2) + "\n", encoding="utf-8")
|
|
continue
|
|
except Exception:
|
|
pass
|
|
shutil.copy2(path, dest)
|
|
|
|
|
|
def export_gguf_to_hf(args) -> None:
|
|
gguf_path = Path(args.gguf).resolve()
|
|
if not gguf_path.is_file():
|
|
raise FileNotFoundError(f"GGUF file does not exist: {gguf_path}")
|
|
|
|
_assign_output(args)
|
|
output_dir = Path(args.output).resolve()
|
|
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.")
|
|
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
log.info("Loading GGUF file: %s", gguf_path)
|
|
log.info("Target format: %s", "Hugging Face (compressed-tensors)" if args.format == "hf" else "ModelOpt / ComfyUI")
|
|
|
|
import gguf
|
|
reader = gguf.GGUFReader(gguf_path)
|
|
arch = _text(reader.fields.get("general.architecture", "llama").contents())
|
|
block_count = 32
|
|
for k in (f"{arch}.block_count", "general.block_count"):
|
|
if k in reader.fields:
|
|
block_count = int(_pythonize(reader.fields[k].contents()))
|
|
break
|
|
|
|
n_head = int(_pythonize(reader.fields.get(f"{arch}.attention.head_count", 32).contents()))
|
|
n_kv_head = int(_pythonize(reader.fields.get(f"{arch}.attention.head_count_kv", n_head).contents()))
|
|
should_unpermute = arch.lower() in ("llama", "mistral", "nanbeige")
|
|
|
|
name_map = build_tensor_name_map(arch, block_count, gguf)
|
|
profile = quant_profile(arch)
|
|
if profile is not None:
|
|
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)
|
|
args.exclude_re = re.compile(args.exclude) if args.exclude else None
|
|
|
|
is_passthrough = args.no_quant or args.outtype == "nvfp4"
|
|
is_quantized = not is_passthrough
|
|
float_dtype = _target_dtype(args.output_dtype if not is_passthrough else (args.outtype if args.outtype != "auto" else "bfloat16"))
|
|
|
|
# Config & Tokenizer generation / copying
|
|
if args.hf_template and Path(args.hf_template).is_dir():
|
|
log.info("Copying configs and tokenizers from template: %s", args.hf_template)
|
|
_copy_template_files(Path(args.hf_template), output_dir, args, is_quantized=is_quantized)
|
|
else:
|
|
log.info("Synthesizing Hugging Face configs and tokenizer from GGUF metadata")
|
|
_synthesize_hf_configs(reader, output_dir, arch, args, is_quantized=is_quantized)
|
|
|
|
# Scale tensors cache for NVFP4 repacking
|
|
scale_tensors: dict[str, float] = {}
|
|
for tensor in reader.tensors:
|
|
if tensor.name.endswith(".scale"):
|
|
base = tensor.name.removesuffix(".scale")
|
|
scale_val = float(np.asarray(tensor.data, dtype=np.float32).reshape(-1)[0])
|
|
scale_tensors[base] = scale_val
|
|
|
|
# Checkpoint dir
|
|
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)
|
|
|
|
quantizer = None
|
|
if is_quantized:
|
|
ensure_cuda_interpreter(args.device)
|
|
args.optimizer = resolve_optimizer(args.optimizer)
|
|
quantizer = TensorQuantizer(args, gguf, load_learned_converter(args.convert_to_quant)[1])
|
|
|
|
quant_count = 0
|
|
passthrough_count = 0
|
|
bias_count = 0
|
|
all_quant_meta: dict[str, dict] = {}
|
|
all_hf_tensors: dict[str, torch.Tensor] = {}
|
|
|
|
log.info("Processing %d tensor(s) from GGUF", len(reader.tensors))
|
|
|
|
for tensor in reader.tensors:
|
|
g_name = tensor.name
|
|
if g_name.endswith(".scale") or g_name.endswith(".input_scale"):
|
|
continue
|
|
|
|
hf_name = _map_gguf_name(g_name, name_map)
|
|
base_name = hf_name.removesuffix(".weight")
|
|
shape = tuple(int(d) for d in tensor.shape[::-1]) # GGUF shapes are reversed
|
|
|
|
# If already NVFP4 in GGUF
|
|
if tensor.tensor_type == gguf.GGMLQuantizationType.NVFP4:
|
|
log.info(" Repacking NVFP4 [%s] %s -> %s", "x".join(str(d) for d in shape), g_name, hf_name)
|
|
raw = np.asarray(tensor.data, dtype=np.uint8).reshape(shape[0], -1)
|
|
per_tensor = scale_tensors.get(g_name, 1.0)
|
|
packed_w, scale_1, scale_2 = unpack_gguf_nvfp4_to_hf(raw, per_tensor)
|
|
|
|
if args.format == "hf":
|
|
all_hf_tensors[f"{base_name}.weight_packed"] = packed_w
|
|
all_hf_tensors[f"{base_name}.weight_scale"] = scale_1
|
|
all_hf_tensors[f"{base_name}.weight_global_scale"] = torch.tensor(1.0 / float(per_tensor), dtype=torch.float32)
|
|
else:
|
|
all_hf_tensors[hf_name] = packed_w
|
|
all_hf_tensors[f"{base_name}.weight_scale"] = scale_1
|
|
all_hf_tensors[f"{base_name}.weight_scale_2"] = scale_2
|
|
all_hf_tensors[f"{base_name}.comfy_quant"] = torch.tensor([1], dtype=torch.uint8)
|
|
|
|
all_quant_meta[base_name] = {"quant_type": "nvfp4", "block_size": 16, "scale_type": "float8_e4m3fn"}
|
|
quant_count += 1
|
|
continue
|
|
|
|
# Convert tensor data to float
|
|
data_float = _to_float(tensor, gguf)
|
|
|
|
# Unpermute RoPE for Q and K if needed
|
|
if should_unpermute:
|
|
if g_name.endswith("attn_q.weight"):
|
|
data_float = unpermute_rope(data_float, n_head, n_head)
|
|
elif g_name.endswith("attn_k.weight"):
|
|
data_float = unpermute_rope(data_float, n_head, n_kv_head)
|
|
|
|
# Check eligibility for NVFP4 quantization
|
|
is_eligible = (
|
|
is_quantized
|
|
and len(shape) == 2
|
|
and hf_name.endswith(".weight")
|
|
and shape[-1] >= 64
|
|
and shape[-1] % 64 == 0
|
|
and not ("norm" in hf_name.lower() or "embed_tokens" in hf_name.lower())
|
|
and (hf_name != "lm_head.weight" or getattr(args, "quantize_output", False))
|
|
and not (args.exclude_re and args.exclude_re.search(hf_name))
|
|
and not ctq_skips_key(hf_name, args.convert_to_quant)
|
|
)
|
|
|
|
if is_eligible and quantizer is not None:
|
|
log.info(" Quantizing [%s] %s -> %s", "x".join(str(d) for d in shape), g_name, hf_name)
|
|
has_bias = False
|
|
|
|
cached = _load_checkpoint(checkpoint_dir, hf_name, shape, need_bias=has_bias)
|
|
if cached is not None:
|
|
codes, block_scale, per_tensor, bias_corr = cached
|
|
log.info(" Loaded from checkpoint: %s", hf_name)
|
|
else:
|
|
rounds = getattr(args, "scale_refinement_rounds", DEFAULT_SCALE_REFINEMENT_ROUNDS)
|
|
codes, block_scale, per_tensor, dequant = quantizer.quantize_matrix(data_float, hf_name)
|
|
codes, block_scale = coordinate_descent_refine(
|
|
data_float, codes, block_scale, per_tensor, bake=False, rounds=rounds, device=quantizer.device,
|
|
)
|
|
codes, block_scale, per_tensor = sweep_per_tensor_scale(
|
|
data_float, codes, block_scale, per_tensor, bake=False, device=quantizer.device,
|
|
)
|
|
_save_checkpoint(checkpoint_dir, hf_name, shape, codes, block_scale, per_tensor, None)
|
|
|
|
packed_w, scale_1, scale_2 = pack_nvfp4_hf(codes, block_scale, per_tensor)
|
|
|
|
if args.format == "hf":
|
|
all_hf_tensors[f"{base_name}.weight_packed"] = packed_w
|
|
all_hf_tensors[f"{base_name}.weight_scale"] = scale_1
|
|
all_hf_tensors[f"{base_name}.weight_global_scale"] = torch.tensor(1.0 / float(per_tensor), dtype=torch.float32)
|
|
else:
|
|
all_hf_tensors[hf_name] = packed_w
|
|
all_hf_tensors[f"{base_name}.weight_scale"] = scale_1
|
|
all_hf_tensors[f"{base_name}.weight_scale_2"] = scale_2
|
|
all_hf_tensors[f"{base_name}.comfy_quant"] = torch.tensor([1], dtype=torch.uint8)
|
|
|
|
all_quant_meta[base_name] = {"quant_type": "nvfp4", "block_size": 16, "scale_type": "float8_e4m3fn"}
|
|
quant_count += 1
|
|
else:
|
|
# Float tensor passthrough
|
|
tensor_t = torch.from_numpy(data_float)
|
|
if tensor_t.ndim <= 1 and hf_name.endswith("norm.weight"):
|
|
tensor_t = tensor_t.float()
|
|
else:
|
|
tensor_t = tensor_t.to(dtype=float_dtype)
|
|
all_hf_tensors[hf_name] = tensor_t
|
|
passthrough_count += 1
|
|
|
|
# Shard writing logic
|
|
total_bytes = sum(t.numel() * t.element_size() for t in all_hf_tensors.values())
|
|
metadata = {}
|
|
if args.format != "hf":
|
|
metadata["_quantization_metadata"] = json.dumps({"format_version": "1.0", "layers": all_quant_meta}, separators=(",", ":"))
|
|
|
|
if total_bytes < MAX_SHARD_BYTES:
|
|
out_file = output_dir / "model.safetensors"
|
|
save_file(all_hf_tensors, out_file, metadata=metadata if metadata else None)
|
|
log.info("Saved single shard: %s", out_file)
|
|
else:
|
|
# Split into shards
|
|
shards_list: list[dict[str, torch.Tensor]] = [{}]
|
|
current_shard_size = 0
|
|
for name, tensor in all_hf_tensors.items():
|
|
t_size = tensor.numel() * tensor.element_size()
|
|
if current_shard_size + t_size > MAX_SHARD_BYTES and shards_list[-1]:
|
|
shards_list.append({})
|
|
current_shard_size = 0
|
|
shards_list[-1][name] = tensor
|
|
current_shard_size += t_size
|
|
|
|
num_shards = len(shards_list)
|
|
weight_map: dict[str, str] = {}
|
|
for idx, shard_dict in enumerate(shards_list, start=1):
|
|
shard_filename = f"model-{idx:05d}-of-{num_shards:05d}.safetensors"
|
|
shard_path = output_dir / shard_filename
|
|
save_file(shard_dict, shard_path, metadata=metadata if metadata else None)
|
|
for k in shard_dict:
|
|
weight_map[k] = shard_filename
|
|
log.info("Saved shard %d/%d: %s", idx, num_shards, shard_path)
|
|
|
|
index_payload = {
|
|
"metadata": {"total_size": total_bytes},
|
|
"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 other tensors)",
|
|
output_dir,
|
|
quant_count,
|
|
passthrough_count,
|
|
)
|
|
|
|
|
|
def self_test(args) -> None:
|
|
"""Self-test NVFP4 unpacking and roundtripping."""
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
from nvfp4gguf import FP4_TABLE
|
|
import gguf
|
|
|
|
log.info("Running gguf_to_nvfp4_hf 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)
|
|
|
|
out_features = shape[0]
|
|
n_blocks = shape[1] // 16
|
|
n_super = shape[1] // 64
|
|
|
|
w_b = packed_w.reshape(out_features, n_blocks, 8)
|
|
vals = torch.stack([w_b & 0x0F, w_b >> 4], dim=-1).reshape(out_features, n_blocks, 16)
|
|
d_ue = scale_1.view(torch.uint8).reshape(out_features, n_blocks) & 0x7F
|
|
qs = (vals[:, :, :8] | (vals[:, :, 8:] << 4)).to(torch.uint8)
|
|
d_grouped = d_ue.reshape(out_features, n_super, 4)
|
|
qs_grouped = qs.reshape(out_features, n_super, 32)
|
|
raw = torch.cat([d_grouped, qs_grouped], dim=-1).reshape(out_features, n_super * 36).numpy()
|
|
|
|
w_rec, s1_rec, s2_rec = unpack_gguf_nvfp4_to_hf(raw, per_tensor)
|
|
assert torch.equal(packed_w, w_rec), "GGUF unpack weights mismatch!"
|
|
assert torch.equal(scale_1.view(torch.uint8) & 0x7F, s1_rec.view(torch.uint8) & 0x7F), "GGUF unpack scales mismatch!"
|
|
log.info("Shape %s GGUF-to-HF unpack roundtrip ok", shape)
|
|
|
|
log.info("self-test passed successfully.")
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Convert and quantize a GGUF model to NVFP4 Hugging Face format.",
|
|
)
|
|
parser.add_argument("gguf", nargs="?", help="Input GGUF file")
|
|
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; write float HF model")
|
|
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("--hf-template", type=Path, default=None,
|
|
help="Reference Hugging Face model directory to copy config and tokenizer files from")
|
|
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 tensor 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.gguf:
|
|
parser.error("GGUF model file 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_gguf_to_hf(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|