Add LSQ NVFP4 export support

modelopt/torch/export/quant_utils.py

@@ -309,6 +309,26 @@ def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") ->
     return get_scaling_factor(weight_quantizer)
 
 
+def assert_lsq_export_supported(weight_quantizer: TensorQuantizer) -> None:
+    """Raise if an LSQ quantizer uses a configuration that cannot be exported.
+
+    Only tied-amax LSQ with FP8-quantized pre and post block scales is exportable: it
+    then produces exactly the scales and packed weights of the static NVFP4 export path.
+    Any other configuration would serialize scales that do not match training.
+    """
+    if not getattr(weight_quantizer, "_lsq", False):
+        return
+    if (
+        not weight_quantizer._tied_amax
+        or not weight_quantizer._quantize_scales
+        or not weight_quantizer._quantize_pre_scale
+    ):
+        raise NotImplementedError(
+            "Dual LSQ export is not supported: only tied-amax LSQ (tied_amax=True, "
+            "quantize_scales=True, quantize_pre_scale=True) can be exported."
+        )
+
+
 def get_weight_scaling_factor_2(module: nn.Module, weight_name: str = "weight") -> torch.Tensor:
     """Returns the secondary weight scaling factor."""
     weight_quantizer = getattr(module, quantizer_attr_names(weight_name).weight_quantizer, None)

modelopt/torch/export/unified_export_hf.py

@@ -97,6 +97,7 @@
     revert_weight_conversion_quant_aware,
 )
 from .quant_utils import (
+    assert_lsq_export_supported,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -585,6 +586,7 @@ def _export_quantized_weight(
     output_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr(
         sub_module, quantizer_attrs.output_quantizer, None
     )
+    assert_lsq_export_supported(weight_quantizer)
 
     if quantization_format == QUANTIZATION_FP8:
         # Convert amax to float32

modelopt/torch/quantization/qtensor/nvfp4_tensor.py

@@ -139,7 +139,11 @@ def get_weights_scaling_factor_from_quantizer(
         if cls._is_static_quantizer(weight_quantizer):
             # Static path: use pre-computed per-block amax values from quantizer
             global_amax = cls._get_static_global_amax(weight_quantizer).float()
-            per_block_amax = weight_quantizer._amax.float()
+            per_block_amax = getattr(weight_quantizer, "_amax", None)
+            if per_block_amax is None:
+                # Tied LSQ deletes ``_amax`` and exposes ``_amax_post`` through ``amax``.
+                per_block_amax = weight_quantizer.amax
+            per_block_amax = per_block_amax.float()
 
             # Compute scales in float
             per_block_scale_max = global_amax / E2M1_MAX

tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py

@@ -21,6 +21,7 @@
 import torch
 
 from modelopt.torch.export.quant_utils import QUANTIZATION_NVFP4, to_quantized_weight
+from modelopt.torch.export.unified_export_hf import _export_quantized_weight
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.nn import NVFP4StaticQuantizer
 from modelopt.torch.quantization.qtensor import NVFP4QTensor
@@ -71,6 +72,96 @@ def _export_round_trip(
     return weight_scale, weight_scale_2, dequant
 
 
+def _make_lsq_quantizer(learnable_amax: list[str], **lsq_overrides) -> NVFP4StaticQuantizer:
+    """LSQ quantizer, tied-amax (the only exportable variant) unless overridden."""
+    cfg = QuantizerAttributeConfig(
+        num_bits=(2, 1),
+        block_sizes={-1: BLOCK_SIZE, "type": "static", "scale_bits": (4, 3)},
+    )
+    q = NVFP4StaticQuantizer(quant_attribute_cfg=cfg)
+    q.amax = torch.ones(8)
+    q.global_amax = torch.tensor(6.0)
+    q.enable_lsq(
+        learnable_amax=learnable_amax,
+        **{
+            "quantize_scales": True,
+            "tied_amax": True,
+            "quantize_pre_scale": True,
+            **lsq_overrides,
+        },
+    )
+    with torch.no_grad():
+        q.amax_post.copy_(torch.full_like(q.amax_post, 3.0))
+    return q
+
+
+@pytest.mark.parametrize("learnable_amax", [["post"], ["pre", "post"], []])
+def test_lsq_tied_nvfp4_export_matches_training_block_scale(learnable_amax):
+    weight = torch.linspace(-4.0, 4.0, 4 * 32, dtype=torch.float32).view(4, 32)
+    q = _make_lsq_quantizer(learnable_amax)
+
+    weight_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(q)
+    weight_scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(
+        q, weight, weight_scale_2
+    )
+
+    # The exported dequantization scale must reproduce the scale used in training.
+    expected = q._block_scale_from_amax(q.amax_post.detach(), True)
+    torch.testing.assert_close(
+        (weight_scale.float() * weight_scale_2.float()).reshape(-1), expected.reshape(-1)
+    )
+
+    packed = to_quantized_weight(
+        weight, weight_scale, QUANTIZATION_NVFP4, weight_scale_2, BLOCK_SIZE
+    )
+    dequant = NVFP4QTensor(weight.shape, weight.dtype, packed).dequantize(
+        scale=weight_scale,
+        double_scale=weight_scale_2,
+        block_sizes={-1: BLOCK_SIZE},
+        dtype=weight.dtype,
+    )
+    assert packed.dtype == torch.uint8
+    assert weight_scale.shape == (4, 2)
+    assert torch.isfinite(dequant).all()
+
+
+def test_lsq_tied_nvfp4_export_quantized_weight_registers_static_scales():
+    weight = torch.linspace(-4.0, 4.0, 4 * 32, dtype=torch.float32).view(4, 32)
+    module = torch.nn.Module()
+    module.weight = torch.nn.Parameter(weight.clone())
+    module.weight_quantizer = _make_lsq_quantizer(["post"])
+
+    expected_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(
+        module.weight_quantizer
+    )
+    expected_scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(
+        module.weight_quantizer, weight, expected_scale_2
+    )
+    expected_packed = to_quantized_weight(
+        weight, expected_scale, QUANTIZATION_NVFP4, expected_scale_2, BLOCK_SIZE
+    )
+
+    _export_quantized_weight(module, torch.float32)
+
+    torch.testing.assert_close(module.weight, expected_packed)
+    torch.testing.assert_close(module.weight_scale, expected_scale)
+    torch.testing.assert_close(module.weight_scale_2, expected_scale_2.squeeze())
+
+
+@pytest.mark.parametrize(
+    "lsq_overrides",
+    [{"tied_amax": False}, {"quantize_pre_scale": False}, {"quantize_scales": False}],
+)
+def test_lsq_nvfp4_export_rejects_unsupported_variants(lsq_overrides):
+    weight = torch.linspace(-4.0, 4.0, 4 * 32, dtype=torch.float32).view(4, 32)
+    module = torch.nn.Module()
+    module.weight = torch.nn.Parameter(weight.clone())
+    module.weight_quantizer = _make_lsq_quantizer(["post"], **lsq_overrides)
+
+    with pytest.raises(NotImplementedError, match="Dual LSQ"):
+        _export_quantized_weight(module, torch.float32)
+
+
 def _layer1_routed_expert_like(
     out_dim: int, in_dim: int, *, n_outliers: int, seed: int
 ) -> torch.Tensor: