add unit tests, update _ensure_weight_quantizer_calibrated to handle NVFP4StaticQuantizer, update changelog

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

CHANGELOG.rst

@@ -19,6 +19,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 - Add support for context parallelism in Eagle speculative decoding for huggingface and megatron core models.
 - Add PTQ support for GLM-4.7, including loading MTP layer weights from a separate ``mtp.safetensors`` file and export as-is.
 - Add support for image-text data calibration in PTQ for Nemotron VL models.
+- Add support for advanced weight scale search for NVFP4 quantization and its export path.
 
 0.41 (2026-01-19)
 ^^^^^^^^^^^^^^^^^

examples/llm_ptq/hf_ptq.py

@@ -868,6 +868,7 @@ def quantize_main(
                 "fp8",
                 "nvfp4",
                 "nvfp4_awq",
+                "nvfp4_mse",
                 "w4a8_awq",
                 "fp8_pb_wo",
                 "w4a8_mxfp4_fp8",

modelopt/torch/export/quant_utils.py

@@ -41,6 +41,7 @@
 from modelopt.torch.quantization.utils import (
     QuantizerAttrNames,
     quantizer_attr_names,
+    reduce_block_amax,
     weight_attr_names,
 )
 from modelopt.torch.utils import clear_cuda_cache
@@ -238,6 +239,36 @@ def get_scaling_factor(quantizer: TensorQuantizer) -> torch.Tensor:
     return scaling_factor
 
 
+def _get_nvfp4_block_size(
+    weight_quantizer: NVFP4StaticQuantizer, weight: torch.Tensor, module_name: str = ""
+) -> int:
+    """Return block size for NVFP4 from quantizer's block_sizes; raise if missing."""
+    prefix = f"NVFP4StaticQuantizer{f' for {module_name}' if module_name else ''}"
+    block_sizes = weight_quantizer.block_sizes
+    if block_sizes is None:
+        raise ValueError(f"{prefix} has no block_sizes; cannot compute per-block amax from weight.")
+    block_size = block_sizes.get(-1) or block_sizes.get(weight.dim() - 1)
+    if block_size is None:
+        raise ValueError(
+            f"{prefix} block_sizes has no -1 or last-dim key; cannot compute per-block amax."
+        )
+    return block_size
+
+
+def _set_amax_from_tensor(weight_quantizer: TensorQuantizer, tensor: torch.Tensor) -> None:
+    """Set quantizer _amax buffer from tensor; copy in-place if same shape, else replace buffer."""
+    if (
+        hasattr(weight_quantizer, "_amax")
+        and weight_quantizer._amax is not None
+        and weight_quantizer._amax.shape == tensor.shape
+    ):
+        weight_quantizer._amax.data.copy_(tensor.to(weight_quantizer._amax.device))
+    else:
+        if hasattr(weight_quantizer, "_amax"):
+            delattr(weight_quantizer, "_amax")
+        weight_quantizer.register_buffer("_amax", tensor.clone().detach())
+
+
 def _ensure_weight_quantizer_calibrated(
     weight_quantizer: TensorQuantizer, weight: torch.Tensor, module_name: str = ""
 ) -> None:
@@ -246,11 +277,34 @@ def _ensure_weight_quantizer_calibrated(
     This is a lazy calibration pattern used during export when weight quantizers
     may not have been calibrated during the main calibration phase.
 
+    For NVFP4StaticQuantizer, _amax is per-block amax and _global_amax is the max over
+    blocks; both are computed from the weight when missing.
+
     Args:
         weight_quantizer: The weight quantizer to calibrate
         weight: The weight tensor to use for calibration
         module_name: Optional module name for better warning messages
     """
+    if isinstance(weight_quantizer, NVFP4StaticQuantizer):
+        need_per_block = not hasattr(weight_quantizer, "_amax") or weight_quantizer._amax is None
+        need_global = (
+            not hasattr(weight_quantizer, "_global_amax") or weight_quantizer.global_amax is None
+        )
+        if not (need_per_block or need_global):
+            return
+        block_size = _get_nvfp4_block_size(weight_quantizer, weight, module_name)
+        warn(
+            f"NVFP4StaticQuantizer{f' for {module_name}' if module_name else ''} was not fully calibrated. "
+            f"Computing per-block amax and global_amax from weights. This may occur if: "
+            f"some experts were not activated during calibration (expected for MoE models), try increasing --calib_size"
+        )
+        per_block_amax = reduce_block_amax(weight, block_sizes={-1: block_size})
+        if need_per_block:
+            _set_amax_from_tensor(weight_quantizer, per_block_amax.to(weight.device))
+        if need_global:
+            weight_quantizer.global_amax = per_block_amax.max()
+        return
+
     if not hasattr(weight_quantizer, "_amax") or weight_quantizer._amax is None:
         warn(
             f"Weight quantizer{f' for {module_name}' if module_name else ''} was not calibrated. "

tests/_test_utils/torch/export/utils.py

@@ -119,6 +119,18 @@ def forward(self, x):
             "axis": None,
             "enable": True,
         },
+        "*.2.weight_quantizer": {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        },
+        "*.2.input_quantizer": {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        },
         "default": {"enable": False},
     },
     "algorithm": "max",

tests/gpu/torch/export/test_export.py

@@ -415,7 +415,7 @@ def test_get_scaling_factor(
         (
             partial_nvfp4_config,
             {
-                "exclude_modules": ["linears.0", "linears.2"],
+                "exclude_modules": ["linears.0"],
                 "group_size": 16,
                 "kv_cache_quant_algo": None,
                 "quant_algo": "NVFP4",

tests/gpu/torch/export/test_export_weight_gpu.py

@@ -17,16 +17,17 @@
 
 import torch
 import torch.nn as nn
-from _test_utils.torch.export.utils import ToyModel, partial_w4a8_config
+from _test_utils.torch.export.utils import ToyModel, partial_nvfp4_config, partial_w4a8_config
 from torch.nn import functional as F
 from torch.nn import init
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+from modelopt.torch.quantization.nn import NVFP4StaticQuantizer
 from modelopt.torch.quantization.nn.modules.quant_module import QuantModule, QuantModuleRegistry
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
 from modelopt.torch.quantization.tensor_quant import QUANT_DESC_8BIT_PER_TENSOR
-from modelopt.torch.quantization.utils import quantizer_attr_names
+from modelopt.torch.quantization.utils import quantizer_attr_names, reduce_block_amax
 
 
 class ToyLinear(nn.Module):
@@ -121,3 +122,61 @@ def test_export_per_block_quantized_weight():
     assert hasattr(model.linears[2], quantizer_attrs.output_quantizer)
     assert not getattr(model.linears[2], quantizer_attrs.output_quantizer).is_enabled
     assert not hasattr(model.linears[2], quantizer_attrs.output_scale)
+
+
+def test_export_nvfp4_static_weight_dynamic_vs_static_match():
+    """Dynamic vs static NVFP4 export: same weight and scales after export even when amaxs are
+    cleared on one layer (lazy calibration via _ensure_weight_quantizer_calibrated fills them from weights).
+    """
+    device = "cuda"
+    dims = [32, 32, 32, 32]
+    block_size = 16
+    calib_input = torch.randn(1, 4, 32, device=device)
+    nvfp4_layer_indices = [1, 2]  # layers with NVFP4 enabled in partial_nvfp4_config
+
+    torch.manual_seed(42)
+    model_dynamic = ToyModel(dims=dims).to(device)
+    mtq.quantize(model_dynamic, partial_nvfp4_config, lambda x: x(calib_input))
+
+    torch.manual_seed(42)
+    model_static = ToyModel(dims=dims).to(device)
+    mtq.quantize(model_static, partial_nvfp4_config, lambda x: x(calib_input))
+
+    # Convert NVFP4 layers to NVFP4StaticQuantizer with per-block and global amax
+    for idx in nvfp4_layer_indices:
+        layer = model_static.linears[idx]
+        weight = layer.weight.data
+        per_block_amax = reduce_block_amax(weight, block_sizes={-1: block_size})
+        tq = layer.weight_quantizer
+        if hasattr(tq, "_amax"):
+            delattr(tq, "_amax")
+        tq.register_buffer("_amax", per_block_amax.to(weight.device).clone().detach())
+        NVFP4StaticQuantizer.from_tensor_quantizer(tq, global_amax=per_block_amax.max())
+
+    # Clear amaxs on layer 1 to exercise lazy calibration during export
+    for linear, is_static in [(model_dynamic.linears[1], False), (model_static.linears[1], True)]:
+        wq = linear.weight_quantizer
+        if hasattr(wq, "_amax"):
+            delattr(wq, "_amax")
+        if is_static and hasattr(wq, "_global_amax"):
+            delattr(wq, "_global_amax")
+
+    quantizer_attrs = quantizer_attr_names("weight")
+    for idx in nvfp4_layer_indices:
+        _export_quantized_weight(model_dynamic.linears[idx], torch.float32, "weight")
+        _export_quantized_weight(model_static.linears[idx], torch.float32, "weight")
+
+    for idx in nvfp4_layer_indices:
+        dyn_linear = model_dynamic.linears[idx]
+        sta_linear = model_static.linears[idx]
+        assert torch.equal(dyn_linear.weight, sta_linear.weight), (
+            f"Layer {idx}: exported NVFP4 weight should match (dynamic vs static)"
+        )
+        assert torch.allclose(
+            getattr(dyn_linear, quantizer_attrs.weight_scale).float(),
+            getattr(sta_linear, quantizer_attrs.weight_scale).float(),
+        ), f"Layer {idx}: weight_scale should match"
+        assert torch.allclose(
+            getattr(dyn_linear, quantizer_attrs.weight_scale_2).float(),
+            getattr(sta_linear, quantizer_attrs.weight_scale_2).float(),
+        ), f"Layer {idx}: weight_scale_2 should match"

tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py

@@ -35,6 +35,7 @@
     [
         ("fp8", "tiny_llama-fp8", True, False, True, True, False),
         ("nvfp4", "tiny_llama-nvfp4", True, False, True, True, False),
+        ("nvfp4_mse", "tiny_llama-nvfp4-mse", True, False, True, True, False),
         ("nvfp4_awq", "tiny_llama-nvfp4-awq", True, False, True, True, False),
         ("int4_awq", "tiny_llama-int4-awq", True, False, True, True, False),
         ("w4a8_awq", "tiny_llama-w4a8-awq", True, False, True, True, False),

tests/unit/torch/export/test_get_quantization.py

@@ -15,11 +15,22 @@
 
 import pytest
 import torch
-from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config
+from _test_utils.torch.export.utils import (
+    ToyModel,
+    partial_fp8_config,
+    partial_nvfp4_config,
+    partial_w4a8_config,
+)
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.layer_utils import get_quantization_format
-from modelopt.torch.export.model_config import QUANTIZATION_FP8, QUANTIZATION_W4A8_AWQ
+from modelopt.torch.export.model_config import (
+    QUANTIZATION_FP8,
+    QUANTIZATION_NVFP4,
+    QUANTIZATION_W4A8_AWQ,
+)
+from modelopt.torch.export.quant_utils import get_quant_config
+from modelopt.torch.quantization.nn import NVFP4StaticQuantizer
 
 
 @pytest.mark.parametrize(
@@ -30,3 +41,22 @@ def test_get_quantization_format(config, expected):
     model = ToyModel()
     mtq.quantize(model, config, lambda x: x(torch.randn(1, 4, 10)))
     assert get_quantization_format(model) == expected
+
+
+def test_nvfp4_static_quantizer_export():
+    """NVFP4StaticQuantizer: get_quantization_format returns NVFP4 and get_quant_config returns export config."""
+    model = ToyModel()
+    mtq.quantize(model, partial_nvfp4_config, lambda x: x(torch.randn(1, 4, 10)))
+
+    # Convert all weight quantizers to NVFP4StaticQuantizer
+    for module in model.modules():
+        tq = getattr(module, "weight_quantizer", None)
+        if tq is not None and hasattr(tq, "_amax") and not isinstance(tq, NVFP4StaticQuantizer):
+            global_amax = tq._amax.max() if tq._amax.dim() > 0 else tq._amax
+            NVFP4StaticQuantizer.from_tensor_quantizer(tq, global_amax=global_amax)
+
+    assert get_quantization_format(model) == QUANTIZATION_NVFP4
+
+    quant_config = get_quant_config(model)
+    assert quant_config["quantization"]["quant_algo"] == "NVFP4"
+    assert quant_config["quantization"]["group_size"] == 16