support static NVFP4 HF export (#858)

## What does this PR do?

**Type of change:** ? new feature

**Overview:** ?
Supports export `NVFP4StaticQantizer` in unified huggingface checkpoint,
as a deployment path for PTQ algorithms such as MSE

## Usage
<!-- You can potentially add a usage example below. -->

```python
# checkpoint generation
python examples/llm_ptq/hf_ptq.py --pyt_ckpt_path Qwen/Qwen3-8B  --qformat nvfp4_mse --export_path test-Qwen3-8B-Instruct-MSE-FP8-sweep-FP4 --kv_cache_qformat none --trust_remote_code 

```

## Testing
Tested generated Qwen3 8B checkpoint with trtllm serve and nv_eval
example in `Model-Optimizer-Internal/examples/nv_eval`.

NV eval results:
```
| Groups |Version|Filter|n-shot|  Metric   |   |Value |   |Stderr|
|--------|-------|------|------|-----------|---|-----:|---|-----:|
|mmlu_str|       |none  |      |exact_match|↑  |0.7186|±  |0.0036|
```


## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added support for static NVFP4 quantizers that utilize pre-computed
calibration scales.
* Introduced new NVFP4 W4A4 quantization configuration with optional FP8
scale sweep.

* **Performance Improvements**
* Static quantizers now skip unnecessary dynamic scaling factor
recalculation.
* Unified quantization handling for improved consistency and efficiency.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: Fridah-nv

CHANGELOG.rst

@@ -15,6 +15,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 - Add ``--moe_calib_experts_ratio`` flag in ``hf_ptq.py`` to specify the ratio of experts to calibrate during forward pass to improve expert coverage during calibration. Default to all the experts.
 - Add sparse attention optimization for transformer models (``modelopt.torch.sparsity.attention_sparsity``). This reduces computational cost by skipping attention computation. Supports calibration for threshold selection on HuggingFace models. See `examples/llm_sparsity/attention_sparsity/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_sparsity/attention_sparsity>`_ for usage.
 - Add support for rotating the input before quantization for RHT.
+- Add support for advanced weight scale search for NVFP4 quantization and its export path.
 
 0.42 (2026-02-xx)
 ^^^^^^^^^^^^^^^^^
@@ -36,6 +37,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 - Add LTX-2 and Wan2.2 (T2V) support in the diffusers quantization workflow.
 - 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.
 - Add PTQ support for Nemotron Parse.
 - Add distillation support for LTX-2. See `examples/diffusers/distillation/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/distillation>`_ for more details.
 

examples/llm_ptq/hf_ptq.py

@@ -81,6 +81,7 @@
     "w4a8_awq": mtq.W4A8_AWQ_BETA_CFG,
     "nvfp4": mtq.NVFP4_DEFAULT_CFG,
     "nvfp4_awq": mtq.NVFP4_AWQ_LITE_CFG,
+    "nvfp4_mse": mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG,
     "fp8_pb_wo": mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
     "fp8_pc_pt": mtq.FP8_PER_CHANNEL_PER_TOKEN_CFG,
     "w4a8_nvfp4_fp8": mtq.W4A8_NVFP4_FP8_CFG,
@@ -890,6 +891,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,11 +41,12 @@
 from modelopt.torch.quantization.utils import (
     QuantizerAttrNames,
     quantizer_attr_names,
+    reduce_block_amax,
     weight_attr_names,
 )
 from modelopt.torch.utils import clear_cuda_cache
 
-from ..quantization.nn import SequentialQuantizer, TensorQuantizer
+from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer
 from .model_config import (
     KV_CACHE_FP8,
     KV_CACHE_INT8,
@@ -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. "
@@ -299,7 +353,7 @@ def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") ->
         return get_scaling_factor(weight_quantizer[0])
 
     quantization_format = get_quantization_format(module)
-    # If NVFP4, we need to return quantized per_block scaling factors
+
     if quantization_format in [
         QUANTIZATION_NVFP4,
         QUANTIZATION_NVFP4_AWQ,
@@ -318,9 +372,10 @@ def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") ->
             weight_scaling_factor_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(
                 weight_quantizer
             )
-        return NVFP4QTensor.get_weights_scaling_factor(
+        # Unified method handles both static and dynamic quantizers
+        return NVFP4QTensor.get_weights_scaling_factor_from_quantizer(
+            weight_quantizer,
             weight,
-            weight_quantizer.block_sizes[-1],
             weight_scaling_factor_2.to(weight.device),
         )[0]
 
@@ -343,27 +398,24 @@ def get_weight_scaling_factor_2(module: nn.Module, weight_name: str = "weight")
 
     quantization_format = get_quantization_format(module)
 
-    # Calibrate weight quantizer if amax is not set for all NVFP4 variants
     if quantization_format in [
         QUANTIZATION_NVFP4,
         QUANTIZATION_NVFP4_AWQ,
         QUANTIZATION_NVFP4_SVDQUANT,
         QUANTIZATION_W4A8_NVFP4_FP8,
     ]:
+        # Calibrate weight quantizer if amax is not set
         weight = getattr(module, weight_name)
         module_name = f"{type(module).__name__}.{weight_name}"
         _ensure_weight_quantizer_calibrated(weight_quantizer, weight, module_name)
 
-    if quantization_format in [
-        QUANTIZATION_NVFP4,
-        QUANTIZATION_NVFP4_AWQ,
-        QUANTIZATION_NVFP4_SVDQUANT,
-    ]:
-        return NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
-    elif quantization_format == QUANTIZATION_W4A8_NVFP4_FP8:
-        # weight_scaling_factor_2 for w4a8 needs to be amax/448, so that the wsf is in range 448/6.
-        # This is because the kernel dequantizes weight to fp8, which is in range 448.
-        return weight_quantizer._amax.float() / 448.0
+        if quantization_format == QUANTIZATION_W4A8_NVFP4_FP8:
+            # weight_scaling_factor_2 for w4a8 needs to be amax/448, so that the wsf is in range 448/6.
+            # This is because the kernel dequantizes weight to fp8, which is in range 448.
+            return weight_quantizer._amax.float() / 448.0
+        else:
+            # Unified method handles both static and dynamic quantizers
+            return NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
 
     # SequentialQuantizer is required
     if not isinstance(weight_quantizer, SequentialQuantizer) or not weight_quantizer[-1].is_enabled:
@@ -735,7 +787,7 @@ def process_layer_quant_config(layer_config_dict):
             layer_config = {"quant_algo": "W8A16"}
         elif v == "int8_sq":
             layer_config = {"quant_algo": "W8A8_SQ_PER_CHANNEL"}
-        elif v == "nvfp4":
+        elif v in ["nvfp4", "nvfp4_static"]:
             layer_config = {
                 "quant_algo": "NVFP4",
                 "group_size": block_size_value,
@@ -1339,6 +1391,18 @@ def preprocess_linear_fusion(modules: list[torch.nn.Module], resmooth_only=False
                 for module in modules:
                     module.weight_quantizer[-1].amax = weight_amax
 
+        # Handle NVFP4StaticQuantizer: unify global_amax for fused layers
+        elif isinstance(modules[0].weight_quantizer, NVFP4StaticQuantizer):
+            global_amax_list = [
+                m.weight_quantizer.global_amax
+                for m in modules
+                if m.weight_quantizer.global_amax is not None
+            ]
+            if global_amax_list:
+                unified_global_amax = torch.max(torch.stack(global_amax_list))
+                for module in modules:
+                    module.weight_quantizer.global_amax = unified_global_amax
+
         elif (
             modules[0].weight_quantizer.is_enabled
             and modules[0].weight_quantizer.amax is not None
@@ -1423,6 +1487,22 @@ def get_quant_config(
             if block_size == 0:
                 block_size = get_weight_block_size(module)
 
+            # Static NVFP4 uses pre-computed per-block scales from MSE calibration
+            if quantization_format == QUANTIZATION_NVFP4:
+                weight_quantizer = getattr(module, "weight_quantizer", None)
+                if weight_quantizer is None:
+                    # Try to get from first weight attribute
+                    for wn in weight_names:
+                        weight_quantizer = getattr(
+                            module, quantizer_attr_names(wn).weight_quantizer, None
+                        )
+                        if weight_quantizer is not None:
+                            break
+                if weight_quantizer is not None:
+                    is_static = isinstance(weight_quantizer, NVFP4StaticQuantizer)
+                    if is_static:
+                        quantization_format = "nvfp4_static"
+
             # Construct per layer config dictionary
             layer_config_dict[name + ".quantization"] = quantization_format
             layer_config_dict[name + ".awq_block_size"] = block_size

modelopt/torch/export/unified_export_hf.py

@@ -50,7 +50,11 @@
 from torch.distributed.fsdp import FSDPModule
 
 from modelopt.torch.quantization import set_quantizer_by_cfg_context
-from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
+from modelopt.torch.quantization.nn import (
+    NVFP4StaticQuantizer,
+    SequentialQuantizer,
+    TensorQuantizer,
+)
 from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor
 from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names
 
@@ -496,6 +500,11 @@ def _export_quantized_weight(
         expert_type in type(sub_module).__name__
         for expert_type in ["Llama4TextExperts", "GptOssExperts"]
     )
+    if is_bmm_expert_weight and isinstance(weight_quantizer, NVFP4StaticQuantizer):
+        raise ValueError(
+            "NVFP4StaticQuantizer with BMM-style expert weights (e.g. Llama4TextExperts, "
+            "GptOssExperts) is not yet supported."
+        )
 
     if quantization_format in [
         QUANTIZATION_NVFP4,
@@ -507,6 +516,7 @@ def _export_quantized_weight(
         weight, _ = maybe_transpose_expert_weight_dimensions(
             weight, is_bmm_expert_weight=is_bmm_expert_weight
         )
+
         weight_scale = NVFP4QTensor.get_weights_scaling_factor(
             weight,
             block_size=block_size,

modelopt/torch/quantization/config.py

@@ -419,6 +419,28 @@
     "algorithm": "max",
 }
 
+NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG = {
+    "quant_cfg": {
+        "*weight_quantizer": {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        },
+        "*input_quantizer": {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        },
+        **_default_disabled_quantizer_cfg,
+    },
+    "algorithm": {
+        "method": "mse",
+        "fp8_scale_sweep": True,
+    },
+}
+
 NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG = {
     "quant_cfg": {
         "*weight_quantizer": {
@@ -751,6 +773,7 @@
     "MAMBA_MOE_NVFP4_AGGRESSIVE_CFG",
     "MAMBA_MOE_FP8_CONSERVATIVE_CFG",
     "MAMBA_MOE_FP8_AGGRESSIVE_CFG",
+    "NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG",
 }
 
 BiasType = Literal["static", "dynamic"]