Fix NVFP4 multi-GPU export device handling - #2197
Conversation
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2197 +/- ##
==========================================
- Coverage 78.97% 78.96% -0.01%
==========================================
Files 522 522
Lines 60606 60612 +6
==========================================
Hits 47862 47862
- Misses 12744 12750 +6
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| 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") |
There was a problem hiding this comment.
[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")
| 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 |
There was a problem hiding this comment.
[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.
| # 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): |
There was a problem hiding this comment.
[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) # <- unguardedAdding 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.
There was a problem hiding this comment.
Claude review — Fix NVFP4 multi-GPU export device handling
Scope: full review (trigger comment carried no scoping instructions). All 5 changed files opened: modelopt/torch/export/unified_export_hf.py, modelopt/torch/quantization/qtensor/nvfp4_tensor.py, and the 3 test files. Traced same_device_as and is_enabled to their definitions and all call sites.
Findings
| Severity | Count |
|---|---|
| CRITICAL | 0 |
| IMPORTANT | 1 |
| SUGGESTION | 2 |
IMPORTANT — test_export_quantized_weight_does_not_repr_input_quantizer does not exercise the line it is meant to cover. partial_fp8_config makes linears[1] per-tensor FP8, so export takes the quantization_format == QUANTIZATION_FP8 branch, which registers input_scale off hasattr(input_quantizer, "_amax") and never evaluates the enabled-state check. The repr(...) → is_enabled hunk lives in the else: branch. The test is green with or without the source change, so the repr-based check could regress silently. Switching to partial_w4a8_config on linears[2] reaches the changed condition.
SUGGESTIONs (non-blocking): the CPU unit test for the device context is tautological — it stubs out same_device_as and asserts the stub ran, which can't detect a wrong device; and the nvfp4_tensor.py guard leaves _real_quantize unprotected for MXFP8/FP8/INT4 while re-indenting ~60 lines for a one-line intent.
What I verified as correct
repr→is_enabledis semantically equivalent and is the established convention.extra_repremits the literal"disabled"iffself._disabled, andis_enabledis exactlynot self._disabled(tensor_quantizer.py:519-521); nothing else inextra_reprcan contain that substring.quant_utils.py:524already uses the identicalinput_quantizer is not None and input_quantizer.is_enabledform for the same purpose in the TRT-LLM path, so this alignsunified_export_hfwith it. Container quantizers delegateis_enabledto member 0 per_QuantizerContainerBase._delegated_properties, andenable/disablebroadcast to all members, so mixed-state containers aren't reachable in practice.- The
_export_quantized_weightwrapper is safe for every entry path.same_device_asshort-circuits tonullcontext()for non-CUDA tensors, so meta-offloaded weights still reach the existingweight.is_metaRuntimeErrorunchanged, and CPU-only unit tests are unaffected. The unconditionalgetattr(sub_module, weight_name)in the wrapper — now ahead of theQUANTIZATION_NONEearly return — is fine: all three call sites inhf_export_handlers.py(142/158/186) gate onget_quantization_format(...) != QUANTIZATION_NONEfirst, and themoe_utils.py:268wrapper module sets.weightimmediately before the call.QTensorWrapperis annn.Parameter, so.is_cuda/.deviceresolve for the compressed path, and FSDP2 callers wrap infsdp2_aware_weight_updatebefore the getattr. - The root cause is correctly identified.
torch.cuda.device()only redirects allocations that key off the current device — inside this file that istorch.ops.trtllm.fp4_quantize, which matches the reported failure._cast_fp4was already device-correct via thee2m1_bounds/e2m1_valuesper-device caches. The guard is thread-local (CUDA current device is per-thread), so the layer-by-layer streaming writer is safe. - No public API, mode registration, config schema, or
modelopt_statesurface is touched; no CHANGELOG entry needed per the repo's own criteria. No CPU-GPU sync added.
Risk
Low. The production change is a narrow, correctly-targeted device-scoping fix with no behavioral change on single-GPU exports, plus a strictly-better replacement of a string-matching hack. The one blocking item is test-only — the multi-GPU behavior is genuinely covered by the two new GPU tests; it is the is_enabled half of the change that currently has no effective regression test.
1 participant
What does this PR do?
Type of change: Bug fix
NVFP4 weight export could allocate intermediate tensors on the process's current CUDA device instead of the input tensor's device. This caused export failures when a model shard lived on a different GPU. This change scopes NVFP4 quantization to the input device while restoring the caller's current device afterward.
The export path also now uses the quantizer's public
is_enabledstate instead of inferring that state from its formatted representation.Usage
N/A — this fixes existing multi-GPU export behavior without changing the public API.
Testing
pre-commit run --files modelopt/torch/export/unified_export_hf.py modelopt/torch/quantization/qtensor/nvfp4_tensor.py tests/gpu/torch/quantization/test_qtensor_cuda.py tests/unit/torch/export/test_export_weight.pypytest tests/unit/torch/export/test_export_weight.py::test_export_quantized_weight_does_not_repr_input_quantizer -q -xCUDA_VISIBLE_DEVICES=1,2 pytest tests/gpu/torch/quantization/test_qtensor_cuda.py::TestQTensor::test_nvfp4_export_uses_input_device -q -xBefore your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
No new dependency or public API is introduced.