From 89ea8751e8faf156b02fa2761398857d5fec6d2e Mon Sep 17 00:00:00 2001 From: ZX-ModelCloud Date: Tue, 14 Jul 2026 11:53:01 +0800 Subject: [PATCH] support ModelOpt NVFP4 dequantization Signed-off-by: ZX-ModelCloud --- gptqmodel/utils/model_dequant.py | 65 +++++++++++++++++++-- tests/test_model_dequant_fp8.py | 96 ++++++++++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 5 deletions(-) diff --git a/gptqmodel/utils/model_dequant.py b/gptqmodel/utils/model_dequant.py index 23945b170..5fb9a2731 100644 --- a/gptqmodel/utils/model_dequant.py +++ b/gptqmodel/utils/model_dequant.py @@ -423,6 +423,7 @@ def _is_quant_auxiliary_tensor_key(key: str) -> bool: "_scale_inv", ".scale", ".weight_scale", + ".weight_scale_2", ".weight_absmax", ".weight_quant_map", ".weight_nested_absmax", @@ -606,18 +607,21 @@ def infer_block_shape(weight_shape: Tuple[int, int], scale_tensor: torch.Tensor) raise ValueError("unable to infer block size from 1D scale tensor") - raise ValueError("unsupported scale tensor rank for block size inference") + raise ValueError("unsupported scale tensor rank for block size inference") def detect_format(model_path: Path, config: dict) -> str: quant_cfg = config.get("quantization_config", {}) or {} method = (quant_cfg.get("method") or quant_cfg.get("quant_method") or "").lower() format_name = (quant_cfg.get("format") or "").lower() + variant_name = (quant_cfg.get("variant") or "").lower() + quant_algo = str(quant_cfg.get("quant_algo") or quant_cfg.get("algorithm") or "").lower() files, _ = list_safetensor_files(model_path) if not files: raise FileNotFoundError("No .safetensors files found in model directory") + # legacy/local checkpoints that do not persist a useful quantization_config. with safe_open(model_path / files[0], framework="pt", device="cpu") as reader: keys = list(reader.keys()) # Prefer dtype-based detection @@ -660,6 +664,22 @@ def detect_format(model_path: Path, config: dict) -> str: ) return "gptq" if has_g else "awq" + if format_name == "nvfp4" or ( + method == "modelopt" + and ( + format_name in {"nvfp4", "fp4"} + or variant_name == "nvfp4" + or quant_algo == "nvfp4" + ) + ): + LOG.debug( + "Detected NVFP4 format via config method=%s format=%s variant=%s quant_algo=%s", + method, + format_name, + variant_name, + quant_algo, + ) + return "nvfp4" if format_name in _FLOAT8_FORMAT_NAMES: LOG.debug("Detected FP8 format via config format=%s", format_name) return "fp8" @@ -1079,9 +1099,11 @@ def convert_nvfp4_shard( reader, target_dtype: torch.dtype, *, + tensor_lookup: Optional[_ShardTensorLookup] = None, ignored_layers: Iterable[str] = (), ) -> Dict[str, torch.Tensor]: tensors: Dict[str, torch.Tensor] = {} + reader_keys = set(reader.keys()) for key in reader.keys(): tensor = reader.get_tensor(key) ignored_tensor = _handle_ignored_tensor(key, tensor, target_dtype, ignored_layers) @@ -1093,9 +1115,41 @@ def convert_nvfp4_shard( if key.endswith(".weight") and tensor.dtype in _NVFP4_STORAGE_DTYPES: scale_key = key + "_scale" - if scale_key not in reader.keys(): + # ModelOpt checkpoints may shard weight and scale tensors separately. + if tensor_lookup is None: + has_scale = scale_key in reader_keys + else: + has_scale = tensor_lookup.has_tensor( + scale_key, + local_reader=reader, + local_keys=reader_keys, + ) + if not has_scale: raise KeyError(f"Missing scale tensor for {key}") - scale = reader.get_tensor(scale_key) + if tensor_lookup is None: + scale = reader.get_tensor(scale_key) + else: + scale = tensor_lookup.get_tensor( + scale_key, + local_reader=reader, + local_keys=reader_keys, + ) + + scale_2_key = key + "_scale_2" + # ModelOpt stores an optional global scale in weight_scale_2. + if tensor_lookup is not None and tensor_lookup.has_tensor( + scale_2_key, + local_reader=reader, + local_keys=reader_keys, + ): + scale_2 = tensor_lookup.get_tensor( + scale_2_key, + local_reader=reader, + local_keys=reader_keys, + ) + scale = scale.to(torch.float32) * scale_2.to(torch.float32) + elif scale_2_key in reader_keys: + scale = scale.to(torch.float32) * reader.get_tensor(scale_2_key).to(torch.float32) LOG.debug("Using scale tensor '%s' for NVFP4 weight '%s'", scale_key, key) deq = dequantize_f4_e2m1( tensor, @@ -1104,7 +1158,7 @@ def convert_nvfp4_shard( target_dtype=target_dtype, ) tensors[key] = finalize_for_save(deq, target_dtype) - elif key.endswith(".weight_scale"): + elif key.endswith((".weight_scale", ".weight_scale_2")): LOG.debug("Dropping auxiliary NVFP4 tensor '%s' after dequantization", key) continue else: @@ -1483,7 +1537,7 @@ def dequantize_model( device=open_device, weight_map=index.get("weight_map", {}) if isinstance(index, dict) else None, ) - if fmt == "fp8" + if fmt in {"fp8", "nvfp4"} else None ) @@ -1515,6 +1569,7 @@ def dequantize_model( tensors = convert_nvfp4_shard( reader, target_dtype, + tensor_lookup=tensor_lookup, ignored_layers=ignored_layers, ) elif fmt == "awq": diff --git a/tests/test_model_dequant_fp8.py b/tests/test_model_dequant_fp8.py index daf5868de..5653f9815 100644 --- a/tests/test_model_dequant_fp8.py +++ b/tests/test_model_dequant_fp8.py @@ -8,6 +8,7 @@ from safetensors import safe_open from safetensors.torch import save_file +from gptqmodel.quantization.dtype import dequantize_f4_e2m1 from gptqmodel.utils.model_dequant import ( convert_awq_file, convert_bitsandbytes_shard, @@ -15,10 +16,17 @@ convert_gptq_file, convert_nvfp4_shard, dequantize_model, + detect_format, finalize_for_save, ) +try: + from torchao.prototype.mx_formats.nvfp4_tensor import nvfp4_quantize +except Exception: + nvfp4_quantize = None + + def _write_index(model_dir, shard_name: str, keys: list[str]) -> None: weight_map = dict.fromkeys(keys, shard_name) (model_dir / "model.safetensors.index.json").write_text( @@ -280,6 +288,94 @@ def test_dequantize_model_fp8_honors_ignored_layers(tmp_path): torch.testing.assert_close(ignored_out, ignored_weight) +def test_detect_format_modelopt_nvfp4_uses_quant_algo_config(tmp_path): + model_dir = tmp_path / "modelopt_nvfp4_config_detect" + model_dir.mkdir() + + config = { + "architectures": ["TestModel"], + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + }, + } + (model_dir / "config.json").write_text(json.dumps(config), encoding="utf-8") + save_file({"dense.weight": torch.ones(2, 2, dtype=torch.bfloat16)}, str(model_dir / "model.safetensors")) + + assert detect_format(model_dir, config) == "nvfp4" + + +@pytest.mark.skipif(nvfp4_quantize is None, reason="torchao NVFP4 support required") +def test_dequantize_model_modelopt_nvfp4_resolves_scales_from_other_shard(tmp_path): + model_dir = tmp_path / "modelopt_nvfp4_cross_shard" + output_dir = tmp_path / "modelopt_nvfp4_cross_shard_out" + model_dir.mkdir() + + config = { + "architectures": ["TestModel"], + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + }, + } + (model_dir / "config.json").write_text(json.dumps(config), encoding="utf-8") + + torch.manual_seed(0) + dense = torch.randn(4, 16, dtype=torch.float32) + scales, packed = nvfp4_quantize(dense, block_size=16) + global_scale = torch.tensor(2.0, dtype=torch.float32) + bias = torch.randn(4, dtype=torch.float32) + + weight_shard = "model-00001-of-00002.safetensors" + scale_shard = "model-00002-of-00002.safetensors" + save_file({"linear.weight": packed.cpu(), "linear.bias": bias}, str(model_dir / weight_shard)) + save_file( + { + "linear.weight_scale": scales.cpu(), + "linear.weight_scale_2": global_scale, + }, + str(model_dir / scale_shard), + ) + (model_dir / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "linear.weight": weight_shard, + "linear.bias": weight_shard, + "linear.weight_scale": scale_shard, + "linear.weight_scale_2": scale_shard, + } + } + ), + encoding="utf-8", + ) + + dequantize_model(model_dir, output_dir, target_dtype=torch.bfloat16, device="cpu") + + with safe_open(output_dir / weight_shard, framework="pt", device="cpu") as reader: + assert set(reader.keys()) == {"linear.weight", "linear.bias"} + weight_out = reader.get_tensor("linear.weight") + bias_out = reader.get_tensor("linear.bias") + + expected_scale = scales.cpu().to(torch.float32) * global_scale + expected = dequantize_f4_e2m1( + packed.cpu(), + scale=expected_scale, + axis=None, + target_dtype=torch.bfloat16, + ) + assert weight_out.dtype is torch.bfloat16 + torch.testing.assert_close(weight_out, expected, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(bias_out, bias.to(torch.bfloat16)) + assert not (output_dir / scale_shard).exists() + + output_index = json.loads((output_dir / "model.safetensors.index.json").read_text(encoding="utf-8")) + assert output_index["weight_map"] == { + "linear.weight": weight_shard, + "linear.bias": weight_shard, + } + + @pytest.mark.skipif( not hasattr(torch, "float8_e8m0fnu"), reason="float8_e8m0fnu dtype not available",