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 cee64c22c05..615638bd2a5 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/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( 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: