support ModelOpt NVFP4 dequantization

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":

tests/test_model_dequant_fp8.py

@@ -8,17 +8,25 @@
 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,
     convert_compressed_pack_file,
     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",