Add LSQ NVFP4 export support

modelopt/torch/export/quant_utils.py

@@ -309,6 +309,77 @@ def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") ->
     return get_scaling_factor(weight_quantizer)
 
 
+def _to_export_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    if hasattr(tensor, "to_local"):
+        tensor = tensor.to_local()
+    return tensor.detach().float()
+
+
+def _lsq_amax_to_scale(
+    amax: torch.Tensor, max_bound: float, min_value: float | torch.Tensor
+) -> torch.Tensor:
+    scale = _to_export_tensor(amax) / max_bound
+    return torch.where(scale <= min_value, min_value, scale)
+
+
+def _lsq_scale_factors(
+    scale: torch.Tensor,
+    quantizer: TensorQuantizer,
+    quantize_scale: bool,
+    expected_shape: torch.Size,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    scale = scale.view(expected_shape)
+    if not quantize_scale:
+        return scale, torch.tensor(1.0, device=scale.device)
+
+    per_tensor_scale = _to_export_tensor(quantizer.global_amax) / quantizer._quant_max_bound
+    scale_2 = per_tensor_scale / 448.0
+    scaled = (scale * 448.0 / per_tensor_scale.view(-1)).clamp(min=2**-9, max=448.0)
+    return scaled.to(torch.float8_e4m3fn), scale_2
+
+
+def get_lsq_weight_scaling_factors(
+    weight_quantizer: TensorQuantizer, weight: torch.Tensor, block_size: int
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return LSQ pre/post scale factors for NVFP4 export.
+
+    The pre factors are used only to pack the FP4 weight codes. The post factors
+    are serialized as the exported dequantization scales.
+    """
+    assert getattr(weight_quantizer, "_lsq", False), "Expected an LSQ quantizer."
+    assert weight.shape[-1] % block_size == 0, (
+        "Weight shape is not divisible for block size for block quantization."
+    )
+
+    expected_shape = torch.Size((*weight.shape[:-1], weight.shape[-1] // block_size))
+    quantize_scales = getattr(weight_quantizer, "_quantize_scales", False)
+    max_bound = float(weight_quantizer._quant_max_bound)
+    per_tensor_scale = (
+        _to_export_tensor(weight_quantizer.global_amax) / max_bound if quantize_scales else None
+    )
+
+    post_min = 2**-9 * per_tensor_scale.view(-1) if per_tensor_scale is not None else 1e-8
+    pre_min = (
+        2**-9 * per_tensor_scale.view(-1)
+        if per_tensor_scale is not None and getattr(weight_quantizer, "_quantize_pre_scale", True)
+        else 1e-8
+    )
+
+    post_scale = _lsq_amax_to_scale(weight_quantizer.amax_post, max_bound, post_min)
+    pre_scale = _lsq_amax_to_scale(weight_quantizer.amax_pre, max_bound, pre_min)
+
+    pre_scale, pre_scale_2 = _lsq_scale_factors(
+        pre_scale,
+        weight_quantizer,
+        quantize_scales and getattr(weight_quantizer, "_quantize_pre_scale", True),
+        expected_shape,
+    )
+    post_scale, post_scale_2 = _lsq_scale_factors(
+        post_scale, weight_quantizer, quantize_scales, expected_shape
+    )
+    return pre_scale, pre_scale_2, post_scale, post_scale_2
+
+
 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

@@ -98,6 +98,7 @@
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
+    get_lsq_weight_scaling_factors,
     get_quant_config,
     get_quantization_format,
     get_weight_block_size,
@@ -583,6 +584,12 @@ def _export_quantized_weight(
     output_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr(
         sub_module, quantizer_attrs.output_quantizer, None
     )
+    is_lsq_nvfp4 = getattr(weight_quantizer, "_lsq", False) and quantization_format in [
+        QUANTIZATION_NVFP4,
+        QUANTIZATION_NVFP4_AWQ,
+        QUANTIZATION_NVFP4_SVDQUANT,
+        QUANTIZATION_W4A16_NVFP4,
+    ]
 
     if quantization_format == QUANTIZATION_FP8:
         # Convert amax to float32
@@ -633,6 +640,8 @@ def _export_quantized_weight(
             sub_module.register_buffer(quantizer_attrs.weight_scale, e8m0_scale)
             if hasattr(weight_quantizer, "_scale") and weight_quantizer._scale is not None:
                 del weight_quantizer._scale
+        elif is_lsq_nvfp4:
+            pass
         else:
             sub_module.register_buffer(
                 quantizer_attrs.weight_scale, get_weight_scaling_factor(sub_module, weight_name)
@@ -650,14 +659,18 @@ def _export_quantized_weight(
                 ).squeeze(),
             )
 
-    if quantization_format in [
-        QUANTIZATION_NVFP4_AWQ,
-        QUANTIZATION_NVFP4_SVDQUANT,
-        QUANTIZATION_NVFP4,
-        QUANTIZATION_W4A16_NVFP4,
-        QUANTIZATION_W4A8_AWQ,
-        QUANTIZATION_W4A8_NVFP4_FP8,
-    ]:
+    if (
+        quantization_format
+        in [
+            QUANTIZATION_NVFP4_AWQ,
+            QUANTIZATION_NVFP4_SVDQUANT,
+            QUANTIZATION_NVFP4,
+            QUANTIZATION_W4A16_NVFP4,
+            QUANTIZATION_W4A8_AWQ,
+            QUANTIZATION_W4A8_NVFP4_FP8,
+        ]
+        and not is_lsq_nvfp4
+    ):
         # Register weight_scale_2
         sub_module.register_buffer(
             quantizer_attrs.weight_scale_2,
@@ -692,7 +705,11 @@ def _export_quantized_weight(
             weight, is_bmm_expert_weight=is_bmm_expert_weight
         )
 
-        if NVFP4QTensor._is_static_quantizer(weight_quantizer):
+        if is_lsq_nvfp4:
+            weight_scale, weight_scale_2, export_weight_scale, export_weight_scale_2 = (
+                get_lsq_weight_scaling_factors(weight_quantizer, weight, block_size)
+            )
+        elif NVFP4QTensor._is_static_quantizer(weight_quantizer):
             weight_scale = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(
                 weight_quantizer,
                 weight,
@@ -712,10 +729,18 @@ def _export_quantized_weight(
             weight_scale_2,
             block_size,
         )
+        if is_lsq_nvfp4:
+            weight_scale = export_weight_scale
+            weight_scale_2 = export_weight_scale_2
 
         quantized_weight, weight_scale = maybe_transpose_expert_weight_dimensions(
             quantized_weight, weight_scale, is_bmm_expert_weight=is_bmm_expert_weight
         )
+        if is_lsq_nvfp4:
+            assert weight_scale is not None
+            assert weight_scale_2 is not None
+            sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale)
+            sub_module.register_buffer(quantizer_attrs.weight_scale_2, weight_scale_2.squeeze())
     elif quantization_format == QUANTIZATION_FP8_PC_PT and is_bmm_expert_weight:
         # For FP8_PC_PT with BMM-style experts, transpose only the weight (not weight_scale)
         weight, _ = maybe_transpose_expert_weight_dimensions(

tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py

@@ -20,7 +20,12 @@
 import pytest
 import torch
 
-from modelopt.torch.export.quant_utils import QUANTIZATION_NVFP4, to_quantized_weight
+from modelopt.torch.export.quant_utils import (
+    QUANTIZATION_NVFP4,
+    get_lsq_weight_scaling_factors,
+    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 +76,90 @@ def _export_round_trip(
     return weight_scale, weight_scale_2, dequant
 
 
+def _make_lsq_quantizer(
+    learnable_amax: list[str],
+    tied_amax: bool,
+    quantize_pre_scale: bool,
+) -> NVFP4StaticQuantizer:
+    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(
+        quantize_scales=True,
+        learnable_amax=learnable_amax,
+        tied_amax=tied_amax,
+        quantize_pre_scale=quantize_pre_scale,
+    )
+    with torch.no_grad():
+        q.amax_post.copy_(torch.full_like(q.amax_post, 3.0))
+        if not tied_amax:
+            q.amax_pre.copy_(torch.full_like(q.amax_pre, 1.5))
+    return q
+
+
+@pytest.mark.parametrize(
+    ("learnable_amax", "tied_amax", "quantize_pre_scale"),
+    [
+        (["post"], False, True),
+        (["pre"], False, True),
+        (["pre", "post"], False, True),
+        (["pre", "post"], True, True),
+        ([], False, True),
+        (["post"], False, False),
+    ],
+)
+def test_lsq_nvfp4_export_uses_pre_scale_for_packing_and_post_scale_for_dequant(
+    learnable_amax, tied_amax, quantize_pre_scale
+):
+    weight = torch.linspace(-4.0, 4.0, 4 * 32, dtype=torch.float32).view(4, 32)
+    q = _make_lsq_quantizer(learnable_amax, tied_amax, quantize_pre_scale)
+
+    pre_scale, pre_scale_2, post_scale, post_scale_2 = get_lsq_weight_scaling_factors(
+        q, weight, BLOCK_SIZE
+    )
+    packed = to_quantized_weight(weight, pre_scale, QUANTIZATION_NVFP4, pre_scale_2, BLOCK_SIZE)
+
+    dequant = NVFP4QTensor(weight.shape, weight.dtype, packed).dequantize(
+        scale=post_scale,
+        double_scale=post_scale_2,
+        block_sizes={-1: BLOCK_SIZE},
+        dtype=weight.dtype,
+    )
+
+    assert packed.dtype == torch.uint8
+    assert post_scale.shape == (4, 2)
+    assert torch.isfinite(dequant).all()
+    if not tied_amax:
+        wrong_packed = to_quantized_weight(
+            weight, post_scale, QUANTIZATION_NVFP4, post_scale_2, BLOCK_SIZE
+        )
+        assert not torch.equal(packed, wrong_packed)
+
+
+def test_lsq_nvfp4_export_quantized_weight_registers_post_scale_and_packs_with_pre_scale():
+    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"], False, quantize_pre_scale=False)
+
+    pre_scale, pre_scale_2, post_scale, post_scale_2 = get_lsq_weight_scaling_factors(
+        module.weight_quantizer, weight, BLOCK_SIZE
+    )
+    expected_packed = to_quantized_weight(
+        weight, pre_scale, QUANTIZATION_NVFP4, pre_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, post_scale)
+    torch.testing.assert_close(module.weight_scale_2, post_scale_2.squeeze())
+
+
 def _layer1_routed_expert_like(
     out_dim: int, in_dim: int, *, n_outliers: int, seed: int
 ) -> torch.Tensor: