From f6cd3fe38fe81c1d1b793890483e967437839e9a Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 14 Aug 2026 19:30:19 +0000 Subject: [PATCH 1/2] fix(export): preserve NVFP4 input device during export Signed-off-by: realAsma --- modelopt/torch/export/unified_export_hf.py | 2 +- .../quantization/qtensor/nvfp4_tensor.py | 123 +++++++++--------- .../torch/quantization/test_qtensor_cuda.py | 31 +++++ tests/unit/torch/export/test_export_weight.py | 15 +++ 4 files changed, 110 insertions(+), 61 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 2a605ed6d9e..72aa2cad490 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -697,7 +697,7 @@ def _export_quantized_weight( if ( input_quantizer is not None - and "disabled" not in repr(input_quantizer) + and input_quantizer.is_enabled and input_quantizer.amax is not None ): sub_module.register_buffer( diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 0d18530b902..fcb5026131d 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -17,6 +17,8 @@ import torch +from modelopt.torch.utils import same_device_as + from ..backends.utils import fp4_compatible from ..qtensor.base_qtensor import BaseQuantizedTensor from ..utils import reduce_amax, reduce_block_amax, reduce_block_padding @@ -275,71 +277,72 @@ def quantize( input_shape = input.shape input_dtype = input.dtype - # pad the input if needed - input = reduce_block_padding(input, block_sizes={-1: block_size}) - - if weights_scaling_factor_2 is None: - weights_scaling_factor_2 = cls.get_weights_scaling_factor_2(input) - - # try call trtllm fp4 quantization if possible - if ( - fp4_compatible() - and weights_scaling_factor is None - and try_tensorrt - and block_size == 16 - and input.is_cuda - and input.dtype in [torch.half, torch.bfloat16] - ): - try: - import tensorrt_llm # noqa: F401 - - # Make sure this utils is available for dequantize - from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( - cutlass_fp4_scale_to_modelopt_fp4_scale, # noqa: F401 + with same_device_as(input): + # pad the input if needed + input = reduce_block_padding(input, block_sizes={-1: block_size}) + + if weights_scaling_factor_2 is None: + weights_scaling_factor_2 = cls.get_weights_scaling_factor_2(input) + + # try call trtllm fp4 quantization if possible + if ( + fp4_compatible() + and weights_scaling_factor is None + and try_tensorrt + and block_size == 16 + and input.is_cuda + and input.dtype in [torch.half, torch.bfloat16] + ): + try: + import tensorrt_llm # noqa: F401 + + # Make sure this utils is available for dequantize + from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( + cutlass_fp4_scale_to_modelopt_fp4_scale, # noqa: F401 + ) + + packed_weight, weights_scaling_factor = torch.ops.trtllm.fp4_quantize( + input, 1.0 / weights_scaling_factor_2, block_size, False + ) + # weights_scaling_factor is ready for nvfp4_gemm to use; + # however, it is different from the non trtllm version, so when dequantize, + # it will be converted. + return ( + cls(input_shape, input_dtype, packed_weight), + weights_scaling_factor, + weights_scaling_factor_2, + ) + except ImportError: + pass + + if weights_scaling_factor is None: + weights_scaling_factor, _ = cls.get_weights_scaling_factor( + input, block_size, weights_scaling_factor_2 ) - packed_weight, weights_scaling_factor = torch.ops.trtllm.fp4_quantize( - input, 1.0 / weights_scaling_factor_2, block_size, False - ) - # weights_scaling_factor is ready for nvfp4_gemm to use; - # however, it is different from the non trtllm version, so when dequantize, - # it will be converted. - return ( - cls(input_shape, input_dtype, packed_weight), - weights_scaling_factor, - weights_scaling_factor_2, - ) - except ImportError: - pass + # Reshape the weight and scale factors + original_shape = input.shape + input = input.view((*tuple(input.shape[:-1]), -1, block_size)) - if weights_scaling_factor is None: - weights_scaling_factor, _ = cls.get_weights_scaling_factor( - input, block_size, weights_scaling_factor_2 + # Scale weights + scaled_weight = input / ( + (weights_scaling_factor.to(torch.float32) * weights_scaling_factor_2).unsqueeze(-1) ) - # Reshape the weight and scale factors - original_shape = input.shape - input = input.view((*tuple(input.shape[:-1]), -1, block_size)) - - # Scale weights - scaled_weight = input / ( - (weights_scaling_factor.to(torch.float32) * weights_scaling_factor_2).unsqueeze(-1) - ) - - # Reshape weights to original - scaled_weight = scaled_weight.view(original_shape) - - if keep_high_precision: - return scaled_weight - # Cast weights to fp4 - q_weight = cls._cast_fp4(scaled_weight) - # Pack weights - packed_weight = (q_weight[..., 1::2] << 4) | q_weight[..., 0::2] - return ( - cls(input_shape, input_dtype, packed_weight), - weights_scaling_factor, - weights_scaling_factor_2, - ) + # Reshape weights to original + scaled_weight = scaled_weight.view(original_shape) + + if keep_high_precision: + return scaled_weight + # Cast weights to fp4 + q_weight = cls._cast_fp4(scaled_weight) + # Pack weights + packed_weight = (q_weight[..., 1::2] << 4) | q_weight[..., 0::2] + return ( + cls(input_shape, input_dtype, packed_weight), + weights_scaling_factor, + weights_scaling_factor_2, + ) def dequantize(self, dtype: torch.dtype = None, fast=False, **kwarg): """Dequantze NVFP4 packed tensor to a target dtype.""" diff --git a/tests/gpu/torch/quantization/test_qtensor_cuda.py b/tests/gpu/torch/quantization/test_qtensor_cuda.py index cfdf38864fb..52213787d7f 100644 --- a/tests/gpu/torch/quantization/test_qtensor_cuda.py +++ b/tests/gpu/torch/quantization/test_qtensor_cuda.py @@ -21,6 +21,8 @@ import torch from _test_utils.torch.misc import set_seed +from modelopt.torch.export.model_config import QUANTIZATION_NVFP4 +from modelopt.torch.export.quant_utils import to_quantized_weight from modelopt.torch.quantization.backends.utils import fp4_compatible from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import TensorQuantizer @@ -397,6 +399,35 @@ def _unpack_tensor(x): # Compare with input tensor assert torch.allclose(deq_x, x, rtol=2e-1, atol=2e-1) + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Test requires two CUDA devices") + def test_nvfp4_export_uses_input_device(self): + with torch.cuda.device(1): + test_input = torch.randn((8, 32), dtype=torch.bfloat16, device="cuda:1") + double_scale = NVFP4QTensor.get_weights_scaling_factor_2(test_input) + scale, _ = NVFP4QTensor.get_weights_scaling_factor( + test_input, block_size=16, weights_scaling_factor_2=double_scale + ) + + with torch.cuda.device(0): + assert torch.cuda.current_device() == 0 + + packed_weight = to_quantized_weight( + test_input, + scale, + QUANTIZATION_NVFP4, + weights_scaling_factor2=double_scale, + block_size=16, + ) + + assert packed_weight.device == test_input.device + assert packed_weight.shape == (8, 16) + assert packed_weight.dtype == torch.uint8 + assert scale.device == test_input.device + assert double_scale.device == test_input.device + assert torch.cuda.current_device() == 0 + torch.cuda.synchronize(test_input.device) + assert torch.cuda.current_device() == 0 + @pytest.mark.parametrize("device", ["cuda"]) @pytest.mark.parametrize( "test_input", diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 6fc17d982e8..feb02e25651 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -102,6 +102,21 @@ def test_export_per_block_quantized_weight(): assert not hasattr(model.linears[2], quantizer_attrs.output_scale) +def test_export_quantized_weight_does_not_repr_input_quantizer(monkeypatch): + model = ToyModel(dims=[32, 256, 32]) + mtq.quantize(model, partial_fp8_config, lambda x: x(torch.randn(1, 4, 32))) + input_quantizer = model.linears[1].input_quantizer + + monkeypatch.setattr( + input_quantizer, + "extra_repr", + lambda: pytest.fail("export should inspect is_enabled without formatting the quantizer"), + ) + + _export_quantized_weight(model.linears[1], torch.float32, "weight") + assert hasattr(model.linears[1], "input_scale") + + class QuantMoELinear(nn.Module): def __init__(self): super().__init__() From 25819dfd4fc43408a651e0c4585463b6d597f7fc Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 14 Aug 2026 21:42:05 +0000 Subject: [PATCH 2/2] fix(export): scope quantized weight export to its device Signed-off-by: realAsma --- modelopt/torch/export/unified_export_hf.py | 12 ++++++++ .../torch/export/test_export_weight_gpu.py | 28 +++++++++++++++++++ tests/unit/torch/export/test_export_weight.py | 24 ++++++++++++++++ 3 files changed, 64 insertions(+) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 72aa2cad490..2bd3937347e 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -61,6 +61,7 @@ from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload +from modelopt.torch.utils import same_device_as from modelopt.torch.utils.dataset_utils import _disable_use_cache from modelopt.torch.utils.distributed import is_fsdp2_model @@ -571,6 +572,17 @@ def _export_quantized_weight( dtype: torch.dtype, weight_name: str = "weight", _tied_cache: dict[int, nn.Module] | None = None, +): + """Export one quantized weight while its device is current.""" + with same_device_as(getattr(sub_module, weight_name)): + return _export_quantized_weight_impl(sub_module, dtype, weight_name, _tied_cache) + + +def _export_quantized_weight_impl( + sub_module: nn.Module, + dtype: torch.dtype, + weight_name: str = "weight", + _tied_cache: dict[int, nn.Module] | None = None, ): """For the given weight attr of the sub_module, export the quantization info of it. diff --git a/tests/gpu/torch/export/test_export_weight_gpu.py b/tests/gpu/torch/export/test_export_weight_gpu.py index 9db2b51114b..4a91bba3c3e 100644 --- a/tests/gpu/torch/export/test_export_weight_gpu.py +++ b/tests/gpu/torch/export/test_export_weight_gpu.py @@ -16,6 +16,7 @@ import copy import math +import pytest import torch import torch.nn as nn from _test_utils.torch.export.utils import ToyModel, partial_w4a8_config @@ -125,6 +126,33 @@ def test_export_per_block_quantized_weight(): assert not hasattr(model.linears[2], quantizer_attrs.output_scale) +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Test requires two CUDA devices") +def test_export_nvfp4_modules_uses_each_weight_device(): + modules = [] + for device_idx in range(2): + device = torch.device("cuda", device_idx) + with torch.cuda.device(device): + module = nn.Linear(32, 32, bias=False, device=device, dtype=torch.bfloat16) + mtq.quantize( + module, + mtq.NVFP4_DEFAULT_CFG, + lambda m: m(torch.randn(2, 32, device=device, dtype=torch.bfloat16)), + ) + modules.append(module) + + for expected_device, module in enumerate(modules): + wrong_device = 1 - expected_device + with torch.cuda.device(wrong_device): + _export_quantized_weight(module, torch.bfloat16) + + assert torch.cuda.current_device() == wrong_device + assert module.weight.device.index == expected_device + assert module.weight_scale.device.index == expected_device + assert module.weight_scale_2.device.index == expected_device + assert module.input_scale.device.index == expected_device + torch.cuda.synchronize(expected_device) + + def test_export_compressed_nvfp4_weight(): """``mtq.compress`` (used by ``hf_ptq --low_memory_mode``) leaves the weight as packed NVFP4 nibbles, so per-block scales cannot be recomputed from it. The export must reuse the scales diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index feb02e25651..8cb01d907c2 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -14,6 +14,8 @@ # limitations under the License. +from contextlib import contextmanager + import pytest import torch import torch.nn as nn @@ -102,6 +104,28 @@ def test_export_per_block_quantized_weight(): assert not hasattr(model.linears[2], quantizer_attrs.output_scale) +def test_export_quantized_weight_uses_weight_device_context(monkeypatch): + model = ToyModel(dims=[32, 32]) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(1, 4, 32))) + linear = model.linears + entered = False + + @contextmanager + def record_device_context(weight): + nonlocal entered + assert weight is linear.weight + entered = True + yield + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf.same_device_as", record_device_context + ) + + _export_quantized_weight(linear, torch.float32) + + assert entered + + def test_export_quantized_weight_does_not_repr_input_quantizer(monkeypatch): model = ToyModel(dims=[32, 256, 32]) mtq.quantize(model, partial_fp8_config, lambda x: x(torch.randn(1, 4, 32)))