Skip to content

Fix NVFP4 multi-GPU export device handling #2197

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
realAsma wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Fix NVFP4 multi-GPU export device handling #2197

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f6cd3fe
fix(export): preserve NVFP4 input device during export
realAsma Aug 14, 2026
25819df
fix(export): scope quantized weight export to its device
realAsma Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion modelopt/torch/export/unified_export_hf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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(
Expand Down
123 changes: 63 additions & 60 deletions modelopt/torch/quantization/qtensor/nvfp4_tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):

@claude claude Bot Aug 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] Format-asymmetric coverage: this guard fixes NVFP4QTensor.quantize only, while the sibling real-quantize paths keep the original exposure.

Within this file the only current-device-sensitive allocation is torch.ops.trtllm.fp4_quantize (_cast_fp4 already keys its lookup tables off weight.device via get_e2m1_bounds, and every other op derives its device from input). The same class of exposure exists for MXFP8QTensor / FP8QTensor / INT4QTensor when reached through TensorQuantizer.forward → _real_quantize, which — unlike _fake_quantize (tensor_quantizer.py:1205, already wrapped in same_device_as) — has no device guard:

            if self.fake_quant:
                with same_device_as(inputs):
                    outputs = self._fake_quantize(inputs)
            elif not self._dequantize:
                outputs = self._real_quantize(inputs)   # <- unguarded

Adding with same_device_as(inputs): around self._real_quantize(inputs) would make the compress/mtq.compress path symmetric with fake-quant for every format in one place, rather than per-QTensor-class. Non-blocking, and it doesn't replace this hunk — to_quantized_weight() calls NVFP4QTensor.quantize directly, which is what the new test_nvfp4_export_uses_input_device exercises.

Minor: wrapping the whole body re-indents ~60 lines, which makes the actual one-line intent hard to see in the diff. An early with same_device_as(input): scoped to just the trtllm branch (or an ExitStack) would keep the diff surgical, per the "prefer simple, surgical changes" guidance in CONTRIBUTING.md.

# 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."""
Expand Down
28 changes: 28 additions & 0 deletions tests/gpu/torch/export/test_export_weight_gpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
31 changes: 31 additions & 0 deletions tests/gpu/torch/quantization/test_qtensor_cuda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down
39 changes: 39 additions & 0 deletions tests/unit/torch/export/test_export_weight.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
# limitations under the License.


from contextlib import contextmanager

import pytest
import torch
import torch.nn as nn
Expand Down Expand Up @@ -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
Comment on lines +107 to +126

@claude claude Bot Aug 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] This test asserts an implementation detail rather than behavior: it replaces same_device_as with a stub that does nothing and then checks the stub was entered. It can only ever verify "_export_quantized_weight calls a module-level name same_device_as with getattr(module, weight_name)" — it cannot detect a wrong device, and it breaks on any refactor that keeps the behavior (e.g. inlining torch.cuda.device(...), or moving the guard into the impl body).

The behavioral coverage already exists in tests/gpu/torch/export/test_export_weight_gpu.py::test_export_nvfp4_modules_uses_each_weight_device and test_qtensor_cuda.py::test_nvfp4_export_uses_input_device. Consider dropping this one, or — if you want a CPU-runnable guard — assert something observable instead, e.g. that all registered scale buffers land on the weight's device.



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")
Comment on lines +129 to +141

@claude claude Bot Aug 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT TestCoverage] This test never reaches the line the PR changes, so it passes identically with and without the fix.

partial_fp8_config gives linears[1] a per-tensor (4,3) weight quantizer, so get_quantization_format() returns QUANTIZATION_FP8 and _export_quantized_weight_impl takes the if quantization_format == QUANTIZATION_FP8: branch (unified_export_hf.py:659). That branch registers input_scale from hasattr(input_quantizer, "_amax") and never evaluates the enabled-state check at all. The repr(...) → is_enabled hunk lives in the else: branch (unified_export_hf.py:709-712), which only runs for non-FP8 formats.

Why it matters: the PR's Testing section cites this test as the evidence for the is_enabled change, but extra_repr was never going to be called on this path — reverting the source hunk leaves the test green, so a future regression back to a repr-based check would not be caught.

Fix: drive a format that lands in the else: branch. partial_w4a8_config on linears[2] (already imported in this file, QUANTIZATION_W4A8_AWQ, input quantizer enabled with an amax) does hit the changed condition:

def test_export_quantized_weight_does_not_repr_input_quantizer(monkeypatch):
    model = ToyModel(dims=[32, 256, 256, 32])
    mtq.quantize(model, partial_w4a8_config, lambda x: x(torch.randn(1, 4, 32)))
    input_quantizer = model.linears[2].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[2], torch.float32, "weight")
    assert hasattr(model.linears[2], "input_scale")



class QuantMoELinear(nn.Module):
def __init__(self):
super().__init__()
Expand Down
Loading