fix/improve hf to hf conversions

gguf_to_nvfp4_hf.py

@@ -0,0 +1,761 @@
+"""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()