[NVBug: 5987078] Fix unified HF export of compressed NVFP4 weights (--low_memory_mode)

[cjluo-nv]

What does this PR do?

Type of change: Bug fix

Fixes unified HF export of already-compressed NVFP4 weights — mtq.compress and, 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 the FP8_PB_REAL branch which consumes weight_quantizer._scale.

1. weight_scale was recomputed from packed data. After compression the weight is a QTensorWrapper of packed NVFP4 nibbles, and QTensorWrapper.__new__ builds the Parameter from _quantized_data, so .shape reports the packed shape (the logical shape survives only in metadata["shape"]). The export derived the block count from weight.shape[-1] and took amax over nibble-pair bytes, so it wrote a scale of half the required size with meaningless values:

	weight	weight_scale written	expected
TinyLlama-1.1B q_proj	[2048, 1024] U8	[2048, 64]	[2048, 128]
DeepSeek-R1-Distill-Llama-70B q_proj	[8192, 4096] U8	[8192, 256]	[8192, 512]

2. An internal quantizer buffer leaked into the checkpoint. postprocess_state_dict strips weight_quantizer.<name> for every name in RealQuantLinear.list_of_scale_tensors, but that list carried "double_scale" where the buffer is _double_scale — a missing underscore. So _scale was stripped and _double_scale was 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:

KeyError: 'layers.0.mlp.down_proj.weight_quantizer._double_scale'
RuntimeError: Engine core initialization failed.

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_2 convention, 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-calibration weight_scale_2 the export writes. Exporting the stored scale as-is loads fine but leaves every block off by a constant factor (~1.96x measured), so the weight_scale * weight_scale_2 product 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_quantize calls NVFP4QTensor.quantize(..., try_tensorrt=True), so on an FP4-capable device with TensorRT-LLM importable the stored _scale is the cutlass-swizzled 1-D uint8 scale rather than the modelopt 2-D E4M3 layout. Confirmed on GB10 in a TRT-LLM container:

logical weight (512, 256)  -> modelopt scale should be (512, 16) e4m3
_scale : (8192,) torch.uint8  (ndim=1)     <- cutlass-swizzled
after cutlass_fp4_scale_to_modelopt_fp4_scale: (512, 16) torch.float8_e4m3fn

The export therefore normalizes it the same way NVFP4QTensor.dequantize does, and raises if tensorrt_llm cannot be imported to convert, rather than writing raw byte values.

Usage

No API change. The previously broken path now works:

python hf_ptq.py --pyt_ckpt_path <local_ckpt_dir> --qformat nvfp4 \
  --low_memory_mode --export_path <out>

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, that weight_scale * weight_scale_2 matches an uncompressed export of the same model, and that postprocess_state_dict strips 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

suite	TensorRT-LLM 1.3.0rc17 container	vLLM 26.05 container
both files above	3 passed	2 passed, 1 skipped

Negative controls (each test fails without the code it guards)

• On main with only the test file applied: test_export_compressed_nvfp4_weight fails.
• With only the un-swizzle conversion neutered, the rest of the fix intact: test_export_compressed_nvfp4_weight_trtllm_scale fails.

End-to-end, TensorRT-LLM container (the environment from the bug report, where _scale is swizzled) — TinyLlama-1.1B, --qformat nvfp4 --low_memory_mode, plus a normal export as control. Exported checkpoints are structurally identical (0 stray _double_scale keys vs. 560 on main; 663 keys each; q_proj.weight_scale [2048, 128] FP8 in both; QKV weight_scale_2 unified in both). Loaded on the TensorRT-LLM PyTorch backend — the backend reported as unusable:

##### trt_normal #####
GEN: France and is the most populous city in the country. It is located on the Seine River...
##### trt_lowmem #####
GEN: France and is the most popular tourist destination in the country. It is a city of art, history...

On main that same load fails with KeyError: '...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:

• before: KeyError: '...weight_quantizer._double_scale', engine fails to start
• after: loads and generates coherently ("Paris is the capital of" → " France and is best known for the awe inspiring Notre-Dame C")
• The fixed checkpoint is structurally equivalent to a normal (non-low-memory) export: identical key set (663 tensors), input_scale identical, and weight_scale_2 identical — including the values unified across fused groups, so the fused-GEMM contract (one shared weight_scale_2 per QKV / gate-up group) still holds.
• DeepSeek-R1-Distill-Llama-70B (the model in the bug) reproduces the same signature on main and 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_mode NVFP4 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_fusion later unifies weight_scale_2 across a fused group to the group max, and the format requires the per-block scale be E4M3, so the best the export can write is round_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, so to_quantized_weight re-quantizes the nibbles after unification.

Measured on TinyLlama-1.1B (weight relative error vs. the source BF16 weights, 154 quantized tensors):

	mean rel. error
normal PTQ	0.090248
--low_memory_mode (this PR)	0.091443

The degradation is confined to exactly 66 of 154 tensors = 22 layers x 3, i.e. the non-max members of each q/k/v and gate/up group (worst observed +0.0066, e.g. layers.11.self_attn.k_proj 0.0898 -> 0.0964). The 88 remaining tensors — group winners plus the unfused o_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"

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌

Additional Information

NVBug 5987078.

Not addressed here: at 70B scale on DGX Spark, --low_memory_mode can also crash before export when the device map offloads, because QTensorWrapper.to() cannot represent a meta tensor:

accelerate/hooks.py: set_module_tensor_to_device(module, name, "meta")
RuntimeError: Attempted to call `variable.set_data(tensor)`, but `variable` and `tensor` have incompatible tensor type.

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.

Summary by CodeRabbit

• Bug Fixes

  • Fixed Hugging Face export for already-compressed NVFP4 weights.
  • Preserved accurate dequantization by reusing and correctly converting stored compression scales.
  • Removed internal quantization buffers from exported checkpoints, preventing downstream loading failures.
  • Corrected handling of TensorRT-LLM swizzled FP4 scale layouts during export.

• Tests

  • Added GPU regression coverage for compressed NVFP4 exports and scale conversion accuracy.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b5ddb1ae-404b-45b2-9356-e8b4a444a756

📥 Commits

Reviewing files that changed from the base of the PR and between 2314cb7 and 1a3616a.

📒 Files selected for processing (2)

• CHANGELOG.rst
• modelopt/torch/export/unified_export_hf.py

🚧 Files skipped from review as they are similar to previous changes (2)

• CHANGELOG.rst
• modelopt/torch/export/unified_export_hf.py

---

📝 Walkthrough

Walkthrough

The Hugging Face exporter now supports compressed NVFP4 weights. It reuses and rescales stored compression scales, preserves dequantization, removes internal quantizer buffers, and adds GPU regression coverage.

Changes

Compressed NVFP4 export

Layer / File(s)	Summary
Preserve compressed NVFP4 scales modelopt/torch/export/unified_export_hf.py, CHANGELOG.rst	The exporter detects compressed weights, converts and rescales stored scales for exported weight_scale_2, and documents the change.
Track internal scale buffers modelopt/torch/quantization/nn/modules/quant_linear.py	RealQuantLinear.list_of_scale_tensors now references _double_scale.
Validate compressed export tests/gpu/torch/export/test_export_weight_gpu.py, tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py	GPU tests compare compressed and uncompressed exports, validate scale conversion and equivalence, and verify internal buffers are absent.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Suggested reviewers: mxino

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

Loading

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the bug fix for unified Hugging Face export of compressed NVFP4 weights.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	The isolated PR diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval/exec, or # nosec usage.

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch chenjiel/fix-lowmem-nvfp4-export

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

modelopt/torch/export/unified_export_hf.py (1)

668-679: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Skip the temporary packed-weight scale calculation.

get_weight_scaling_factor(...) already runs before this code detects QTensorWrapper. 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

📥 Commits

Reviewing files that changed from the base of the PR and between a23390d and c11d9cb.

📒 Files selected for processing (3)

• CHANGELOG.rst
• modelopt/torch/export/unified_export_hf.py
• tests/gpu/torch/export/test_export_weight_gpu.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.93%. Comparing base (2f6e77f) to head (1a3616a).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2038       +/-   ##
===========================================
+ Coverage   66.94%   77.93%   +10.98%     
===========================================
  Files         519      519               
  Lines       59401    59419       +18     
===========================================
+ Hits        39767    46308     +6541     
+ Misses      19634    13111     -6523     

Flag	Coverage Δ
examples	43.11% <55.00%> (-0.19%)	⬇️
gpu	58.77% <100.00%> (+37.82%)	⬆️
regression	15.00% <20.00%> (+0.06%)	⬆️
unit	55.15% <55.00%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

for readability: should we name this branch condition as real_nvfp4_quantized or something similar?

[meenchen]

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:

1. weight_quantizer._scale is not always the modelopt per-block E4M3 layout. TensorQuantizer._real_quantize() calls NVFP4QTensor.quantize(..., try_tensorrt=True); when fp4_compatible() (SM100+), block_size == 16, and tensorrt_llm is importable, it returns the cutlass-swizzled 1-D uint8 scale factor instead (see NVFP4QTensor.dequantize, which explicitly detects scale.dtype == torch.uint8 and scale.ndim == 1 and converts via cutlass_fp4_scale_to_modelopt_fp4_scale). The new export branch has no such guard, so on Blackwell inside a TRT-LLM container --low_memory_mode NVFP4 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 in tests/gpu (not tests/gpu_trtllm), so CI won't either. Please either convert or raise loudly.

2. Root cause of the _double_scale leak looks like a typo in RealQuantLinear.list_of_scale_tensors ("double_scale" instead of "_double_scale"). postprocess_state_dict already strips weight_quantizer._scale via that list and would strip _double_scale too 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.

3. Minor: the new .to(torch.float8_e4m3fn) skips the clamp that _cast_per_block_scale_to_fp8 applies on every other NVFP4 path.

No licensing or design concerns; size and CHANGELOG are fine.

[meenchen]

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.

[meenchen]

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.

[meenchen]

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.

[meenchen]

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.

[cjluo-nv]

Thanks — all four addressed in 2314cb7f0.

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):

fp4_compatible: True   tensorrt_llm: 1.3.0rc17   trtllm.fp4_quantize op: True
logical weight (512, 256)  -> modelopt scale should be (512, 16) e4m3
_scale : (8192,) torch.uint8  (ndim=1)     <- cutlass-swizzled
after cutlass_fp4_scale_to_modelopt_fp4_scale: (512, 16) torch.float8_e4m3fn

My earlier validation ran in the vLLM container, where tensorrt_llm doesn't import, so NVFP4QTensor.quantize fell through to the dense E4M3 path — the bug reporter's environment takes the other branch, so the original patch would have been broken exactly where it mattered. New _compressed_per_block_scale() mirrors the guard in NVFP4QTensor.dequantize: converts when the scale is 1-D uint8, and raises if tensorrt_llm can't be imported to convert rather than writing raw bytes.

2. list_of_scale_tensors typo — agreed, that's the root cause; fixed there instead. "double_scale" -> "_double_scale", and I dropped my export-path deletion loop. Corroborating evidence from the broken checkpoint: weight_quantizer._double_scale leaked but weight_quantizer._scale did not — exactly what a missing underscore predicts, since _scale matched the list and was stripped. Note this also affects megatron.py:505,511, which filter on the same list; that looks like the desired behavior (these are internal buffers) but it is a behavior change in that path, so flagging it explicitly rather than slipping it in.

3. Missing clamp — fixed. Now routed through _cast_per_block_scale_to_fp8. Worth noting the live risk here is underflow, not overflow: fusion unifies weight_scale_2 to the group max, so the rescale ratio is <= 1 and scales shrink — which is what the 2**-9 floor guards.

4. Redundant computation — fixed. get_weight_scaling_factor() is now skipped when the compressed scale will be used (it was both wasted block-amax work over packed bytes and a value immediately overwritten). Scoped to the NVFP4 family so other compressed formats still get their scale registered.

Testing. New tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py covers the swizzled path and asserts as a precondition that the environment really produces a 1-D uint8 scale, so it can't silently degrade into the dense case. Results on GB10:

• TRT-LLM container (swizzled path): 3 passed
• vLLM container (dense path): 2 passed, 1 skipped
• Negative control: neutering only the un-swizzle conversion makes the new test fail, so it guards this specific regression rather than just the original bug.

[meenchen]

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

1. Cutlass-swizzled _scale (correctness, would have written raw byte values on Blackwell + TRT-LLM). New _compressed_per_block_scale() mirrors the guard in NVFP4QTensor.dequantize (1-D uint8 → cutlass_fp4_scale_to_modelopt_fp4_scale, hard ImportError if tensorrt_llm can't be imported to convert). Verified the detection condition matches the producing branch in NVFP4QTensor.quantize(..., try_tensorrt=True). New tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py pins the swizzled layout and asserts the precondition so it can't silently degrade into the dense case.
2. Root cause of the _double_scale leak. Fixed at the source: RealQuantLinear.list_of_scale_tensors "double_scale" → "_double_scale", and the export-path delattr loop was dropped. I checked the other consumers of that list (megatron.py _parameter_to_keep_in_quantizer_state_dict uses k in key, _get_shard_axis_dict uses endswith("double_scale")) — both already matched ._double_scale before, so the megatron path is unchanged, not a behavior change. postprocess_state_dict's exact "weight_quantizer." + q_key match is the only place the typo mattered, which is consistent with _scale being stripped and _double_scale leaking in the bug report. The gpu test asserts both keys are gone.
3. Missing E4M3 clamp. Now routed through _cast_per_block_scale_to_fp8, so the 2**-9 floor / 448 ceiling apply as on every other NVFP4 path — important here since group-unified weight_scale_2 makes the rescale ratio ≤ 1 (underflow direction).
4. Redundant packed-data reduction. get_weight_scaling_factor() is now skipped via elif not use_compressed_scale, scoped to the NVFP4 family so other formats still register their scale. Confirmed the later if weight_scale is not None: register_buffer(...) still writes the buffer, and to_quantized_weight() short-circuits on QTensorWrapper, 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.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-08-03 17:46 UTC