[NVBug: 5987078] Fix unified HF export of compressed NVFP4 weights (--low_memory_mode) - #2038
[NVBug: 5987078] Fix unified HF export of compressed NVFP4 weights (--low_memory_mode)#2038cjluo-nv wants to merge 1 commit into
Conversation
📝 WalkthroughWalkthroughThe Hugging Face exporter now supports already-compressed NVFP4 weights. It reuses and rescales stored compression scales, preserves dequantization, removes internal quantizer buffers, and adds GPU regression coverage. ChangesCompressed NVFP4 export
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant QTensorWrapper
participant unified_export_hf
participant TensorRTLLM
participant Checkpoint
QTensorWrapper->>unified_export_hf: Provide compressed NVFP4 weights and stored scales
unified_export_hf->>TensorRTLLM: Convert swizzled block scales when required
unified_export_hf->>unified_export_hf: Rescale block scales for weight_scale_2
unified_export_hf->>Checkpoint: Write exported weights without internal scale buffers
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
modelopt/torch/export/unified_export_hf.py (1)
668-679: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSkip the temporary packed-weight scale calculation.
get_weight_scaling_factor(...)already runs before this code detectsQTensorWrapper. For a compressed NVFP4 weight, Lines 716-720 replace that result. Detect the compressed weight before generic scale registration and skip the packed-data reduction on this path. This removes an unnecessary full-weight pass during low-memory export.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 668 - 679, Move the QTensorWrapper detection and retrieval of _scale/_double_scale in the export flow before get_weight_scaling_factor(...) and generic scale registration. For compressed weights, reuse the quantizer’s stored scales and skip the packed-data reduction entirely; preserve the existing generic calculation for uncompressed weights.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 668-679: Move the QTensorWrapper detection and retrieval of
_scale/_double_scale in the export flow before get_weight_scaling_factor(...)
and generic scale registration. For compressed weights, reuse the quantizer’s
stored scales and skip the packed-data reduction entirely; preserve the existing
generic calculation for uncompressed weights.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5f099001-1f77-4423-a9f0-9d7b51416b4e
📒 Files selected for processing (3)
CHANGELOG.rstmodelopt/torch/export/unified_export_hf.pytests/gpu/torch/export/test_export_weight_gpu.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2038 +/- ##
==========================================
- Coverage 75.76% 69.65% -6.12%
==========================================
Files 518 518
Lines 58269 58287 +18
==========================================
- Hits 44148 40597 -3551
- Misses 14121 17690 +3569
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:
|
b96737c to
f749100
Compare
| ) | ||
|
|
||
| if NVFP4QTensor._is_static_quantizer(weight_quantizer): | ||
| if ( |
There was a problem hiding this comment.
for readability: should we name this branch condition as real_nvfp4_quantized or something similar?
cjluo-nv
changed the title
cjluo-nv
changed the title
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
The rescale approach is correct for the environment the author tested (no TensorRT-LLM in the process), and the new GPU test is a meaningful regression test. Three findings:
-
weight_quantizer._scaleis not always the modelopt per-block E4M3 layout.TensorQuantizer._real_quantize()callsNVFP4QTensor.quantize(..., try_tensorrt=True); whenfp4_compatible()(SM100+),block_size == 16, andtensorrt_llmis importable, it returns the cutlass-swizzled 1-D uint8 scale factor instead (seeNVFP4QTensor.dequantize, which explicitly detectsscale.dtype == torch.uint8 and scale.ndim == 1and converts viacutlass_fp4_scale_to_modelopt_fp4_scale). The new export branch has no such guard, so on Blackwell inside a TRT-LLM container--low_memory_modeNVFP4 export would write a 1-D scale built from raw uint8 byte values — the same class of bug this PR fixes, just in a different environment. The DGX Spark / vLLM validation wouldn't hit it, and the new test lives intests/gpu(nottests/gpu_trtllm), so CI won't either. Please either convert or raise loudly. -
Root cause of the
_double_scaleleak looks like a typo inRealQuantLinear.list_of_scale_tensors("double_scale"instead of"_double_scale").postprocess_state_dictalready stripsweight_quantizer._scalevia that list and would strip_double_scaletoo if the name matched. Worth confirming whether the one-word fix there is the better place (it also covers any other consumer of that list) rather than adding a second, export-path-only deletion mechanism. -
Minor: the new
.to(torch.float8_e4m3fn)skips the clamp that_cast_per_block_scale_to_fp8applies on every other NVFP4 path.
No licensing or design concerns; size and CHANGELOG are fine.
| # ``preprocess_linear_fusion`` unifies ``weight_scale_2`` over a fused group the ratio | ||
| # below is 1 only for the member owning the group max; the others take one extra E4M3 | ||
| # rounding (<= half-ULP, 6.25%). Avoiding that needs a shared scale at compress time. | ||
| weight_scale = ( |
There was a problem hiding this comment.
Bot comment.
_scale isn't guaranteed to be the modelopt per-block E4M3 tensor. TensorQuantizer._real_quantize() calls NVFP4QTensor.quantize(..., try_tensorrt=True), and when fp4_compatible() (device capability >= (10, 0)), block_size == 16, the weight is CUDA fp16/bf16, and tensorrt_llm imports, it returns the cutlass-format scale — a 1-D uint8 tensor in the 128x4 swizzled layout — and stores that in _scale. NVFP4QTensor.dequantize handles this case explicitly:
if kwarg["scale"].dtype == torch.uint8 and kwarg["scale"].ndim == 1:
kwarg["scale"] = cutlass_fp4_scale_to_modelopt_fp4_scale(kwarg["scale"], self.metadata["shape"][-2:])Here compressed_weight_scale.float() on that tensor yields raw byte values and a 1-D shape, so the exported weight_scale would again be wrong (wrong shape and meaningless values) — on Blackwell + TRT-LLM, which is a primary target for hf_ptq.py. The DGX Spark/vLLM validation and the new tests/gpu test both avoid this branch (no tensorrt_llm in-process), so neither catches it.
Suggest either converting via cutlass_fp4_scale_to_modelopt_fp4_scale before the rescale, or at minimum asserting compressed_weight_scale.dtype == torch.float8_e4m3fn and compressed_weight_scale.dim() == weight_scale_2.dim() + ... and raising a clear NotImplementedError, so the failure is loud rather than a silently malformed checkpoint.
| compressed_weight_scale.float() | ||
| * compressed_weight_scale_2.float().to(compressed_weight_scale.device) | ||
| / weight_scale_2.float().to(compressed_weight_scale.device) | ||
| ).to(torch.float8_e4m3fn) |
There was a problem hiding this comment.
Bot comment.
Every other NVFP4 path funnels the fp8 cast through _cast_per_block_scale_to_fp8 (in modelopt/torch/quantization/qtensor/nvfp4_tensor.py), which clamps to [2**-9, E4M3_MAX] specifically to avoid underflow→0 / overflow→NaN. Since weight_scale_2 here is the (possibly group-unified, i.e. >=) scale, the ratio is <= 1 and small block scales can underflow to 0 after the cast — a zero per-block scale silently zeroes that block at dequant. Reusing the existing helper (or at least clamping min=2**-9) would keep this consistent with the uncompressed path.
|
|
||
| # Internal to the quantizer: left registered they leak into the checkpoint as | ||
| # ``*.weight_quantizer._double_scale``, which downstream loaders reject with a KeyError. | ||
| if is_compressed_weight: |
There was a problem hiding this comment.
Bot comment.
postprocess_state_dict already has a mechanism for exactly this:
if any(key.endswith("weight_quantizer." + q_key) for q_key in RealQuantLinear.list_of_scale_tensors)and RealQuantLinear.list_of_scale_tensors = ["_scale", "double_scale", "_scale_zeros"] — note "double_scale" is missing its leading underscore, which is precisely why _scale never leaked but _double_scale did (matching the 560 stray *.weight_quantizer._double_scale keys and the absence of _scale keys in the bug report). Fixing that entry looks like the actual root cause and covers every consumer of that list, rather than adding a second deletion mechanism in this one export path. If you keep the delattr here (it does free memory earlier, which is a fair reason for --low_memory_mode), please still fix the list entry and say so in the CHANGELOG, otherwise the typo stays as a latent leak.
| weight_quantizer = getattr(compressed_module, quantizer_attrs.weight_quantizer) | ||
| assert getattr(weight_quantizer, "_scale", None) is None | ||
| assert getattr(weight_quantizer, "_double_scale", None) is None | ||
| assert not any("_double_scale" in key for key in compressed_module.state_dict()) |
There was a problem hiding this comment.
Bot comment.
This test only exercises the pure-PyTorch compression path. If tensorrt_llm is importable on a SM100+ runner, _real_quantize takes the torch.ops.trtllm.fp4_quantize branch and _scale becomes a 1-D uint8 cutlass scale — the assertions here would fail (or the export silently writes a bad scale in production). Consider adding coverage under tests/gpu_trtllm/ or explicitly monkeypatching fp4_compatible/the trtllm import to cover both layouts, so the behaviour is pinned in both environments.
mtq.compress (used by hf_ptq --low_memory_mode) replaces the weight with packed NVFP4 nibbles and stores the per-block scale on the quantizer. The NVFP4 export branch had no handling for this and recomputed the scale from the weight, which is now uint8 packed data: QTensorWrapper reports the packed shape, so the block count came out as in_dim/32 instead of in_dim/16 and the amax was taken over nibble-pair bytes. On DeepSeek-R1-Distill- Llama-70B this wrote weight_scale [8192, 256] instead of [8192, 512] with meaningless values, which dequantizes to garbage. The internal _scale / _double_scale buffers also stayed registered and were serialized as *.weight_quantizer._double_scale entries (560 of them in the 70B checkpoint), so vLLM and the TensorRT-LLM PyTorch backend failed with KeyError before loading any weight. Reuse the compression-time per-block scale, rescaled into the exported weight_scale_2 convention so the weight_scale * weight_scale_2 product that dequantization consumes is preserved, and drop the internal buffers after use. Verified on DGX Spark: the exported checkpoint now matches a normal export (identical key set, weight_scale_2 and input_scale identical, all 154 quantized tensors at the same dequantization error) and serves correctly instead of failing to load. NVBug 5987078 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
f749100 to
2314cb7
Compare
|
Thanks — all four addressed in 1. Swizzled cutlass scale — confirmed and fixed. You were right, and it reproduces on the hardware I was testing on. In a TRT-LLM container on GB10 (sm121): My earlier validation ran in the vLLM container, where 2. 3. Missing clamp — fixed. Now routed through 4. Redundant computation — fixed. Testing. New
|
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review at 2314cb7f0: all four previously-raised items are addressed, and I found no new correctness problems.
Critical items — resolved
- Cutlass-swizzled
_scale(correctness, would have written raw byte values on Blackwell + TRT-LLM). New_compressed_per_block_scale()mirrors the guard inNVFP4QTensor.dequantize(1-D uint8 →cutlass_fp4_scale_to_modelopt_fp4_scale, hardImportErroriftensorrt_llmcan't be imported to convert). Verified the detection condition matches the producing branch inNVFP4QTensor.quantize(..., try_tensorrt=True). Newtests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.pypins the swizzled layout and asserts the precondition so it can't silently degrade into the dense case. - Root cause of the
_double_scaleleak. Fixed at the source:RealQuantLinear.list_of_scale_tensors"double_scale"→"_double_scale", and the export-pathdelattrloop was dropped. I checked the other consumers of that list (megatron.py_parameter_to_keep_in_quantizer_state_dictusesk in key,_get_shard_axis_dictusesendswith("double_scale")) — both already matched._double_scalebefore, so the megatron path is unchanged, not a behavior change.postprocess_state_dict's exact"weight_quantizer." + q_keymatch is the only place the typo mattered, which is consistent with_scalebeing stripped and_double_scaleleaking in the bug report. The gpu test asserts both keys are gone. - Missing E4M3 clamp. Now routed through
_cast_per_block_scale_to_fp8, so the2**-9floor / 448 ceiling apply as on every other NVFP4 path — important here since group-unifiedweight_scale_2makes the rescale ratio ≤ 1 (underflow direction). - Redundant packed-data reduction.
get_weight_scaling_factor()is now skipped viaelif not use_compressed_scale, scoped to the NVFP4 family so other formats still register their scale. Confirmed the laterif weight_scale is not None: register_buffer(...)still writes the buffer, andto_quantized_weight()short-circuits onQTensorWrapper, so the packed nibbles are passed through untouched.
Also checked: BMM-style fused experts (gate_up_proj) are never compressed by pack_real_quantize_weight (no module.weight), so the transpose-orientation mismatch I was worried about for stored _scale can't be reached; the naming ask from cjluo-nv is covered by uses_compressed_nvfp4_scale / use_compressed_scale; the new file's header matches LICENSE_HEADER verbatim, so no licensing concern; size (+186/-3) and CHANGELOG are fine.
Minor, non-blocking: the CHANGELOG entry still describes the buffers as "removed after use" and doesn't mention the list_of_scale_tensors typo (the actual root cause) or the swizzled-scale conversion — worth a one-line touch-up on the way in. And if weight_scale_2 were ever None for an NVFP4 format, the compressed branch would fall through to recomputing from packed data (the original bug) rather than failing loudly; practically unreachable since weight_scale_2 is registered unconditionally for these formats.
2 participants
What does this PR do?
Type of change: Bug fix
Fixes unified HF export of already-compressed NVFP4 weights —
mtq.compressand, through it,examples/llm_ptq/hf_ptq.py --low_memory_mode(NVBug 5987078). Two defects, one root cause: the NVFP4 export branch has no handling for weights that were already real-quantized, unlike theFP8_PB_REALbranch which consumesweight_quantizer._scale.1.
weight_scalewas recomputed from packed data. After compression the weight is aQTensorWrapperof packed NVFP4 nibbles, andQTensorWrapper.__new__builds the Parameter from_quantized_data, so.shapereports the packed shape (the logical shape survives only inmetadata["shape"]). The export derived the block count fromweight.shape[-1]and took amax over nibble-pair bytes, so it wrote a scale of half the required size with meaningless values:weightweight_scalewrittenq_proj[2048, 1024]U8[2048, 64][2048, 128]q_proj[8192, 4096]U8[8192, 256][8192, 512]2. An internal quantizer buffer leaked into the checkpoint.
postprocess_state_dictstripsweight_quantizer.<name>for every name inRealQuantLinear.list_of_scale_tensors, but that list carried"double_scale"where the buffer is_double_scale— a missing underscore. So_scalewas stripped and_double_scalewas not, and it reached the checkpoint as*.weight_quantizer._double_scale(560 entries in the 70B checkpoint). Downstream loaders reject it before loading any weight:This is why only the TensorRT backend appeared usable in the bug report — its converter tolerates the stray key, then produces
!!!!!!output from the broken scales, while the PyTorch backend (vLLM / TensorRT-LLM) fails to load outright.The fix reuses the per-block scale captured at compression time, rescaled into the exported
weight_scale_2convention, and corrects the typo above. The rescale matters: compression normalizes per-block FP8 scales against the global scale it captured at that moment, which is not the post-calibrationweight_scale_2the export writes. Exporting the stored scale as-is loads fine but leaves every block off by a constant factor (~1.96x measured), so theweight_scale * weight_scale_2product that dequantization consumes must be preserved.Note the typo fix also affects
modelopt/torch/quantization/plugins/megatron.py:505,511, which filter on the same list — these are internal buffers so excluding them looks correct, but calling it out since it is a behavior change outside the export path.Compression-time scale layout.
TensorQuantizer._real_quantizecallsNVFP4QTensor.quantize(..., try_tensorrt=True), so on an FP4-capable device with TensorRT-LLM importable the stored_scaleis the cutlass-swizzled 1-D uint8 scale rather than the modelopt 2-D E4M3 layout. Confirmed on GB10 in a TRT-LLM container:The export therefore normalizes it the same way
NVFP4QTensor.dequantizedoes, and raises iftensorrt_llmcannot be imported to convert, rather than writing raw byte values.Usage
No API change. The previously broken path now works:
Testing
All runs on DGX Spark (GB10, sm121, aarch64). Both compression-time scale layouts are covered, since the layout depends on whether TensorRT-LLM is importable in the process:
New tests
tests/gpu/torch/export/test_export_weight_gpu.py::test_export_compressed_nvfp4_weight— dense E4M3 path. Asserts the per-block scale covers the logical input dim, thatweight_scale * weight_scale_2matches an uncompressed export of the same model, and thatpostprocess_state_dictstrips both internal buffers.tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py::test_export_compressed_nvfp4_weight_trtllm_scale— cutlass-swizzled path. Asserts as a precondition that the environment really produced a 1-D uint8 scale, so it cannot silently degrade into the dense case when TensorRT-LLM is absent.Results
Negative controls (each test fails without the code it guards)
mainwith only the test file applied:test_export_compressed_nvfp4_weightfails.test_export_compressed_nvfp4_weight_trtllm_scalefails.End-to-end, TensorRT-LLM container (the environment from the bug report, where
_scaleis swizzled) — TinyLlama-1.1B,--qformat nvfp4 --low_memory_mode, plus a normal export as control. Exported checkpoints are structurally identical (0 stray_double_scalekeys vs. 560 onmain; 663 keys each;q_proj.weight_scale[2048, 128]FP8 in both; QKVweight_scale_2unified in both). Loaded on the TensorRT-LLM PyTorch backend — the backend reported as unusable:On
mainthat same load fails withKeyError: '...weight_quantizer._double_scale'before a single weight is read.End-to-end, vLLM container (dense scale path) — TinyLlama-1.1B, NVFP4 +
--low_memory_mode:KeyError: '...weight_quantizer._double_scale', engine fails to start"Paris is the capital of"→" France and is best known for the awe inspiring Notre-Dame C")input_scaleidentical, andweight_scale_2identical — including the values unified across fused groups, so the fused-GEMM contract (one sharedweight_scale_2per QKV / gate-up group) still holds.mainand is the source of the numbers in the table above.Known limitation: fused groups are slightly less accurate than a full-memory PTQ
This restores a working checkpoint, but
--low_memory_modeNVFP4 is not numerically identical to a normal PTQ, and cannot be made so at export time.The nibbles are packed at load time against the layer's own global scale, so the effective per-block scale baked into them is
fp8_own * ws2_own.preprocess_linear_fusionlater unifiesweight_scale_2across a fused group to the group max, and the format requires the per-block scale be E4M3, so the best the export can write isround_fp8(fp8_own * ws2_own / ws2_unified). For the group member owning the max amax that ratio is exactly 1 and the round trip is bit-exact; for the others it costs one extra E4M3 rounding, bounded by a half-ULP (6.25%). The uncompressed path never pays this because its weights are still high precision at export, soto_quantized_weightre-quantizes the nibbles after unification.Measured on TinyLlama-1.1B (weight relative error vs. the source BF16 weights, 154 quantized tensors):
--low_memory_mode(this PR)The degradation is confined to exactly 66 of 154 tensors = 22 layers x 3, i.e. the non-max members of each
q/k/vandgate/upgroup (worst observed +0.0066, e.g.layers.11.self_attn.k_proj0.0898 -> 0.0964). The 88 remaining tensors — group winners plus the unfusedo_proj/down_proj— are bit-exact.Removing this requires compressing a fusion group against one shared scale in the compress-on-load path (
RealQuantParameterDict), where the group max first becomes known; that is a larger change and is left as a follow-up.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
NVBug 5987078.
Not addressed here: at 70B scale on DGX Spark,
--low_memory_modecan also crash before export when the device map offloads, becauseQTensorWrapper.to()cannot represent ametatensor:It is reproducible on demand under low free GPU memory (crashes at 41.2 GB and 58.5 GB free; succeeds at 88.1 GB) and is easy to hit on Spark's unified memory, where page cache from reading the checkpoint counts against
torch.cuda.mem_get_info(). That is an independent defect and will be filed separately.