Fix NVFP4 multi-GPU export device handling

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.
 
@@ -697,7 +709,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(

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."""

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

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",

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,43 @@ 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)))
+    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__()