From c3132bc7180a191243f169ec9fc35827e6013cdd Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 22 Jul 2026 23:08:44 +0000 Subject: [PATCH 1/2] Add LSQ NVFP4 export support Signed-off-by: realAsma --- modelopt/torch/export/quant_utils.py | 71 +++++++++++++++ modelopt/torch/export/unified_export_hf.py | 43 +++++++-- .../test_nvfp4_static_export_cpu.py | 91 ++++++++++++++++++- 3 files changed, 195 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..5de27d91dab 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/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) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f51eea17b1a..c788e516c81 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -100,6 +100,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, @@ -585,6 +586,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 @@ -635,6 +642,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) @@ -652,14 +661,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, @@ -694,7 +707,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, @@ -714,10 +731,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( 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..5b6e128f4d7 100644 --- a/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py +++ b/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: From 540d2fc38cc07a9d7861b027a831dbbb8307eaab Mon Sep 17 00:00:00 2001 From: realAsma Date: Thu, 30 Jul 2026 02:46:37 +0000 Subject: [PATCH 2/2] Restrict NVFP4 export to tied LSQ quantizers Signed-off-by: realAsma --- modelopt/torch/export/quant_utils.py | 83 +++------------ modelopt/torch/export/unified_export_hf.py | 45 ++------ .../quantization/qtensor/nvfp4_tensor.py | 6 +- .../test_nvfp4_static_export_cpu.py | 100 +++++++++--------- 4 files changed, 83 insertions(+), 151 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 5de27d91dab..0556fe2dcc3 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -309,75 +309,24 @@ 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 assert_lsq_export_supported(weight_quantizer: TensorQuantizer) -> None: + """Raise if an LSQ quantizer uses a configuration that cannot be exported. - -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. + 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. """ - 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 + 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: diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index c788e516c81..5c08e0d1d15 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -97,10 +97,10 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import ( + assert_lsq_export_supported, 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, @@ -586,12 +586,7 @@ 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, - ] + assert_lsq_export_supported(weight_quantizer) if quantization_format == QUANTIZATION_FP8: # Convert amax to float32 @@ -642,8 +637,6 @@ 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) @@ -661,18 +654,14 @@ 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, - ] - and not is_lsq_nvfp4 - ): + if quantization_format in [ + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_NVFP4_SVDQUANT, + QUANTIZATION_NVFP4, + QUANTIZATION_W4A16_NVFP4, + QUANTIZATION_W4A8_AWQ, + QUANTIZATION_W4A8_NVFP4_FP8, + ]: # Register weight_scale_2 sub_module.register_buffer( quantizer_attrs.weight_scale_2, @@ -707,11 +696,7 @@ def _export_quantized_weight( weight, is_bmm_expert_weight=is_bmm_expert_weight ) - 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): + if NVFP4QTensor._is_static_quantizer(weight_quantizer): weight_scale = NVFP4QTensor.get_weights_scaling_factor_from_quantizer( weight_quantizer, weight, @@ -731,18 +716,10 @@ 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( 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 5b6e128f4d7..a2f0ce28295 100644 --- a/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py +++ b/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py @@ -20,11 +20,7 @@ import pytest import torch -from modelopt.torch.export.quant_utils import ( - QUANTIZATION_NVFP4, - get_lsq_weight_scaling_factors, - to_quantized_weight, -) +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 @@ -76,11 +72,8 @@ 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: +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)}, @@ -89,75 +82,84 @@ def _make_lsq_quantizer( 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, + **{ + "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)) - 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 -): +@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, tied_amax, quantize_pre_scale) + 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 + ) - pre_scale, pre_scale_2, post_scale, post_scale_2 = get_lsq_weight_scaling_factors( - q, weight, BLOCK_SIZE + # 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, pre_scale, QUANTIZATION_NVFP4, pre_scale_2, BLOCK_SIZE) + packed = to_quantized_weight( + weight, weight_scale, QUANTIZATION_NVFP4, weight_scale_2, BLOCK_SIZE + ) dequant = NVFP4QTensor(weight.shape, weight.dtype, packed).dequantize( - scale=post_scale, - double_scale=post_scale_2, + scale=weight_scale, + double_scale=weight_scale_2, block_sizes={-1: BLOCK_SIZE}, dtype=weight.dtype, ) - assert packed.dtype == torch.uint8 - assert post_scale.shape == (4, 2) + assert weight_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(): +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"], False, quantize_pre_scale=False) + module.weight_quantizer = _make_lsq_quantizer(["post"]) - pre_scale, pre_scale_2, post_scale, post_scale_2 = get_lsq_weight_scaling_factors( - module.weight_quantizer, weight, BLOCK_SIZE + 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, pre_scale, QUANTIZATION_NVFP4, pre_scale_2, BLOCK_SIZE + 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, post_scale) - torch.testing.assert_close(module.weight_scale_2, post_scale_2.squeeze()) + 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(