"""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: ./-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()