diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..0556fe2dcc3 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/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) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f51eea17b1a..5c08e0d1d15 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/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 diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 0d18530b902..5d454c74b2f 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/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 diff --git a/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py b/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py index ded38fa9dc2..a2f0ce28295 100644 --- a/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py +++ b/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: